Repository navigation
feat(wan): load the UMT5 text encoder from GGUF builds - #9967
Pfannkuchensack wants to merge 10 commits into
Conversation
city96's UMT5-XXL GGUFs now install as Wan T5 encoders and load with packed weights and a tokenizer built from the GGUF fields; before, they were taken for FLUX T5 encoders. The llama.cpp T5 helpers move to util/llamacpp_t5.py, and ModelOnDisk.metadata() returns typed GGUF scalars.
# Conflicts: # docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx # invokeai/frontend/api/schema.ts
# Conflicts: # invokeai/frontend/api/schema.ts
There was a problem hiding this comment.
A few things to address before this can merge. I reviewed by reading the code and did not run the branch locally; CI is green at this head.
1. Merge conflict. invokeai/frontend/api/schema.ts conflicts with main (from #10094). Please merge main and regenerate schema.ts and openapi.json, even though openapi.json itself merges cleanly.
2. Use unpack_ggml_at_load in the Wan GGUF loader. Details are inline. This removes the embedding peak the PR lists as a limitation for a later PR, without changing the shared helper.
3. Low-VRAM, multi-GPU and FLUX paths haven't been run. The real run in the description loaded the encoder fully on an RTX 4090, and test_wan_t5_encoder_gguf_loader.py doesn't go through the model cache. Reading the code, the tied embedding and patched forwards follow the same pattern as the FLUX T5 GGUF loader, so I expect them to work, but nothing shows it yet. Low-VRAM users are the main audience for this feature. Please run and report:
- Partial load: a Wan prompt with the encoder only partly in VRAM, for example with a VRAM cap, checking that the output matches a full load.
- Multi-GPU: a second device adopting the shared CPU weights. If you don't have the hardware, please list it as untested in the description.
- FLUX or SD 3: a real generation with a T5 GGUF, since this PR moves that loader's key mapping, config inference and feed-forward patch into shared code and changes the cast in that patch.
4. Docs caveat (docs/src/content/docs/Users Guide/05.Video Generation/wan-2.2.mdx:58). The page offers "any other quantization" with no warning, but your own Q4_K_M measurement shows 21 % relative L2 error and a worst-token cosine of 0.40. Please add a sentence saying that lower quantizations change the conditioning, and that Q8_0 stays closest to the Diffusers encoder.
5. Out-of-date PR description.
- The "
ModelOnDisk.metadata()for GGUF" section describes a change that isn't in this diff:model_on_disk.pyisn't among the changed files, because that behavior is already onmain. - The "Starter" section names "Wan 2.2 T5 Encoder (GGUF Q8_0)", but the code uses "Wan T5 Encoder (UMT5-XXL, GGUF Q8_0)".
- The embedding reservation limitation will also need updating once item 2 is done.
| with accelerate.init_empty_weights(): | ||
| model = UMT5EncoderModel(infer_t5_encoder_config(sd, UMT5Config)) | ||
|
|
||
| # Not `unpack_ggml_at_load` with the model's declared skip patterns, as other GGUF loaders do: |
There was a problem hiding this comment.
The skip patterns are an argument to this helper, so they aren't a reason to avoid it. ltx2.py:746 already calls it with () for a text encoder:
unpack_ggml_at_load(sd, model, (), self._ram_cache.make_room)Called after the model is built and before load_state_dict_ignoring_extras, it keeps quantized Linear weights packed and unpacks everything else plus unquantized Linears, one tensor at a time. That replaces the loop below and materialize_quantized_embedding; sd.clear() is no longer needed to release the originals. The embed_tokens re-tie stays.
Why it matters:
- Decode peak. The embedding is the one tensor here decoded whole.
unpack_ggml_tensordecodes in row slices, and the comment above_UNPACK_CHUNK_ELEMENTSdescribes this exact case. This is the limitation the PR defers, and it doesn't need a change tomaterialize_quantized_embedding. - Reservation. Every cold load currently reserves 11.7 GiB, which can evict the RAM cache. The helper reserves what the load actually holds: about 7.8 GiB for Q8_0 by my calculation.
- Unreserved widening. On a device without bf16, an F16 file widens to float32 in this loop (about 10 GiB more) with nothing reserved. The helper accounts for it.
Numbers for Q8_0:
| Now | With the helper | Basis | |
|---|---|---|---|
| Decode rise | ~3.9 GiB | ~2.2 GiB | Measured on CPU at 1/8 scale with get_dequantized_tensor vs the sliced decode, then scaled up (the slice overhead is fixed, so 2.2 is an upper estimate) |
| Total peak | 9.6 GiB | ~7.8 GiB | Your measurement, and that minus the decode difference |
For Q4_K_M, your 11.5 GiB peak implies about 6 GiB saved, but I didn't measure that. I also haven't run this change against your tests, so please confirm the packed/unpacked and tied-embedding tests still pass.
…ml_at_load Replaces the manual unpack loop and embedding reservation with the shared helper; docs note that lower quantizations change the conditioning.
|
Thanks, all addressed (82fb271, b290123):
|
# Conflicts: # invokeai/frontend/api/schema.ts
Summary
city96's GGUF builds of UMT5-XXL, Wan 2.2's text encoder (
city96/umt5-xxl-encoder-gguf), now install as Wan T5 encoders and work in Wan 2.2 generation.Until now, these files were claimed by the FLUX / SD 3 T5 encoder config. The keys follow the same llama.cpp
t5encoderlayout, so the file installed, but it produced unusable embeddings and could not be selected for Wan.Install. A new
WanT5Encoder_GGUF_Config:t5conversions;T5Encoder_GGUF_Configno longer claims UMT5 or mT5. Extended-vocabulary T5 v1.1 builds still match.Load. A new
WanT5EncoderGGUFLoaderbuilds transformers'UMT5EncoderModel:woincluded.unpack_ggml_at_load(as the LTX-2 text encoder uses) unpacks the rest one tensor at a time. The embedding is decoded in row slices and the RAM reservation comes from the helper.shared).Tokenizer. The tokenizer is built from the GGUF's own fields (vocabulary, scores, flags), so no tokenizer file has to be downloaded or vendored. It matches the Wan 2.2 Diffusers
tokenizer.jsonexactly.Shared T5 code. The llama.cpp T5 key mapping, the config inference and the feed-forward patch move from the FLUX loader into
util/llamacpp_t5.py. Both loaders now use it.wo's storage dtype, which made city96's F16 T5 GGUFs fail on a dtype mismatch.Old records. A UMT5 file that an earlier version installed as a FLUX T5 encoder is refused at load, with a hint to use "Re-identify model".
Nodes.
Starter. A new starter, "Wan T5 Encoder (UMT5-XXL, GGUF Q8_0)" (~6 GB). Bundles are unchanged.
Docs. The Wan 2.2 page notes that lower quantizations change the conditioning and that Q8_0 stays closest to the Diffusers encoder.
QA Instructions
Automated checks
ruff check/ruff format --checkon the changed files: clean.pytest -n 4ontests/backend/model_manager,tests/backend/quantization,tests/model_identification,tests/backend/ltx2, plus the Wan model loader, packed-layer encoder, SD 3 encoder and Wan denoise tests: 3099 passed, 157 skipped, 1 xfailed.openapi.json/schema.tsregenerated (new config class); the api package lint and test pass. webv2pnpm lint(format, oxlint, tsc, architecture) passes.New tests
UMT5EncoderModelbuilt from the same weights;</s>, padding, truncation.prepend_scheme="first";Tokenizer parity
Checked against the Wan 2.2 TI2V-5B Diffusers
tokenizer.json, using the Wan node's tokenizer arguments: 0 mismatches.On an RTX 4090, through the loader with the app's eval mode
Per-token cosine of the last hidden state against the Wan 2.2 TI2V-5B Diffusers UMT5 in bf16, 8 prompts (EN / ZH / DE / full-width / emoji), token IDs identical:
Browser (webv2)
Review
Material findings resolved:
unpack_ggml_at_load: the embedding is decoded in row slices, and the RAM reservation covers what the load holds, including the float32 widening on a device without bf16.T5Encoder_GGUF_Configstops claiming mT5.Remaining limitations:
unpack_ggml_at_load(about 7.8 GiB estimated for Q8_0, no separate embedding reservation). The embedding is decoded in row slices, so its decode peak is lower than a whole-tensor decode. The estimate is not re-measured on a GPU run.AutoTokenizerunder transformers 5.10) does not reproducetokenizer.json. It maps unknown characters to<s>instead of<unk>. The GGUF tokenizer matchestokenizer.json. Not changed here.Partial load, FLUX T5 GGUF and multi-GPU (RTX 4090, through the real
ModelCacheand loaders)max_vram_cache_size_gb=3, which left about 3.9 GiB of the 8.5 GiB (Q8_0) or 6.5 GiB (Q4_K_M) encoder on the GPU. For both Q8_0 and Q4_K_M, the two embeddings are bit-identical (max abs diff 0).HFEncoder(as the FLUX text-encoder node does) for 3 prompts at 512 tokens, onmain(e28937b) and on this branch. The embeddings are bit-identical, so the moved key mapping, config inference and feed-forward patch leave FLUX and SD 3 conditioning unchanged. On the branch, a partial load (2 GiB cap) gave the same bit-identical result.Compatibility / Rollout
WanT5Encoder_GGUF_ConfigjoinsAnyModelConfig, andopenapi.json/schema.tsare regenerated. No migration is needed.Checklist
What's Newcopy (if doing a release after this PR)Need help on this PR? Tag
@codesmith-botwith what you need.