[Pfannkuchensack]

Summary

Wan 2.2 held far more system RAM than the model cache accounts for. Three separate causes, all in Wan code, each fixed here:

• Expert swap. The outgoing expert stayed referenced by _ExpertSwapper.get() and by the denoise loop's transformer local while the next expert loaded. The cache evicted it, but its tensors could not be freed, so both experts were resident during every swap. The default cache is about 1x VRAM, below two experts, so this hits default settings.
• Diffusers folders. Wan's Diffusers weights are fp32 on disk and _keep_in_fp32_modules leaves some tensors fp32. Those stay views of the mmap, and each pins its whole shard for the life of the model, while the cache only counts the bf16 size. After from_pretrained, the loader now copies only the tensors still on the mapping that are not bf16. bf16 tensors stay on the mapping as on main, because they are the model itself, so bf16 folders load exactly as before.
• Single files. The loader used to read the whole file before it cast anything or dropped bundled VAE and text-encoder weights. It now decides everything from the header first: the prefix, which keys are dropped, how the file is quantized, and the cache reservation at the size after the fold and the cast. Then it reads one tensor at a time through pread and folds or casts each one before the next. Dropped keys and scale metadata are never read, and no tensor stays a view of a file-wide mapping. ComfyUI int8 files reserve before their first read too.

Peak and final RAM (working set), Windows, CPU-only loads, TI2V-5B:

Case	main	first version of this PR	now
Single file fp16 → bf16 (9.3 GB)	19.5 / 10.2 GB	10.3 / 10.2 GB	10.3 / 10.2 GB
Single file fp32 with bundled VAE (21.3 GB)	28.8 / 10.2 GB	22.1 / 10.2 GB	10.3 / 10.2 GB
Diffusers fp32, 5 shards (18.6 GB)	28.7 / 28.4 GB	16.8 / 10.2 GB	28.7 / 10.2 GB
Diffusers bf16, unsharded (10 GB)	10.2 / 10.2 GB	19.5 / 10.2 GB	10.2 / 10.2 GB

Checksums match across all three versions in every case. The fp32 Diffusers peak equals main's, see the known limits.

From the first version, with the expert-swap fix (unchanged since):

Case	Before	After
T2V 14B int8 pair, max_cache_ram_gb: 16, real clips	29.7 GB	16.2 GB

Related Issues / Discussions

Follow-up announced in #10185. Both PRs edit the cast block in WanCheckpointModel._load_from_singlefile, so the second to merge needs a small rebase. The RAM figures in wan-2.2.mdx from #10185 describe the state before this PR and need updating once both are in.

QA Instructions

Peak RAM above is process working set. Windows commit reads higher because VRAM counts toward it 1:1 (12 GB on the device added 12 GB of commit and no working set, with either CUDA allocator).

• Clips: int8 pair from feat(wan): load ComfyUI int8 convrot checkpoints #10185 with this fix applied, 832x480, 17 frames, 8 steps, use_cache=false. All three seeds are bit-identical before and after (same SHA-256). 65 s per clip after, 68-72 s before.
• Load figures: real loaders in a fresh process, CPU only, with main and the first version run from git archive trees. The working set is read after touching every byte, as a device move would. The bundled-VAE single file was built locally from the TI2V-5B folder, because no real one was at hand.
• For pread, reading from disk showed no penalty: 462 MiB/s against 431 MiB/s for mmap on a cold 19 GiB file.
• Tests use tiny real safetensors files and record reads at the safe_open boundary:
  • bundled keys are never read;
  • the one reservation comes before the first payload read, on the fp8/dense and the int8 paths;
  • each tensor is cast before the next one is read;
  • after a load nothing reads the file anymore, checked by overwriting the payload bytes in place, with a positive control through mmap. These probes are skipped where safetensors' mapping does not show in-place writes, as on the macOS runners;
  • an fp32 folder leaves nothing on the mapping, and a bf16 folder keeps exactly its bf16 weights there;
  • an MXFP8 exponent grid stored as F8_E8M0 gets its refusal by name.
  • Each of these fails when the matching part of the fix is reverted.
• Tests: tests/backend/model_manager/load, tests/backend/model_manager/test_model_on_disk.py, tests/backend/model_manager/configs, tests/model_identification and tests/backend/quantization: 3499 passed, 161 skipped. Wan loader and invocation tests: 411 passed, 1 skipped. uv lock --check, scripts/check_platform_locks.py and Ruff are clean.

Not run: Linux, macOS, ROCm, multi-GPU, the family E2E (Wan has no manifest yet). The 14B fp32 expert was not re-measured after the review round: its peak would not fit the free RAM of the test machine. No real scaled-fp8 Wan file was at hand, so the fold is covered by tests only.

Review

Fixed during review:

• The first reservation came up short for weights the fold widens. It now matches the post-cast size exactly across scaled fp8, plain and non-float tensors.
• The single-file loader read the whole file through pread before it dropped or cast anything; it now works from the header and reads tensor by tensor.
• The blanket disable_mmap doubled unsharded bf16 folders; it is gone, and only the tensors that pin a mapping are copied.
• A Wan file storing its MXFP8 exponent grid as F8_E8M0 failed with a bare KeyError in the header read; it reaches its refusal by name again, and any other unreadable dtype is named.

Known limits:

