Repository navigation
Add disable_mmap kwarg to from_pretrained with hf-mount auto-detection - #45547
Merged
ArthurZucker merged 3 commits intoApr 22, 2026
Merged
ArthurZucker merged 3 commits into
ArthurZucker merged 3 commits into
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
ArthurZucker
approved these changes
Apr 21, 2026
ArthurZucker
left a comment
ArthurZucker
left a comment
Collaborator
There was a problem hiding this comment.
sounds great ty for adding
Collaborator
|
@bot /style |
Contributor
|
Style fix bot fixed some files and pushed the changes. |
rtrompier
marked this pull request as ready for review
April 21, 2026 09:43
XciD
added a commit
to huggingface/hf-mount
that referenced
this pull request
Nov 13, 2026
…ans from window plans (#246) Fixes #234. Supersedes #241 and #254. ## Summary Parallel page faults on one memory-mapped file ran one at a time, at about 3 MB/s. transformers, torch and safetensors load checkpoints this way. The per-handle `PrefetchState` mutex was held across network I/O, and interleaved streams replaced each other in its single buffer, so almost every fault became a new 256 KiB range download of about 50 ms. This is the stall that made transformers turn mmap off on hf-mount (huggingface/transformers#45547). Many cold readers at once also wedged the mount (#234). Each read stream held xet download buffer for data that only a FUSE worker could consume, and the workers were blocked on streams that could not get buffer for their first bytes. Reads timed out at offset 0 and ended in EIO. This PR replaces `PrefetchState` with a `RemoteReader` per lazy handle (`src/virtual_fs/remote_reader.rs`): - **Blocks, fetched concurrently.** The file is cached in 256 KiB blocks. Missing blocks are fetched by background tasks, each streaming one bounded range. No lock is held across network I/O, and a read of a block in flight waits for it instead of fetching it again. - **Read-ahead per stream.** Up to 32 sequential streams are tracked per handle. Each keeps a window of blocks requested ahead of its reads: 2 MiB at first (64 MiB for a read from the start of the file), growing fourfold then twofold up to 512 MiB, and topped up in fetches that download in parallel. A fetch close to the reader is short (1 MiB, then as long as its distance to the reader) so that the first blocks arrive early; fetches far ahead are 32 MiB, and no fetch crosses a multiple of 32 MiB, so that every pass over a file asks for the same ranges and a re-read finds them in the chunk cache. A stream's read-ahead stops where another active stream ahead of it started, also inside its last block, so parallel copies of one tensor do not fetch each other's ranges. A stream that starts up to 16 MiB past the furthest block read (the next tensor, or the next record header of a torch zip) reads ahead at once; further reads (the slices of a tensor-parallel rank) wait for a second read, as with the old prefetch code. - **Fetching goes on past a slow fetch, within what the link carries.** A xet range delivers its data only when the whole range has arrived, so fetches finish out of order. At most half of the read-ahead budget is in flight, and each completed fetch tops its stream's window up again while the stream reads. While a read waits for a slow fetch, the next ones keep downloading instead of the network idling. xet downloads the terms of the fetches in the order they were asked for, so on a slow link a deep pipeline would make the last fetches wait past the fetch timeout: a reader caps the read-ahead in flight that no read waits for yet (64 MiB at first, plus the bytes of each read-ahead fetch that arrives within a quarter of the fetch timeout, halved by each slower or failed one, between 8 MiB and half its limit). On a fast link the cap soon reaches half the limit. - **Bounded memory.** A handle holds at most 512 MiB of pending or unread blocks. Read-ahead is reserved in a budget per mount before its fetches start (1 GiB, `--read-ahead-mb`), and each handle may hold an even share of it. A handle that holds read-ahead or that the budget turns away counts in the split; one whose only fetches are blocks its reads wait for does not. A handle above its share gives fetched read-ahead back at its next read (first what no stream heads to, then the blocks furthest ahead). Blocks already read (for re-reads and for streams that meet in a block; the page cache keeps the rest) are capped at 8 MiB per handle and 256 MiB across the handles of a mount, keeping the last block of each live stream. A handle that has no read waiting and read no new block for 10 s gives back its read-ahead beyond its share; after 120 s it gives back all of it and the blocks it read, and its next read starts small windows again. So a reader that reads in bursts keeps what its next burst needs, and a file kept open does not keep a share. Every 256 reads, a handle also drops the unread blocks of streams that stopped. Blocks are copied out of xet terms, so that a cached block does not keep a whole 64 MiB term alive. - **Fetches do not wait for readers.** A fetch copies its data into blocks as it arrives, whether a reader waits for it or not, so no xet download buffer stays held until a FUSE worker consumes it. This breaks the cycle of #234. - **FUSE reads answer from a task when they wait.** A read of data already fetched is still served on the worker thread. A read that waits for the network finishes in a task, so it does not hold one of the `--max-threads` workers while other faults queue behind it. - **xet settings from #241.** The download buffer limit goes from 256 MiB to 1 GiB, as one fast reader alone keeps up to 256 MiB in flight, and the cap of the adaptive download concurrency goes from 64 to 124, unless `HF_XET_FIXED_DOWNLOAD_CONCURRENCY` is set. The per-stream download buffers of #241 are not needed: fetches are bounded and drained as they arrive. Retries (3 attempts in a row that deliver nothing, resuming at the first missing block), the per-chunk stall timeout (`--read-fetch-timeout-ms`), EOF handling, and `--direct-io` (a re-read fetches the bytes again, even inside a block read in part) keep their behavior, with one change to the timeout: a fetch that gets no chunk within it waits on while other fetches of the mount receive chunks. The link serves fetches in the order they were asked for, so on a slow link the first read of a cold file can wait behind the read-ahead of other files: it waits in line rather than stalls. Each stream notes the bytes the mount's other streams still wait for when it opens, and a fetch waits as long as the mount still receives them, and at least four timeouts. On a 3 MB/s link, 8 cold files of 256 MiB read at once finish at the speed of the link, where a fixed allowance of four timeouts made them all end in EIO. With nothing received on the mount, one timeout still fails the attempt. A read that waits for several fetches fails as soon as one of them gives up; a read whose failed fetch it only joined tries once more with fetches of its own. New unit tests cover parallel reads, sharing a block in flight, interleaved streams fetching each byte once, read-ahead up to the end of the file, forward scans, refill while a read waits, the budgets of a handle and of a mount (with a handle the budget turns away), the read-ahead of an idle handle (and of one whose read waits for a slow fetch), blocks a read waits for staying at hand until it ends, the block of a cancelled read, a read that ends after its failed block was fetched again, a fetch task that ends early, reads of 8 and 40 blocks whose fetches fail and stall at once, fetches spawned on a runtime that shuts down, the eviction queue and partial re-reads with `--direct-io`, cancellation on close, streams that meet inside a block, sparse forward reads, bursts with pauses and a read after a long pause, refill once a stream stops, readers with only demand fetches, the read-ahead of a stopped stream, blocks already read split across a mount and kept at the end of a live stream, attempts that deliver blocks, a read that joined a failed fetch, fetch runs within cells, blocks reads wait for and the cap, and a slow link (a mock link that downloads terms in the order they were asked for, as xet does), for one file and for cold readers of several files (3 files at 3 MB/s, 16 at 12 MB/s, 8 of 256 MiB at 3 MB/s, 24 at 5 MB/s, 3 at 1 MB/s), a stalled fetch that waits in line while the mount receives data, but not forever, and a fetch behind more bytes than four timeouts carry. `test_fuse_parallel_mmap_read` has 16 threads copy one memory-mapped 48 MiB file through a mount with 4 worker threads and checks the content. `test_fuse_parallel_cold_reads`, from #241, has 24 cold readers of 64 MiB files go through a mount with 4 worker threads and a 5 s fetch timeout, and fails on any read error or fetch timeout. ## Benchmarks m6i.2xlarge in us-east-1, `--no-disk-cache`, page cache dropped and a fresh mount for every run. `Qwen/Qwen2.5-1.5B-Instruct/model.safetensors` (3.09 GB) unless noted. Throughput to CAS changed a lot during the runs (a plain sequential read of the same file went from about 260 to 830 MB/s from one run to the next), so the sequential rows interleave the binaries and give medians. Only the last table includes the xet settings commit. | Workload | main | This PR | |---|---|---| | safetensors, 4 threads, `natural` key order (transformers-like) | 2.7 MB/s (77 of 338 tensors in 411 s) | 195 to 224 MB/s | | safetensors, 1 thread, file order (torch copies each tensor with 8 threads) | 5.9 MB/s (161 of 338 tensors in 302 s) | 301 to 364 MB/s | | `torch.load(mmap=True)` then copy, 4 threads, `EleutherAI/pythia-1.4b` (2.93 GB) | 2.0 MB/s (47 of 364 tensors in 358 s) | 68 to 82 MB/s (36 to 43 s) | | 8 threads copying random 4 MiB ranges of one mapping | 1.8 MB/s, p50 17.9 s per copy | 157 to 210 MB/s, p50 152 to 195 ms | | Sequential `read()`, 54 MB file (median of 3) | 110 MB/s | 166 MB/s | | Sequential `read()`, 91 MB file (median of 3) | 149 MB/s | 275 MB/s | | Sequential `read()`, 268 MB file (median of 3) | 285 MB/s | 290 MB/s | | Sequential `read()`, 3.09 GB file (median of 4) | 590 MB/s | 946 MB/s | | Sequential mmap scan, 3.09 GB file (median of 4) | 575 MB/s | 893 MB/s | | Random 4 KiB page faults, p50 / p99 (3 rounds) | 62 to 63 ms / 183 to 227 ms | 63 to 67 ms / 195 to 281 ms | The runs on main stopped at their deadline, hence the partial tensor counts. The small files are `google/electra-small-discriminator/pytorch_model.bin`, `sentence-transformers/all-MiniLM-L6-v2/model.safetensors` and `distilbert/distilbert-base-uncased/model.safetensors`. Peak RSS of the daemon during a sequential read of the 3.09 GB file: 875 to 959 MiB with the first commits, about 720 to 830 MiB after the review fixes (below), against 698 to 747 MiB on main (read-ahead budget of 512 MiB per handle). Bytes read through a mount were compared with a copy downloaded directly from the Hub, for 16 threads reading contiguous ranges of one mapping, 16 threads reading random ranges, and a safetensors load tensor by tensor: all identical. ### Many cold readers (#234) Shards of `Qwen/Qwen2.5-72B-Instruct` read at once from a cold mount, each from its start, with `--max-threads 4 --read-fetch-timeout-ms 5000` like `test_fuse_parallel_cold_reads`: | Load | main | This PR | |---|---|---| | 24 files, first 64 MiB of each | 177 s, EIO on 19 files | 2.6 s (619 MB/s), no errors | | 37 files, first 256 MiB of each | 301 s, EIO on 32 files | 15 s (661 MB/s), no errors | In the second run with this PR, one fetch got no data for 5 s and its retry succeeded. ### xet settings 3 rounds with the default thread count and fetch timeout, peak RSS of the daemon in parentheses: | Workload | Without the settings | With the settings | |---|---|---| | 37 cold readers, first 256 MiB of each | 978 to 1106 MB/s (1156 to 1185 MiB) | 1097 to 1216 MB/s (1172 to 1258 MiB) | | Sequential `read()` of the 3.09 GB file | 577 to 845 MB/s (771 to 907 MiB) | 669 to 976 MB/s (832 to 895 MiB) | For one reader of a small file, the settings make no measurable difference. Whole-file `read()` as in the CI benchmark, medians of 5 rounds: | File | main | This PR without the settings | This PR | |---|---|---|---| | `google/electra-small-discriminator/pytorch_model.bin` (54 MB) | 112 MB/s | 192 MB/s | 182 MB/s | | `distilbert/distilbert-base-uncased/model.safetensors` (268 MB) | 308 MB/s | 350 MB/s | 353 MB/s | | `openai-community/gpt2/model.safetensors` (548 MB) | 368 MB/s | 548 MB/s | 515 MB/s | The CI benchmark varies more than that between runs of the same code: three runs of its 500 MB FUSE read gave 893, 524 and 1000 MB/s. ### Review fixes The mount budget reserved before fetches start, idle handles, and the bookkeeping cleanups (4552bf4 to 32f3b26), against the code before them with the same xet settings. 3 rounds, peak RSS of the daemon in parentheses: | Workload | Before | After | |---|---|---| | 37 cold readers, first 256 MiB of each | 829 to 1173 MB/s (1190 to 1246 MiB) | 1148 to 1236 MB/s (1095 to 1118 MiB) | | 24 cold readers of 64 MiB, `--max-threads 4 --read-fetch-timeout-ms 5000` | 761 to 953 MB/s | 675 to 845 MB/s | | Sequential `read()` of the 3.09 GB file | 794 to 889 MB/s (756 to 784 MiB) | 702 to 864 MB/s (741 to 831 MiB) | | safetensors, 4 threads, `natural` key order (1 round) | 210 MB/s (822 MiB) | 234 MB/s (905 MiB) | The cleanup that followed (eb85607) reads at the same speed as 32f3b26, within the noise: over 3 rounds, 628 to 779 against 707 to 789 MB/s with 37 readers, 574 to 698 against 648 to 754 MB/s for one reader, and 189 to 225 against 189 to 219 MB/s for safetensors with 4 threads. No fetch timed out. The same binary varies by up to 30% between rounds here (the code before gave 599 to 965 MB/s in the 24-reader case over two runs), so only the 37-reader case moved beyond the noise. When many readers start at once, the first ones give back the end of their window as the others arrive and fetch it again later. ### Second review round PR before the fixes of the second review (506487e) against this PR, m6i.2xlarge, 5 rounds interleaved, medians: | Workload | Before | After | |---|---|---| | safetensors, 4 threads, `natural` key order | 254 MB/s | 230 MB/s | | Sequential `read()`, 548 MB / 3.09 GB | 546 / 966 MB/s | 572 / 959 MB/s | | fio sequential, 548 MB, FUSE / NFS | 555 / 527 MB/s | 591 / 518 MB/s | | 37 cold readers, first 256 MiB of each | 1283 MB/s | 1250 MB/s | | Headers of 37 shards | 3.3 s | 3.2 s | The safetensors load is about 9% slower in every comparison of the day; variants that restore the 64 MiB forward-scan gap, drop the cap of read-ahead in flight or drop the cell alignment did not bring it back, and CDN stalls during the runs (30 s without data on some fetches, for both binaries) blurred the smaller differences. With all traffic into the instance shaped to 3 MB/s (`tc`), reading 400 MiB of the 3.09 GB file: main read about 3 MiB in 15 minutes (its 128 MiB prefetch windows got no data within 30 s, EIO every 3 minutes); the PR before these fixes had 143 fetch timeouts and an EIO at 165 MiB; this PR reads it in 136 s (3.08 MB/s, the link rate), with no timeout, no error and 275 MiB of peak RSS. ## Window plans (`--window-plans`, off by default) Depends on huggingface/xet-core#1064 (`Client::get_chunk_layout`): xet-core is pinned to that branch (`1d4b8487`, on xet-core main after v1.8.1) for now, to repin on xet-core main once it merges. xet asks CAS for the plan of each read-ahead fetch in three pieces (8, 16 and 8 MiB), so the remote reader above makes about 96 reconstruction requests per GiB read, and a memory-mapped load with 16 threads sends 200 to 250 per second, above the rate of the CAS WAF limit (13,000 requests per minute per IP). With `--window-plans on`, hf-mount asks for the plan and the chunk layout (`GET /v2/chunk-layout`, huggingface-internal/xetcas#1414) of an aligned 256 MiB window, and answers every plan request inside it locally. The first and last terms are trimmed on chunk boundaries with the uncompressed sizes. Each term is fetched as a sub-range of the window's signed ranges, located with the serialized sizes. The result is the plan CAS would return for the range, so the downloads do not change. - Windows are asked in the background, so no read waits for one. A reader that starts at the beginning of the file, or comes from the previous window, asks its window at once with the next 4 (1 GiB ahead): a window takes about 0.3 s plus 1 s per GiB to arrive, and a fast reader goes through a window in less than that. - Elsewhere, a window is asked on the second plan request that falls in it, so a lone random read never pulls one. - Reads through the mount give the file size: no window is asked past the end of a file, and files under 256 MiB keep exact plans, since they are read before a window arrives. - Plans stay exact when CAS does not serve chunk layouts (a 404 turns the feature off for the mount), for direct S3 URLs (they sign the exact range), and for a window that failed (asked again after a minute). - At most 128 windows are kept, by last use, and a window is replaced after 50 minutes from its plan request, before its signed URLs expire. A window in flight keeps its slot until its fetch ends, so at most 128 window fetches run at once, and none of them twice. A window used in the last 30 seconds is not dropped for another one: with more large files read at once than the cache holds (37 files at 32 MB/s), a request elsewhere keeps its exact plan, instead of dropping windows in use that are then fetched again (9,216 window fetches for 592 windows before, 394 now). - `--window-plans check` derives the plans, compares them with the plans CAS returns (terms, offset, fetched chunks, and byte ranges against the layout), logs any difference, and uses the CAS plans. This is the first rollout step. Production CAS now serves `/v2/chunk-layout`: a `check` mount there read a 3.1 GB file and 16 threads of random mmap reads with 282 derived plans, all equal to the CAS plans. The option is off by default. RFC: https://app.notion.com/p/3f01384ebcac81d384e0cca29335116d xet-core#1064 also builds reqwest with gzip, so every CAS request asks for compressed responses. CAS already compresses its read responses: a 256 MiB chunk layout goes from 54 to 25 KB and a 1 GiB plan from 180 to 26 KB, so windows arrive 70 to 350 ms sooner from Paris. ### Window plan results Cold mounts with `--no-disk-cache` against production CAS, variants interleaved, medians of 3 rounds. 1,482 mounts in all, with no data difference between variants, and no plan difference in `check` mode (5,416 derived plans compared on the final build). Paris client (88 ms to CAS), no windows / with windows: | Load | Time | CAS requests | |---|---|---| | 3.1 GB file, sequential | 3.9 / 3.2 s | 281 / 59 | | safetensors mmap, 16 threads | 6.8 / 3.7 s | 1,367 / 116 | | 2,000 random 4 KiB reads | 16.3 / 2.1 s | 1,854 / 113 | | 8 shards of 3.9 GB in parallel | 20.0 / 15.4 s | 2,677 / 348 | | 19.9 GB file, sequential | 23.0 / 15.4 s | 1,779 / 212 | The windows were faster in 40 of 42 paired rounds in Paris and 30 of 42 in us-east-1 (2 ms to CAS, a noisier network), where random reads are 2.6 times faster and CAS requests 3 to 35 times fewer. Window sizes from 128 MiB to 1 GiB with 1 to 4 windows ahead, xet-core reconstruction and concurrency settings, fetch size, read-ahead and buffer sizes were also compared: the current hf-mount settings stay the best. ### Window plan notes for review - The window plan is asked through the raw CAS client, not through the plan cache, so window plans do not fill the exact-plan cache. - The chunk layout is asked through the xet-core client, so it shares the mount's CAS auth, retries and user agent. A 404 means the CAS does not serve the endpoint. - Peak RSS rises by 250 to 750 MB on fast reads, from the extra bytes in flight; the windows themselves weigh a few MB. - Tests build plans the way CAS builds them for a deduplicated file and check that the derived plans are equal for about a hundred ranges, that `check` reports a wrong byte range, the fetch, 404 and direct S3 paths, the windows asked ahead, that small files and files of unknown size keep exact plans, and that a CAS that never answers chunk layouts gets at most 128 window fetches, none of them twice. ## Notes - A sequential read makes about three CAS reconstruction queries per 32 MiB fetch (304 for the 3.09 GB file): xet splits each fetch in 8 and 16 MiB reconstruction blocks (`HF_XET_RECONSTRUCTION_MIN_RECONSTRUCTION_FETCH_SIZE`, unchanged) that it downloads in parallel and delivers in turn. The single stream of main made a few queries for the whole file. With one 32 MiB block per fetch, `torch.load(mmap=True)` opened faster (12 to 14 s instead of 23 to 30 s), but small files read 40% slower. Fetches now fall on the same 32 MiB cells on every pass, so the plan cache serves a re-read within the hour; a first read can only make fewer queries once `CachedXetClient` can derive tight range plans from a full plan, which needs chunk sizes in the terms of xet-core (`derive_range_response`). - #228 moves every FUSE operation to a task. This PR only moves the reads that wait for the network: a read of data already fetched is answered on the worker thread, without a task handoff for each 128 KiB request. Both change the same `read` handler. - transformers still reads safetensors fully into memory on hf-mount (`_is_on_hf_mount`). Its first forward pass over mmap-backed weights touches them one parameter at a time, in module order, not file order, so each parameter costs at least one round trip: about 55 s for this model, against about 8 s for the full read. Dropping that workaround needs more than this PR (for example, filling the page cache ahead of the reads). --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds a new
disable_mmapkwarg toPreTrainedModel.from_pretrained(and plumbs it throughload_state_dict/LoadStateDictConfig/_load_pretrained_model). When enabled, safetensors checkpoints are read fully into memory and parsed viasafetensors.torch.load(bytes)instead of being memory-mapped withsafe_open.Why
hf-mount: captured kernel stacks showed Python hanging on page-faults with a saturated FUSE queue.diffusersalready ships the exact same kwarg (disable_mmapindiffusers/models/model_loading_utils.py); adding it totransformersmakes the two libraries consistent and unblocks users who currently hitTypeError: ...__init__() got an unexpected keyword argument 'disable_mmap'when they try to pass it through.How
_is_on_hf_mount(path)helper parses/proc/mountson Linux and returnsTrueifpathresolves under a mountpoint whose device string ishf-mount(returnsFalseon non-Linux).load_state_dict(...)gains adisable_mmap: bool | None = Nonekwarg. WhenNone, it auto-detects via_is_on_hf_mount. When truthy (andmap_location != "meta"), the safetensors file is loaded withsafetensors.torch.load(open(path, "rb").read()).LoadStateDictConfiggets a matchingdisable_mmapfield;from_pretrainedforwards its new kwarg into it; the multi-shard safetensors branch in_load_pretrained_modelmirrors the same no-mmap path per file.hf-mountFUSE mount, so there's no behavior change for existing users.Tests
Adds
DisableMmapLoadingTestintests/utils/test_modeling_utils.pycovering:_is_on_hf_mountmatches when/proc/mountslistshf-mountfor the path._is_on_hf_mountdoes not match for a regular filesystem._is_on_hf_mountshort-circuits toFalseon non-Linux platforms.load_state_dict(..., disable_mmap=True)returns tensors equal to the mmap path.