main gained a config-driven Qwen3 text encoder (Qwen3ForCausalLM + Qwen3TextConfig)
via the Flux2 Klein port (#1349). Reuse it for Z-Image instead of the parallel
~540-line bespoke Qwen3Model this PR originally added: drop the duplicate encoder
+ config, and map Z-Image-Turbo's "Qwen3Model" architecture string to the shared
encoder in the registry. update_model_arch populates Z-Image's dims from config.json.
Validating the reuse on L40S surfaced a real, batch-only divergence in the shared
encoder: it built position_ids as [1, seq_len], but the rotary layer flattens
positions to num_tokens and reshapes q/k to (num_tokens, -1, head_dim). For
batch>1 that folded the batch dim into the head dim and misaligned RoPE (fp32
mean_diff 0.22 at batch=2; batch=1 was fine). Fix: expand position_ids to
[batch_size, seq_len]. batch=1 is byte-identical, so Flux2 Klein is unaffected;
this also fixes a latent batch bug in the shared encoder.
Encoder parity now PASSES on L40S (Z-Image-Turbo): fp32 bit-exact against the
shared encoder (both batch elements, max=0.0000); bf16 within the existing
thresholds (last_hidden mean ~0.016, pre-norm mean ~0.07-0.08).
Also in this PR (unchanged): the gated scheduler option
use_reference_discrete_timesteps (default False) for Z-Image timestep parity,
the Z-Image component parity tests (encoder/scheduler/tokenizer/VAE), and the
encoder parity test's strict-load allowlist + OOM-safe two-model handling
(free HF ref + gc.collect before loading FastVideo).
Component parity validated on Tongyi-MAI/Z-Image-Turbo on NVIDIA A40
(46068 MiB, driver 565.57.01) as of 2026-05-12. All thresholds pass
with 1.5-1.6x headroom over observed:
last_hidden_state bf16 mean (worst): 0.0168 < 0.025
last_hidden_state bf16 median (worst): 0.0127 < 0.020
hidden_states[-2] bf16 mean (worst): 0.0739 < 0.120
hidden_states[-2] bf16 median (worst): 0.0625 < 0.100
Handoff notes updated to reflect that component parity is done; next
port-stack steps (DiT, pipeline, conversion, SSIM) are scoped to a
future PR, not #1339.
Per-layer bf16 diagnostic test on Z-Image-Turbo (Qwen3 35 layers, not the
24 originally assumed) confirmed the divergence between FastVideo's
Qwen3 port and HF Qwen3 is a textbook bf16-tail signature, not a code
bug:
- fp32 forward is bit-exact (last_hidden_state max=0.0000)
- bf16 per-layer mean drift grows smoothly from 5e-4 to 7.5e-2 over 35
layers; no single-layer spike (consecutive layer ratio ≤ 1.5x)
- embedding (idx 0) is bit-exact in bf16 too — weight load is right
Root cause is well-understood: FastVideo's TP-aware fused linears
(QKVParallel, MergedColumnParallel, SiluAndMul) execute matmul ops in a
different order than HF's unfused equivalents. In fp32 the per-element
differences cancel statistically (~zero mean drift); in bf16 the
rounding-order accumulates into a long max tail.
Per add-model-02-parity calibration block: replace element-wise
assert_close (meaningless when the max tail is 4.0) with distribution
checks (mean + median). Thresholds calibrated to empirical observation:
last_hidden_state (post-norm): mean < 0.025, median < 0.020
hidden_states[-2] (pre-norm): mean < 0.120, median < 0.100
Headroom is ~1.6x over observed (0.0152, 0.0117, 0.0754, 0.0625).
fp32 path keeps strict assert_close at atol=1e-4.
Also revert `torch_dtype=dtype` -> `dtype=dtype`: transformers 4.57.3
emits a deprecation warning recommending `dtype=`. The earlier switch
(Copilot review) was wrong-direction for the pinned transformers version.
The two config-key strips that mention `torch_dtype` stay — they're
stripping the legacy serialized config key, not the constructor kwarg.
PORT_STATUS.md updated: pinned Z-Image SHA, pinned HF id (Z-Image-Turbo),
empirical bf16 signature recorded in the decisions table, 35-layer
correction reflected in the parity table.
Diagnostic-only test that prints abs-diff (max/mean/median/p99) for each
of the 25 hidden states (embedding + 24 layer outputs + post-norm) in
bf16. Lets a reviewer distinguish:
- Healthy bf16 tail: monotonic growth per layer, median ≪ atol
- Real op mismatch: single-layer spike where one layer's diff is
>>2x the previous
Always passes; output is captured via `-s` and inspected manually.
Run after the parametrized parity test (which currently fails on the
strict mean-drift threshold) to decide whether to relax thresholds or
investigate a specific layer.
scheduler:
- Restore numpy default (float64) on the non-reference timestep linspace
path. The earlier `dtype=np.float32` silently changed rounded timestep
values for every model already using this scheduler. Float32 is now
only applied on the new `use_reference_discrete_timesteps=True` branch.
- Document `use_reference_discrete_timesteps` in the class docstring with
default + when-to-enable + backward-compat notes.
qwen3 encoder:
- `torch.arange(...)` for position_ids now passes `dtype=torch.long`
explicitly so embedding/RoPE indices are unambiguous.
- Drop the redundant `isinstance(x.device.type, str)` check; `device.type`
is always a str. Keep the mps -> cpu fallback for torch.autocast.
encoder parity test:
- `AutoModel.from_pretrained` uses `torch_dtype=` to match repo convention
(the rest of the repo and existing transformers callers all use this
keyword). transformers 4.57.3 accepts both, but `torch_dtype` keeps the
test green across the older transformers versions that some downstream
forks pin.
Required by `.agents/skills/add-model/contracts/port_state.md` and
`add-model-10-pr-review` "Quality and evidence lane" before any handoff.
README captures: reference assets, weight-dir layout, run commands,
per-component status, and live blockers (lm_head allowlist, scheduler
config-flag pin requirement, tokenizer `text_len` reconciliation).
PORT_STATUS captures: component matrix, parity commands + last result,
open questions (pin Z-Image clone SHA, pin HF id), issues I001-I005
(lm_head allowlist resolved, scheduler-config pin still open, transformer
+ pipeline NOT_STARTED), and the bf16 / scheduler-config-widening
decisions made in this PR.
The fixture only forwarded `num_train_timesteps`, `shift`, and
`use_dynamic_shifting` from `scheduler_config.json`. Any future field in
the on-disk config (e.g. `time_shift_type`, `invert_sigmas`, etc.) would
be silently dropped on both sides — parity would pass while the produced
schedulers diverged from the file Z-Image actually ships.
Spread the full dict (minus Diffusers loader keys) into the constructor
on both sides. The dynamic-shifting test follows the same shape.
All sibling family dirs (`encoders/`, `pipelines/`, `transformers/`,
`vaes/`, `sd35/`) ship an `__init__.py`. Z-Image was the only port-in-
progress dir missing one; conftest setups that walk packages will skip
the dir without it.
The original test forced fp32 because the strict `atol=1e-4` failed in bf16
across 24 transformer blocks. Per add-model-02-parity's calibration block,
deep encoders in bf16 accumulate per-GEMM epsilon into a long tail — the
right fix is to keep both dtypes covered, log per-batch
max/mean/median/p99 diagnostics, and require median≈0 + abs-mean drift
below threshold so a real bug (mean ≫ atol) cannot pass.
Also: assert that the unexpected-key surface from the production loader
equals `Qwen3Model.ALLOWED_UNEXPECTED_KEYS` so future `Qwen3ForCausalLM`
checkpoint changes are caught. Fix stale docstring filename. Fix typo
`torch.bloat16` (removed entirely with the dtype rewrite).
SolitaryThinker