From e6ae346415813945a4778440adff0359b96f6a63 Mon Sep 17 00:00:00 2001 From: Gavin Qi Date: Wed, 7 Oct 2026 20:43:03 -0700 Subject: [PATCH] fix(qwen3.5): average TP-replicated linear-attention gradients across TP Qwen3.5's linear-attention layers (linear_attn and input_layernorm) are replicated on every tensor-parallel rank rather than sharded, but nothing synchronizes their gradients across TP. Each rank computes the same full gradient mathematically, yet nondeterministic kernels make the replicas drift apart a little more every step, so TP ranks end up training slightly different weights and a checkpoint keeps only one of them. Mark these parameters with average_gradients_across_tp_domain so Megatron's finalize_model_grads all-reduces their gradients across the TP group. --- slime_plugins/models/qwen3_5.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/slime_plugins/models/qwen3_5.py b/slime_plugins/models/qwen3_5.py index d7551f9a94..8c28a324af 100644 --- a/slime_plugins/models/qwen3_5.py +++ b/slime_plugins/models/qwen3_5.py @@ -186,6 +186,13 @@ def __init__( self.input_layernorm = RMSNorm(self.hf_config.hidden_size, eps=self.hf_config.rms_norm_eps) + # linear_attn and input_layernorm are replicated, not sharded, across TP ranks. + # Average their gradients across TP so the replicas stay identical; without this, + # nondeterministic kernels let them drift apart a little more every step. + for module in (self.linear_attn, self.input_layernorm): + for param in module.parameters(): + param.average_gradients_across_tp_domain = True + def hf_forward(self, hidden_states, packed_seq_params): hidden_states = self.input_layernorm(hidden_states) hidden_states = self.linear_attn(