fix(qwen3.5): average TP-replicated linear-attention gradients across TP

[gavinqii]

Summary

In the Qwen3.5 spec, slime_plugins/models/qwen3_5.py adds the linear-attention layers (linear_attn = Qwen3_5GatedDeltaNet, plus input_layernorm) as a HuggingfaceAttention. These parameters are replicated on every tensor-parallel rank, not sharded, but nothing synchronizes their gradients across TP.

hf_attention.py gathers the sequence-parallel input before these modules, so in exact arithmetic every TP rank computes the same full gradient. In practice the kernels are not bitwise deterministic across ranks. The replicas drift apart a little more every step, so TP ranks end up training slightly different weights, and a distributed checkpoint keeps only one of the copies.

This PR marks these parameters with average_gradients_across_tp_domain = True. Megatron's finalize_model_grads (_allreduce_non_tensor_model_parallel_grads) then averages their gradients across the TP group, so the replicas stay identical. With TP=1 it is a no-op: that path returns early when the TP group size is 1.

Evidence

A/B on Qwen3.5-9B SFT, 2 nodes × 8 B200, TP=2 / PP=2, 30 training steps, same data and seed. The only difference is the gradient flag. A test-only probe compared each replicated parameter across the 2 TP ranks after every step (all-reduce MAX − MIN):

step	without this change	with this change
0 (after init)	0 / 120 params differ	0 / 120
2	74 / 120 differ, max abs diff 3.1e-5	0 / 120
5	82 / 120, max abs 1.2e-4	0 / 120
30	85 / 120, max abs 2.4e-4 (max rel 0.19%)	0 / 120

We first saw this in a strict save/restore check that compares per-rank tensor fingerprints (Qwen3.5-35B, TP2). At save time, the source TP ranks already held different values for these replicated weights. The checkpoint stores a single copy, so the restored TP1 no longer matched the source TP1. In a 122B A/B, the extra all-reduce cost +0.24% step time.

These runs used a slime 0.3.0-based build. On main, the Attention class touched here differs only by a new name argument, and the sequence-parallel gather/scatter in hf_attention.py is the same. I haven't rerun the A/B on main.

Notes

• qwen3_next.py builds its linear attention the same way and may need the same change. I haven't measured it, so it's not included here.
• No new test: reproducing the drift needs a multi-GPU TP run. Happy to add a test if there's a preferred harness for this.

[gavinqii]

Additional evidence, still on a slime 0.3.0-based build (P6 B200, 10 steps each).

This change on four Qwen3.5 layouts. The flag is set in Attention.__init__, and the probe only measures. Counts are summed over all ranks.

Layout	replicated params / flagged / differing across TP (step 10)
35B-A3B TP2 PP1 EP8 SP	300 / 300 / 0
122B-A10B TP2 PP4 EP8 SP	360 / 360 / 0
122B-A10B TP4 PP4 EP8 SP	2880 / 2880 / 0
397B-A17B TP2 PP8 VPP2 EP8 SP	3600 / 3600 / 0

Numerics, 35B-A3B. Same seed and data: two runs with this change (F1, F2) and one without it (U).

pair	loss max / mean abs diff	grad_norm max / mean abs diff
F1 vs F2 (run-to-run noise)	1.65e-3 / 7.5e-4	1.08 / 0.17
F1 vs U	2.18e-3 / 4.5e-4	0.51 / 0.07
F2 vs U	1.22e-3 / 6.0e-4	1.11 / 0.15

Without the change, 1656 of 2400 replicated rank-params differ across TP by step 10. With it, none do. Step time is unchanged (72.4 / 71.9 s vs 71.9 s).

CI. The only failing check, test_swe_workspace_diff.py::test_patch_uses_actual_image_baseline_without_changing_git_index[unstaged], looks like a flaky race in the test added in #2457. Git's background maintenance removed .git/objects/maintenance.lock during copytree. The test passes on main at the same commit and is unrelated to this change, which only touches qwen3_5.py. Could a maintainer re-run it?

Ready for review whenever someone has time.