[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]

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.