[refactor]: Z-Image reuses the shared Qwen3 encoder instead of a bespoke port

main gained a config-driven Qwen3 text encoder (Qwen3ForCausalLM + Qwen3TextConfig)
via the Flux2 Klein port (hao-ai-lab#1349). It returns BaseEncoderOutput(last_hidden_state,
hidden_states), supports GQA + the same qkv/gate_up fusion, and update_model_arch
populates Z-Image-Turbo's dims (2048/24/16) from its config.json — so our parallel
Qwen3Model was ~540 lines of duplicate.

Drop the bespoke encoder + config (resolved during the rebase onto main; the files
are now identical to main) and wire Z-Image to the shared encoder:
- registry: map Z-Image-Turbo's "Qwen3Model" architecture string -> Qwen3ForCausalLM.
- encoder parity test: use Qwen3TextConfig; move the strict-load allowlist
  ({"lm_head.weight"}) from a class attribute to a test-side constant, since the
  shared encoder's loader is intentionally lenient (serves multiple models).
- PORT_STATUS / README: mark the encoder "reused"; the prior A40 parity PASS was
  measured against the removed bespoke encoder, so text-encoder parity is flagged
  RE-VALIDATION PENDING against the shared class. Scheduler/tokenizer/VAE parity
  are unaffected and still PASS.

The gated scheduler option (use_reference_discrete_timesteps, default False) is
unchanged and still needed for Z-Image timestep parity.

Author: Mister-Raggs

tests/local_tests/zimage/PORT_STATUS.md

@@ -14,12 +14,12 @@
 - phase: Phase 4 (component parity)
 - status: in_progress
 - owner: parity
-- last_updated: 2026-05-12
+- last_updated: 2026-06-21
 
 ## Component Matrix
 | Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues |
 |---|---|---|---|---|---|---|---|---|---|
-| Text encoder (Qwen3) | text_encoder | ported | `zimage.Qwen3Model` (HF-Qwen3 layout) | `Qwen3ForCausalLM` checkpoint at `<weights>/text_encoder/` | `fastvideo/models/encoders/qwen3.py` + `fastvideo/configs/models/encoders/qwen3.py` | DONE | not_needed (raw safetensors load) | PASS (fp32 + bf16) | I001, I003 |
+| Text encoder (Qwen3) | text_encoder | reused | `zimage.Qwen3Model` (HF-Qwen3 layout) | `Qwen3ForCausalLM` checkpoint at `<weights>/text_encoder/` | shared `fastvideo/models/encoders/qwen3.py::Qwen3ForCausalLM` (+ `Qwen3TextConfig`), added for Flux2 Klein; Z-Image's `Qwen3Model` arch string routes to it in the registry | DONE | not_needed (raw safetensors load) | RE-VALIDATION PENDING (shared encoder) | I001, I003 |
 | Tokenizer | tokenizer | reused | `AutoTokenizer` | `<weights>/tokenizer/` | `fastvideo/models/loader/component_loader.py::TokenizerLoader` | DONE | not_needed | PASS | none |
 | VAE | vae | reused | `zimage.AutoencoderKL` (Diffusers-compatible) | `<weights>/vae/` | `fastvideo/models/vaes/autoencoder_kl.py` | DONE | not_needed | PASS (decode only) | encode-path parity deferred to pipeline |
 | Scheduler | scheduler | reused (with extension) | `zimage.FlowMatchEulerDiscreteScheduler` | `<weights>/scheduler/scheduler_config.json` | `fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py` (added `use_reference_discrete_timesteps`) | DONE | not_needed | PASS | I002 |
@@ -30,7 +30,7 @@
 - conversion_script: `<none — encoder loads raw safetensors via tests/local_tests/zimage/test_zimage_encoder_parity.py helpers; VAE/tokenizer/scheduler use HF subfolder loaders>`
 - converted_weights_dir: `<none>`
 - source_layout: diffusers