• fp32 Diffusers peak. diffusers keeps every shard mapped until from_pretrained returns, so the fp32 copies are dropped only then. The peak is main's: on Linux the extra pages are clean and reclaimable but count toward RSS and cgroup limits, and on Windows the mapping is charged to commit. For an A14B expert that is about 57 GB of shards plus about 28.5 GB of bf16, against a 57 GB reservation. Dropping them per shard would mean patching diffusers' loading loop.
• safetensors holds the GIL while pread reads a tensor. A cold read stalls other threads for up to one tensor's read time: 1.9 s measured on a USB disk. This matters most with multiple GPUs.
• A file the fold refuses now evicts cached models before it raises, because the reservation comes first.
• The reservation for Diffusers folders still uses file size, twice the model, so it evicts more than needed. Unchanged here.
• Other families that call load_file and cast, or use _keep_in_fp32_modules, are unchanged.

Compatibility / Rollout

safetensors>=0.8.0 is a new floor, because backend= does not exist before it. uv.lock was already on 0.8.0, so only its metadata line changes. No migrations, no API change.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

---

^{Open in [code]smith · Need help on this PR? Tag @codesmith-bot with what you need.}

[joshistoast]

[P1] Checkpoint loading now materializes the entire source before shrinking it.

load_file(..., backend="pread") eagerly allocates every tensor before FP32-to-BF16 conversion or removal of bundled VAE/text-encoder weights. On RAM-constrained Linux systems, previously reclaimable mapped source pages become anonymous allocations, so loading can exhaust RAM even when the resulting model fits. Cache eviction cannot solve an oversized source allocation. Filter keys before reading payloads, then read and convert tensors incrementally.

[Pfannkuchensack]

Fixed in 5aa9b11. The single-file loader now settles everything from the header before it reads any payload: the prefix, the dropped VAE and text-encoder keys, the quantization checks, and the cache reservation at the size after the fold and the cast. Then it reads one tensor at a time through pread and folds or casts each one before reading the next. Dropped keys and scale metadata are never read. The int8 branch also reserves before its first read now.

On a 21 GB fp32 TI2V-5B single file with a bundled VAE (9.3 GB model), the peak working set is 10.3 GB, against 22.1 GB for the first version and 28.8 GB on main. Tests check that bundled keys are never read, that the reservation comes before the first read, and that each tensor is cast before the next one is read.

[joshistoast]

[P2] Diffusers loading unnecessarily doubles whole-file allocations.

Blanket disable_mmap=True also affects supported unsharded BF16 folders, where the cited FP32 backing-storage problem does not apply. Diffusers holds the complete file bytes alongside deserialized tensors, while cache reservation covers only one file’s size. Previously working loads can now OOM or swap heavily. Use bounded reading or selectively detach problematic mapped tensors.

[Pfannkuchensack]

Agreed, the blanket disable_mmap is gone (5aa9b11). Diffusers loads through mmap again, exactly as on main. Afterwards the loader copies only the tensors left on the mapping that are not bf16, which are the fp32 _keep_in_fp32_modules views that pinned the fp32 shards. A bf16 folder keeps its weights on the mapping: the unsharded 10 GB TI2V-5B folder peaks at 10.2 GB working set again, against 19.5 GB with disable_mmap. The fp32 5-shard folder now ends at 10.2 GB, where main stays at 28.4 GB.

One limit, also in the description: during from_pretrained the fp32 peak is still main's (28.7 GB here), because diffusers keeps every shard mapped until it returns and the copies run after that. Releasing them per shard would mean patching diffusers' loading loop, so I left that out of this PR.

[joshistoast]

Overall, the expert-reference fix is small and sound, and the dense streaming approach and regression tests are strong. The documented fp32 Diffusers peak remains a limitation.

[joshistoast]

[P2] Previously readable fp8 checkpoints now fail.

Unconditionally uses a header reader whose dtype table omits F8_E4M3FNUZ and F8_E5M2FNUZ. Both worked through the previous loader. Identification still accepts them through its fallback, so installation can succeed while generation fails. Real safetensors CPU probes reproduced this. Add the missing mappings and regression coverage.

[Pfannkuchensack]

Fixed in ce8458f. SAFETENSORS_DTYPES now maps F8_E4M3FNUZ, F8_E5M2FNUZ and C64, which is everything safetensors.torch reads except F4. F4 is still refused by name, since its header shape counts packed elements. The new test test_every_dtype_safetensors_writes_reads_back_from_the_header writes every torch dtype save_file accepts and reads it back from the header. It cannot skip a dtype the table maps. Both FNUZ formats, plain and scaled, now load through the Wan single-file path bit-identical to load_file + cast/fold.

[joshistoast]

[P2] Mixed int8 checkpoints can exceed the RAM reservation.

Reads every tensor before casting, despite reserving only the final size. A real mixed fp32/int8 fixture reserved 4.21 MB but held 8.40 MB of payloads before conversion. Unlike the previous reclaimable mmap pages, these are anonymous allocations, increasing OOM risk. Cast dense tensors as they arrive; current int8 coverage uses float32 compute and misses this case.

[Pfannkuchensack]

Fixed in ce8458f. The read loop now casts each dense tensor as it arrives; the marked layers' payload and scale stay as stored for the split. No local keeps the stored tensor alive past its cast either. On your kind of fixture (int8 + fp32 dense, bf16 compute) the peak during the read stays within the reservation. The new test test_a_header_read_never_holds_more_than_the_reservation runs at bf16 and also counts handed-out payloads still referenced outside the state dict. It fails on the old order (264 KB held vs 133 KB reserved). Through the real Wan loader the weights are bit-identical to before, with the same single make_room.