Qwen3 Text Encoder: Fixing FP8 and GGUF Crashes

Loading Qwen3 text encoders in quantized formats—specifically FP8 Safetensors and GGUF—often triggers immediate crashes when paired with the Z-Image model. I've been digging into the tracebacks, and it boils down to strict key checking and dtype mismatches in the attention mechanism.

Here is a deep dive into the two main bugs and how to patch them.

The FP8 "Missing Key" Crash

When loading FP8 Safetensors, you'll likely hit Exception: Missing keys: ['lm_head.weight'].

This happens because many FP8 checkpoints strip the lm_head.weight to save space, since it's tied to model.embed_tokens.weight. The fast_load_transformers_model function in mmgp/offload.py performs a strict key check and panics when it doesn't find the head weight.

The Fix: Use a state dictionary preprocessor to alias the memory reference. By mapping model.embed_tokens.weight to the lm_head.weight slot during the load process, you bypass the strict check without actually duplicating the tensor in VRAM.

The GGUF SDPA Dtype Mismatch

If you're using GGUF encoders (like Q4_K_M), you'll run into a RuntimeError regarding SDPA. The error usually looks like this: Expected query, key, and value to have the same dtype, but got query.dtype: float key.dtype: float and value.dtype: struct c10::BFloat16 instead.

The issue is that the GGUF integration provides the Value tensor in bfloat16, but the Hugging Face transformers library computes RoPE for Query and Key in float32. PyTorch's native SDPA is incredibly picky and crashes if these don't match.

The Fix: Force the entire text encoder to a uniform precision immediately after loading.

In models/z_image/z_image_main.py, the fix is to call:

text_encoder.to(dtype)

right after offload.fast_load_transformers_model finishes. This aligns all projections and lets the native SDPA run without crashing.

These changes ensure that the text encoder doesn't choke on quantized weights, making the whole AI workflow much more stable for those of us not running on massive A100 clusters.