Scope one-round FP32 SwiGLU to the exact Qwen3.5 contract - #51
Scope one-round FP32 SwiGLU to the exact Qwen3.5 contract#51kiddyboots216 wants to merge 4 commits into
Conversation
Broly Security ScanNote ✅ Clean scan Note Re-scan this PR anytime with
|
Lane-doctrine gate (trainability-trap class): the exact-contract Qwen3_5MLP forward/backward must produce non-None finite nonzero grads for every parameter and the input, with forward bytes equal to the one-round reference program and grads matching the autograd reference trajectory. Complements the op-level backward test shipped with the cherry-pick.
Keep the historical two-round fused SwiGLU program as the default and expose the one-round FP32 program through an exact-only operator. Qwen3.5 dense and MoE MLPs select it only under the resolved exact contract, preserving non-exact caller bytes. Conventional scoping, policy, operator, and gradient-engagement tests cover both sides of the dispatch.
020a338 to
e7beb67
Compare
459f8a5 to
51f9665
Compare
Model resolution now stamps config._exact_contract_family (the shared family vocabulary: glm52, qwen3_5_dense, qwen3_5_moe, or None) and config._exact_one_round_swiglu alongside the existing family-branded flags. The Qwen3.5 MLPs key their one-round FP32 SwiGLU selection off the neutral stamp, falling back to the legacy _qwen35_exact_contract flag for configs that predate stamping, so every existing construction path resolves identically. GLM stays on its certified two-operation activation; adopting one-round there remains gated on its own byte evidence.
51f9665 to
4d2688b
Compare
Complete the trainer half of serving's universal one-round switch (xorl-sglang f10b907d8) for the GLM exact contract. Serving's exact mode resolves SiluAndMul.forward_exact to fp32_silu_and_mul whenever rl_on_policy_target is set, while the trainer's GLM exact sites still computed the two-operation program with a BF16 intermediate store after SiLU. Adopt exact_fp32_silu_and_mul at the GLM exact dense and shared-expert value sites, including the active-LoRA value and VJP paths. Routed experts stay on the raw silu_and_mul kernel serving did not change. Fail closed on non-SiLU GLM exact configs at both canonical model admission and stamped MLP construction, and retain the pinned sampler's v2 local-partial policy identifier. Bump the TP16 shared-expert and TP1 dense-MLP active-LoRA contracts from v1 to v2. Gate trainer-versus-serving bytes at both GLM geometries, two-op sensitivity, concatenated separate-tensor dispatch, gradients, stamp engagement, activation refusal, and trainer-sampler policy identity.
qywu
left a comment
qywu
left a comment
There was a problem hiding this comment.
The implementation and tests here are self-consistent, but the PR description states the opposite of what the code does, which is the risky direction — a reviewer approving on the summary would be approving the inverse of what ships.
The description is falsified by this branch's own tests
The body says:
Preserve the historical two-round bytes for every non-exact caller.
[…] non-exact callers are byte-preserved by test.
tests/ops/test_fused_silu_and_mul_one_round.py, added by this PR, says:
Universal one-round FP32 SwiGLU is the intended end-state […]
fused_silu_and_mulcomputes SiLU and the multiply in fp32 and rounds ONCE, for every caller — exact and non-exact alike
and test_nonexact_qwen35_mlp_is_one_round asserts exactly the byte change the body denies.
The code agrees with the test. _silu_and_mul_kernel changes from silu_gate.to(gate.dtype) * up to silu_gate * up.to(tl.float32) with no exact-contract gate, and silu_and_mul() gains an unconditional _fp32_silu_and_mul early return for anything that isn't contiguous bf16/fp16 on SM 9.0.
Measured on an H100 (SM 9.0, the admitted path), 4096×4096 bf16, fused_silu_and_mul output vs. the historical two-round program:
this branch 27.44% of output elements differ
#52 0.00% (byte-identical; adds gated exact_fp32_silu_and_mul instead)
fused_silu_and_mul is the shared entry point for Llama3, OLMo2, Qwen2, Qwen3, Qwen3-MoE, Qwen3.5, Qwen3.5-MoE, GLM4-MoE, GLM5, and DeepSeek-V3. This changes training numerics for all of them.
Separately, the hardcoded (major, minor) != (9, 0) gate in _use_fp32_fused_swiglu means the numerics differ again on any non-Hopper device, which seems worth an explicit decision given the FA4/MXFP4/native-FP8 paths elsewhere in the stack.
Redundant with #43, superseded by #52
#43 already lands this same global change (measured identically at 27.44%), and #52 implements the scoped design this PR's body describes: fused_silu_and_mul byte-identical to base, with a separate exact_fp32_silu_and_mul behind _qwen35_exact_contract. So the stack currently carries two contradictory intents for the same operator; whichever merges second silently overwrites the other.
Also
pre-commit run --all-filesfailsruff-formaton 5 files at this head (4 of them added here). The body says "Changed Python files: Ruff passed." Lint CI does not run on this PR becauselint.ymltriggers only onbranches: [main].- This branch sits on the diverged #49 base — see my review there.
Requested changes
Simplest resolution is to close this in favour of #52's scoping and keep #43's global change as the single place that decision is made and documented. If the universal one-round program is genuinely the intent, that's fine — but the body needs to say so, and the blast radius across the ten model families above should be called out explicitly rather than described as byte-preserving.
4d2688b to
d19cd28
Compare
|
Superseded by #57, which consolidates the complete exact trainer-serving stack directly onto main. |
2 participants
Scope one-round FP32 SwiGLU to the exact Qwen3.5 contract
Stack position: after the radix weight-sync fix.
Problem
The serving side of the exact Qwen3.5 pair evaluates SwiGLU with a one-round
FP32 program, while the trainer side still used the historical two-round
program. The first divergent tensor was the layer-0 MLP activation, producing
nonzero K3 on the otherwise exact dense pair.
Change
paths.
MLP sites that share the serving value program.
tests.
Validation
at the restacked head.
byte-identical, K3 = 0.0, with all 149 compared components bitwise equal.
byte-preserved by test.
This PR contains production code and conventional tests only. Private replay
captures and diagnostic dumps are not included.
Generality follow-up (2026-08-13)
Model resolution now stamps the family-neutral
_exact_contract_familyand_exact_one_round_swiglukeys alongside the compatibility flags. Qwen construction reads the neutral stamp, with a legacy fallback for direct configs that predate resolution.GLM-5.2 exact dense and shared-expert trainer sites now use the one-round FP32 SwiGLU program already selected by exact serving, including active-LoRA value and VJP paths. Routed experts remain on their unchanged serving-paired program; exact non-SiLU configs fail closed; the local-partial policy token remains v2 to match
GLM52_SAMPLER_LOCAL_POLICY. Trainer-versus-serving outputs are bitwise identical at both GLM geometries, while the former two-operation program differs on 18,036/65,536 and 215,706/786,432 elements.