[faresobeid]

Summary

Qwen3.8 Flash Next cloned three Qwen3.5 blocks. They now live once in trainer/models/layers/, and both archs use them:

• layers/gated_delta_net.py::GatedDeltaNet (was Qwen3_5GatedDeltaNet and Qwen3.8 GatedDeltaNet). It takes keyword args. Qwen3.5 passes hidden_act/output_gate_type from its config. Qwen3.8 passes "silu"/"sigmoid", the values it hardcoded before.
• layers/norms.py::ZeroCenteredRMSNorm (was Qwen3_5RMSNorm and Qwen3.8 RMSNorm).
• layers/mlp.py::SigmoidGatedFeedForward (was Qwen3_5SharedExpert and SigmoidOutputGatedFeedForward).

The per-arch copies are deleted (qwen3_5/{gated_delta_net,norm}.py, qwen3_8_flash_next/{gated_delta_net,norm}.py). State-dict keys, shapes and dtypes are unchanged for qwen3_5, qwen3_5_moe, the qwen3_5 VLM and qwen3_8.

CP conv cu_seqlens: Qwen3.5 passed cu_seqlens to the CP conv call and Qwen3.8 did not. Neither one is a bug. When cp_context is set, FLA's causal_conv1d returns early through causal_conv1d_cp and ignores cu_seqlens (fla 0.5.2, fla/modules/conv/causal_conv1d.py). The shared module keeps the Qwen3.5 single call. Qwen3.8 also passed initial_state=None, output_final_state=False explicitly. Those are the defaults, so the calls are the same.

Behavior change

• Qwen3.5 (dense, MoE, VLM): none (bitwise on CPU, see below).
• Qwen3.8: forward is bitwise identical. Two small differences:
  • The shared expert's input gradient can differ in the last fp32 bit. The old Qwen3.8 code computed sigmoid(gate) * ffn and Qwen3.5 computes ffn * sigmoid(gate). The product is the same, but autograd adds the two input-gradient terms in a different order. The shared module keeps the Qwen3.5 order. In a full CPU model, the largest weight-grad difference was about 3e-8.
  • Inside linear_attn, Qwen3.8 now registers parameters in the Qwen3.5 order (conv1d first). The keys are the same. Only the order of the state-dict / named_parameters() output changes, and the random init for a given seed. Models loaded from a checkpoint are not affected.

What I ran

• CPU unit/model equivalence vs origin/main (equiv_qwen38blocks.py). FLA kernels were replaced by pure-torch reference conv / gated-delta-rule / gated-norm, and every kernel call's kwargs were recorded. CP was faked with a stub context. Results:
  • qwen3_5 and qwen3_5_moe: all norms, shared experts and GDNs (CP and non-CP), plus the full-model logits and all grads: bitwise identical.
  • qwen3_5 VLM: keys and units identical.
  • Standalone norm in fp32/bf16/fp16: bitwise identical.
  • qwen3_8: everything is bitwise identical except the shared-expert input grad described above. With the old multiply order the full model is bitwise identical, so that order is the only difference.
• /tmp/prm-equiv/equiv.py check (CPU archs): all IDENTICAL.
• pytest tests/unit/train tests/unit/test_configs.py tests/unit/utils/test_weights.py -m "not gpu": 356 passed. The 2 test_cpu_adam failures are a missing ninja in my env; they also fail on main.
• ruff check / format: clean.

GPU checks (1×8 H200, SLURM job 2386)

• qwen3_5 and qwen3_5_moe, branch vs. main on GPU with the real FLA kernels: bitwise identical (forward, logprobs, entropy, all grads, buffers). main vs. main is also identical.
• qwen3_8 cannot be built in the equivalence script on this node (prime_kernels has no indexed_attention).
• -m gpu tests/unit/train/models/qwen3_5, test_qwen3_8_flash_next.py and test_cp_mamba.py: 11 passed. On main, the first two give 9 passed; test_cp_mamba was not part of the main run.

Conflicts

#3833 rewrites the qwen3_5 / qwen3_8 modeling files. This PR touches only their imports and constructor calls, so expect a small rebase conflict there.

---

Note

Medium Risk
Touches core Qwen3.5/Qwen3.8 decoder paths (linear attention, norms, MoE shared expert) but is largely a refactor with stated bitwise equivalence for Qwen3.5 and near-identical Qwen3.8 behavior.

Overview
Consolidates duplicated Qwen3.5 and Qwen3.8 building blocks into shared trainer/models/layers/ modules and rewires both architectures to import them.

GatedDeltaNet replaces per-model linear-attention copies (Qwen3_5GatedDeltaNet and the Qwen3.8 local module), with explicit constructor kwargs (activation, output_gate_activation, etc.) instead of a Qwen3.5 config object. ZeroCenteredRMSNorm and SigmoidGatedFeedForward replace arch-specific RMSNorm and sigmoid-gated shared-expert FFN implementations; Qwen3.5 drops Qwen3_5SharedExpert in favor of the shared MLP class.

Per-arch files qwen3_5/{gated_delta_net,norm}.py and qwen3_8_flash_next/{gated_delta_net,norm}.py are removed. Qwen3.8 MoE now uses SigmoidGatedFeedForward from layers (local SigmoidOutputGatedFeedForward deleted). Intended checkpoint compatibility: state-dict keys unchanged; Qwen3.5/VLM behavior unchanged; Qwen3.8 forward matches with only minor autograd ordering on shared-expert input grads and linear_attn parameter registration order for fresh inits.

^{Reviewed by Cursor Bugbot for commit 1d66345. Bugbot is set up for automated code reviews on this repo. Configure here.}