-- strict_load_status: `pass_with_documented_exclusions` (Qwen3Model.ALLOWED_UNEXPECTED_KEYS = {"lm_head.weight"})
+- strict_load_status: `pass_with_documented_exclusions` (allowlist `{"lm_head.weight"}` asserted in the encoder parity test; the shared encoder is body-only and owns no LM head)
 - passthrough_components: VAE config, scheduler config, tokenizer assets
 - retry_history: `<none>`
 
@@ -40,23 +40,23 @@
 | Scheduler | `pytest tests/local_tests/zimage/test_zimage_scheduler_parity.py -v -s` | PASS (2/2) on Z-Image-Turbo, A40, 2026-05-12 | full `scheduler_config.json` now forwarded (was 3 keys) |
 | Tokenizer | `pytest tests/local_tests/zimage/test_zimage_tokenizer_parity.py -v -s` | PASS (2/2) on Z-Image-Turbo, A40, 2026-05-12 | tokenizer resolves to `Qwen2TokenizerFast`; `apply_chat_template` parity included |
 | VAE decode | `pytest tests/local_tests/zimage/test_zimage_vae_parity.py -v -s` | PASS (1/1) on Z-Image-Turbo, A40, 2026-05-12 | encode-path deferred |
-| Text encoder fp32 | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[fp32]` | PASS on Z-Image-Turbo, A40, 2026-05-12 | bit-exact (`last_hidden_state` max=0.0000, `hidden_states[-2]` max=0.0012) across both batches — FastVideo Qwen3 port is numerically correct |
-| Text encoder bf16 | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[bf16]` | PASS on Z-Image-Turbo, A40, 2026-05-12 | empirical (worst across 2 batches): `last_hidden_state` mean=0.0168 median=0.0127 (thresholds 0.025 / 0.020, 1.5x headroom); `hidden_states[-2]` mean=0.0739 median=0.0625 (thresholds 0.120 / 0.100, 1.6x headroom). Per-layer diag confirms monotonic accumulation across 35 layers, no single-layer spike — textbook bf16-tail signature |
-| Per-layer bf16 diag | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_per_layer_bf16_diagnostic -v -s` | PASS (informational only) on Z-Image-Turbo, A40, 2026-05-12 | prints 37 hidden-state diffs (embedding + 35 layers + post-norm) for future debugging |
+| Text encoder fp32 | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[fp32]` | RE-VALIDATION PENDING (shared encoder) | Prior PASS (A40, 2026-05-12, bit-exact: `last_hidden_state` max=0.0000, `hidden_states[-2]` max=0.0012) was against the now-removed bespoke `Qwen3Model`. The test now drives the shared `Qwen3ForCausalLM`; re-run on A40 to confirm parity carries over. |
+| Text encoder bf16 | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_parity_forward[bf16]` | RE-VALIDATION PENDING (shared encoder) | Prior PASS (A40, 2026-05-12): `last_hidden_state` mean=0.0168 median=0.0127 (thresholds 0.025 / 0.020); `hidden_states[-2]` mean=0.0739 median=0.0625 (thresholds 0.120 / 0.100). Measured against the bespoke encoder; thresholds may need recalibrating against the shared one. Re-run on A40. |
+| Per-layer bf16 diag | `pytest tests/local_tests/zimage/test_zimage_encoder_parity.py::test_zimage_qwen3_encoder_per_layer_bf16_diagnostic -v -s` | RE-VALIDATION PENDING (shared encoder) | Informational. Prior run (A40, 2026-05-12) printed 37 hidden-state diffs (embedding + 35 layers + post-norm); confirm the shared encoder's `hidden_states` length/indexing matches HF on re-run. |
 
 ## Open Questions
 | ID | Question | Owner | Needed By Phase | Status | Resolution |
 |---|---|---|---|---|---|
 | Q001 | Pin a Z-Image reference clone SHA in the README before handoff | prep | Phase 1 | resolved | Pinned `Tongyi-MAI/Z-Image@26f23eda626ffadda020b04ff79488e1d72004cd` (2026-05-12) |
 | Q002 | Final HF id for published Z-Image weights | prep | Phase 1 | resolved | `Tongyi-MAI/Z-Image-Turbo` (6B, 8 NFE, fits 16 GB) and `Tongyi-MAI/Z-Image` (full, 32.9 GB). Both Apache-2.0 |
-| Q003 | Does Z-Image use Qwen3 chat-template tokenization at pipeline runtime? Currently `Qwen3Config.is_chat_model=False` | pipeline | Phase 6 | open | |
+| Q003 | Does Z-Image use Qwen3 chat-template tokenization at pipeline runtime? The shared `Qwen3TextConfig` (from Flux2 Klein) defaults `is_chat_model=True`; Z-Image's removed bespoke config had `False`. Reconcile when wiring the pipeline preset. | pipeline | Phase 6 | open | |
 
 ## Issues And Blockers
 | ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution |
 |---|---|---|---|---|---|---|---|---|
-| I001 | Phase 4 | text_encoder | medium | `Qwen3ForCausalLM` checkpoints ship `lm_head.weight`; encoder-only `Qwen3Model` does not own an LM head. Strict-load must allowlist this key. | `fastvideo/models/encoders/qwen3.py::Qwen3Model.ALLOWED_UNEXPECTED_KEYS`; encoder parity test asserts the unexpected-key set ⊆ allowlist | parity | resolved | Allowlist landed in this PR; loader raises on any other unexpected key. |
+| I001 | Phase 4 | text_encoder | medium | Z-Image-Turbo ships a full Qwen3 checkpoint with `lm_head.weight`; the shared encoder (`Qwen3ForCausalLM`) is body-only and owns no LM head, so that key goes unmatched. Must be allowlisted; anything else is a real silent drop. | encoder parity test asserts the unexpected-key set ⊆ `_ALLOWED_UNEXPECTED_KEYS = {"lm_head.weight"}` | parity | resolved | Allowlist enforced in the parity test (the shared encoder's loader is intentionally lenient as it serves multiple models, so the contract lives test-side). |
 | I002 | Phase 5/6 | scheduler | high | `scheduler_config.json` at `<weights>/Z-Image/scheduler/` does not pin `use_reference_discrete_timesteps=True`. Stock loaders will silently fall back to default timestep mode (numerically different — parity tests prove the divergence). | `tests/local_tests/zimage/test_zimage_scheduler_parity.py` sets the flag programmatically | pipeline | open | Pin the flag in `scheduler_config.json` when wiring the pipeline preset. |
-| I003 | Phase 6 | text_encoder | low | `Qwen3ArchConfig.text_len=512` → `tokenizer_kwargs.max_length=512`, but parity tests tokenize at 96/128. Pipeline preset must reconcile. | `fastvideo/configs/models/encoders/base.py::TextEncoderArchConfig.__post_init__` | pipeline | open | Set the correct `text_len` from the Z-Image preset config when adding the pipeline preset. |
+| I003 | Phase 6 | text_encoder | low | The shared `Qwen3TextArchConfig.text_len=512` → `tokenizer_kwargs.max_length=512`, but parity tests tokenize at 96/128. Pipeline preset must reconcile. | `fastvideo/configs/models/encoders/base.py::TextEncoderArchConfig.__post_init__` | pipeline | open | Set the correct `text_len` from the Z-Image preset config when adding the pipeline preset. |
 | I004 | Phase 4 | transformer | high | `ZImageTransformer2DModel` not yet ported | PR body, in-progress | port | open | Future PR. |
 | I005 | Phase 6 | pipeline | high | No FastVideo pipeline class, registry entry, preset, or example yet | n/a | pipeline | open | Future PR. |
 
@@ -67,11 +67,13 @@
 ## Decisions
 | Date | Decision | Rationale | Impact |
 |---|---|---|---|
+| 2026-06-21 | Reuse the shared `Qwen3ForCausalLM` encoder (+ `Qwen3TextConfig`) instead of the bespoke `Qwen3Model` this PR originally added. | `main` gained a config-driven Qwen3 text encoder via the Flux2 Klein port (#1349); it returns `BaseEncoderOutput(last_hidden_state, hidden_states)`, supports GQA + the same qkv/gate_up fusion, and `update_model_arch` populates Z-Image-Turbo's dims (2048/24/16) from its `config.json`. Keeping our parallel encoder would have been ~540 lines of duplicate maintenance. | Dropped `fastvideo/models/encoders/qwen3.py` + `fastvideo/configs/models/encoders/qwen3.py` from this PR; added a 1-line registry map (`"Qwen3Model" → Qwen3ForCausalLM`). Encoder parity needs re-validation on A40 against the shared class (see Parity Commands). |
 | 2026-05-12 | bf16 encoder parity uses distribution checks (mean + median) instead of element-wise `assert_close`; thresholds calibrated to empirical Z-Image-Turbo numbers on A40 + 1.5–1.6x headroom. | Z-Image-Turbo's Qwen3 text encoder is 35 layers (not the 24 originally assumed). Cross-kernel bf16 (FastVideo's fused QKVParallel + MergedColumnParallel + SiluAndMul vs HF's unfused equivalents) accumulates into a long max tail (~4.0 at layer 34) but median stays low (0.06). Per-layer diagnostic test confirmed growth is smooth and monotonic with no single-layer spike — textbook bf16-tail signature, fp32 is bit-exact. Element-wise `assert_close` is meaningless on this profile; mean + median + the per-layer diag together detect real bugs (which push mean ≫ atol AND median > 0.01). | Final assertion shape: `last_hidden_state` mean < 0.025, median < 0.020; `hidden_states[-2]` mean < 0.120, median < 0.100. Validated on NVIDIA A40 (driver 565.57.01, 46068 MiB) 2026-05-12. |
 | 2026-05-12 | `AutoModel.from_pretrained` uses `dtype=` (not `torch_dtype=`). | transformers 4.57.3 emits `torch_dtype is deprecated! Use dtype instead!` warning. Mrinaald's original `dtype=` kwarg was correct; the temporary switch to `torch_dtype=` (in response to a Copilot review comment) was reverted. | – |
 | 2026-05-12 | Scheduler parity forwards the full `scheduler_config.json` dict (minus Diffusers loader keys), not 3 hand-picked keys. | Future on-disk fields (`time_shift_type`, `invert_sigmas`, etc.) would have been silently dropped. | Makes parity reflect the actual on-disk config. |
 
 ## Handoff Notes
-- Component parity (scheduler / tokenizer / VAE / Qwen3 fp32 + bf16) is **fully validated** on `Tongyi-MAI/Z-Image-Turbo` weights on A40 as of 2026-05-12.
-- Loader-side strictness is contract-asserted in the encoder parity test; do not relax `ALLOWED_UNEXPECTED_KEYS` without updating the test.
+- Scheduler / tokenizer / VAE parity are **fully validated** on `Tongyi-MAI/Z-Image-Turbo` weights on A40 (2026-05-12) and are unaffected by the encoder-reuse pivot.
+- Text-encoder parity (fp32 + bf16) needs **re-running on A40** against the shared `Qwen3ForCausalLM` — the 2026-05-12 PASS numbers were against the removed bespoke `Qwen3Model`. The fp32 path is expected to stay bit-exact (the shared encoder shipped validated for Flux2 Klein), but bf16 thresholds may need recalibration; confirm before claiming PASS.
+- Loader-side strictness is contract-asserted in the encoder parity test via `_ALLOWED_UNEXPECTED_KEYS`; do not relax it without updating the test.
 - Next port-stack steps (separate PR, not in #1339 scope): `ZImageTransformer2DModel` port (I004), pipeline preset including `use_reference_discrete_timesteps=True` pinned in `scheduler_config.json` (I002), conversion-or-direct-load story, SSIM media regression (blocked on PR #1321's `media_extension` helper landing for T2I `.png` output).

tests/local_tests/zimage/README.md

@@ -62,23 +62,26 @@ pytest tests/local_tests/zimage/ -v -s
 | Scheduler (`FlowMatchEulerDiscreteScheduler` + `use_reference_discrete_timesteps`) | [`test_zimage_scheduler_parity.py`](./test_zimage_scheduler_parity.py) | full `scheduler_config.json` forwarded; pipeline must also pin the new flag at load time | PASS |
 | Tokenizer (`TokenizerLoader` vs `AutoTokenizer`) | [`test_zimage_tokenizer_parity.py`](./test_zimage_tokenizer_parity.py) | `apply_chat_template` parity included | PASS |
 | VAE decode (`AutoencoderKL`) | [`test_zimage_vae_parity.py`](./test_zimage_vae_parity.py) | decode-only; encode path deferred until pipeline | PASS |
-| Text encoder (`Qwen3Model`) | [`test_zimage_encoder_parity.py`](./test_zimage_encoder_parity.py) | parametrized fp32 + bf16; bf16 uses calibrated `atol=0.05` + diagnostic prints + abs-mean drift check | PASS |
+| Text encoder (shared `Qwen3ForCausalLM`, reused) | [`test_zimage_encoder_parity.py`](./test_zimage_encoder_parity.py) | parametrized fp32 + bf16; bf16 uses calibrated distribution checks + diagnostic prints. Z-Image's `Qwen3Model` checkpoint routes to the shared encoder via the registry | RE-VALIDATION PENDING (shared encoder) |
 
 ## Known Blockers / Open Items
 
 See [`PORT_STATUS.md`](./PORT_STATUS.md) for the live tracker. Highlights:
 
-- `Qwen3Model.ALLOWED_UNEXPECTED_KEYS = {"lm_head.weight"}` — `Qwen3ForCausalLM`
-  checkpoints carry an LM head the encoder bucket does not need. Loader raises
-  if any other unexpected key appears.
+- The text encoder reuses the shared `Qwen3ForCausalLM` (added for Flux2 Klein,
+  #1349); Z-Image-Turbo's `Qwen3Model` architecture string routes to it via the
+  model registry. Z-Image-Turbo's full Qwen3 checkpoint carries an `lm_head.weight`
+  the body-only encoder does not own, so the encoder parity test allowlists exactly
+  that key (`_ALLOWED_UNEXPECTED_KEYS`) and fails on any other unmatched key.
 - `tests/local_tests/zimage/test_zimage_scheduler_parity.py` builds the
   FastVideo scheduler with `use_reference_discrete_timesteps=True` programmatically.
   When the pipeline lands, `<repo_root>/official_weights/Z-Image/scheduler/scheduler_config.json`
   must pin this flag, otherwise stock loaders will silently fall back to the
   default Diffusers timestep mode (numerically different).
-- `Qwen3ArchConfig.text_len = 512` derives `tokenizer_kwargs.max_length = 512`
+- The shared `Qwen3TextArchConfig.text_len = 512` derives `tokenizer_kwargs.max_length = 512`
   via `TextEncoderArchConfig.__post_init__`, but the parity tests tokenize at
-  `max_length=96..128`. Reconcile when the pipeline preset lands.
+  `max_length=96..128`. Reconcile when the pipeline preset lands. (Also note the
+  shared config defaults `is_chat_model=True`, vs Z-Image's removed bespoke `False`.)
 
 ## Review Notes
 

tests/local_tests/zimage/test_zimage_encoder_parity.py

@@ -4,7 +4,9 @@
 
 This compares:
 1) direct transformers AutoModel output, and
-2) FastVideo Qwen3Model wrapper output,
+2) FastVideo's shared Qwen3 encoder (``Qwen3ForCausalLM``, added for Flux2
+   Klein and reused here — Z-Image-Turbo's ``Qwen3Model`` checkpoint routes
+   to it via the model registry),
 
 using identical local checkpoint, tokenization, and inputs.
 
@@ -24,10 +26,17 @@
 from transformers import AutoModel, AutoTokenizer
 from safetensors.torch import safe_open
 
-from fastvideo.configs.models.encoders.qwen3 import Qwen3Config
+from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
 from fastvideo.distributed.parallel_state import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
 from fastvideo.models.registry import ModelRegistry
 
+# Strict-load contract: the shared Qwen3 encoder is body-only (embed_tokens +
+# layers + norm, no lm_head). Z-Image-Turbo ships a full Qwen3 checkpoint, so
+# `lm_head.weight` is the only key that may go unmatched; anything else means a
+# real silent drop. Enforced here in the test since the shared encoder's loader
+# is intentionally lenient (it's used by multiple models).
+_ALLOWED_UNEXPECTED_KEYS = {"lm_head.weight"}
+
 
 REPO_ROOT = Path(__file__).resolve().parents[3]
 ZIMAGE_TEXT_ENCODER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "text_encoder"
@@ -172,17 +181,17 @@ def test_zimage_qwen3_encoder_parity_forward(dtype: torch.dtype):
     for k in ("_name_or_path", "transformers_version", "model_type", "torch_dtype"):
         cfg_raw.pop(k, None)
 
-    cfg = Qwen3Config()
+    cfg = Qwen3TextConfig()
     cfg.update_model_arch(cfg_raw)
     fv = fv_cls(cfg).eval()
     loaded = fv.load_weights(_iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR))
     assert loaded, "No Qwen3 weights were loaded into FastVideo model"
     fv = fv.to(device=device, dtype=dtype)
 
-    # Per add-model-02-parity strict-load contract: assert the production
-    # loader's allowlist matches the actual checkpoint surface. Non-encoder
-    # heads like `lm_head.weight` (from `Qwen3ForCausalLM` checkpoints) are
-    # the only acceptable unexpected keys; anything else means a real drop.
+    # Strict-load contract: assert the checkpoint surface matches the model's
+    # params modulo the allowlist. Non-encoder heads like `lm_head.weight`
+    # (the body-only encoder has no lm_head) are the only acceptable unmatched
+    # keys; anything else means a real silent drop.
     checkpoint_keys = {name for name, _ in _iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR)}
     checkpoint_keys = {k[len("model."):] if k.startswith("model.") else k for k in checkpoint_keys}
     param_names = {name for name, _ in fv.named_parameters()}
@@ -198,9 +207,9 @@ def test_zimage_qwen3_encoder_parity_forward(dtype: torch.dtype):
         else:
             mapped_keys.add(ckpt_name)
     unexpected = mapped_keys - param_names - {"rotary_emb.inv_freq", "rotary_emb.cos_cached", "rotary_emb.sin_cached"}
-    assert unexpected <= fv_cls.ALLOWED_UNEXPECTED_KEYS, (
-        f"Unexpected checkpoint keys not in ALLOWED_UNEXPECTED_KEYS: "
-        f"{sorted(unexpected - fv_cls.ALLOWED_UNEXPECTED_KEYS)}"
+    assert unexpected <= _ALLOWED_UNEXPECTED_KEYS, (
+        f"Unexpected checkpoint keys not in allowlist: "
+        f"{sorted(unexpected - _ALLOWED_UNEXPECTED_KEYS)}"
     )
 
     with torch.no_grad():
@@ -301,7 +310,7 @@ def test_zimage_qwen3_encoder_per_layer_bf16_diagnostic():
     cfg_raw = _load_json(ZIMAGE_TEXT_ENCODER_DIR / "config.json")
     for k in ("_name_or_path", "transformers_version", "model_type", "torch_dtype"):
         cfg_raw.pop(k, None)
-    cfg = Qwen3Config()
+    cfg = Qwen3TextConfig()
     cfg.update_model_arch(cfg_raw)
     fv = fv_cls(cfg).eval()
     fv.load_weights(_iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR))