Fix: resolve Qwen3 text encoder loading issues for FP8 and GGUF formats

Description

This PR resolves two separate crashing issues that occur when loading the Qwen3 text encoder in different quantized formats (FP8 Safetensors and GGUF) alongside the Z-Image model.

Issue 1: Missing lm_head.weight in FP8 Safetensors

Error Traceback: Exception: Missing keys: ['lm_head.weight']

• Root Cause: Many quantized FP8 checkpoints intentionally drop the lm_head.weight key to save disk space, as it is tied to model.embed_tokens.weight. However, when fast_load_transformers_model is called, mmgp/offload.py performs a strict key check and crashes.
• Solution: Implemented a unified_preprocessor (with a _fix_qwen_fp8_sd module-level helper) that intercepts the state dictionary as it loads. It dynamically aliases the model.embed_tokens.weight memory reference into the lm_head.weight slot before the strict check occurs, bypassing the error without wasting VRAM on duplicated tensors.

Issue 2: SDPA Dtype Mismatch in GGUF

Error Traceback: RuntimeError: 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.

• Root Cause: When using GGUF text encoders (e.g., Q4_K_M), the underlying integration provides the Value tensor in bfloat16, while Hugging Face transformers computes the RoPE for Query and Key in float32. PyTorch's native SDPA strictly requires all three to share the exact same data type, causing an immediate crash.
• Solution: Forced the text encoder to a uniform precision state directly after loading by calling text_encoder.to(dtype). This aligns all attention projections, allowing HF's native SDPA to run seamlessly without needing to alter the core attention backend.

Changes Made in models/z_image/z_image_main.py

• Injected a state dictionary preprocessor to handle missing tied embedding weights on the fly.
• Added a .to(dtype) cast immediately after offload.fast_load_transformers_model completes to enforce strict dtype uniformity.

Fix: resolve Qwen3 text encoder loading issues for FP8 and GGUF formats #1904

arifanchan posted on Jun 13, 2026

Description

This PR resolves two separate crashing issues that occur when loading the Qwen3 text encoder in different quantized formats (FP8 Safetensors and GGUF) alongside the Z-Image model.

Issue 1: Missing lm_head.weight in FP8 Safetensors

Error Traceback: Exception: Missing keys: ['lm_head.weight']

• Root Cause: Many quantized FP8 checkpoints intentionally drop the lm_head.weight key to save disk space, as it is tied to model.embed_tokens.weight. However, when fast_load_transformers_model is called, mmgp/offload.py performs a strict key check and crashes.
• Solution: Implemented a unified_preprocessor (with a _fix_qwen_fp8_sd module-level helper) that intercepts the state dictionary as it loads. It dynamically aliases the model.embed_tokens.weight memory reference into the lm_head.weight slot before the strict check occurs, bypassing the error without wasting VRAM on duplicated tensors.

Issue 2: SDPA Dtype Mismatch in GGUF

Error Traceback: RuntimeError: 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.

• Root Cause: When using GGUF text encoders (e.g., Q4_K_M), the underlying integration provides the Value tensor in bfloat16, while Hugging Face transformers computes the RoPE for Query and Key in float32. PyTorch's native SDPA strictly requires all three to share the exact same data type, causing an immediate crash.
• Solution: Forced the text encoder to a uniform precision state directly after loading by calling text_encoder.to(dtype). This aligns all attention projections, allowing HF's native SDPA to run seamlessly without needing to alter the core attention backend.

Changes Made in models/z_image/z_image_main.py

• Injected a state dictionary preprocessor to handle missing tied embedding weights on the fly.
• Added a .to(dtype) cast immediately after offload.fast_load_transformers_model completes to enforce strict dtype uniformity.

View on GitHub