GDN input RMSNorm gradients are multiplied by TP size during Qwen3.8-27B training

Read src/mcore_bridge/model/modules/gated_delta_net.py and trace the sequence-parallel toggling for the combined in_proj; then inspect Megatron-Core's _allreduce_non_tensor_model_parallel_grads finalizer and the cited Transformer Engine layer norm backward path. Run the provided four-rank reproduction with the listed dependencies and compare the gamma gradient before and after finalization. Done when TP finalization leaves it equal to the unsharded reference instead of multiplying it by TP size.

Checklist / 检查清单

• I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

Environment and observed behavior

While training Qwen3.8-27B with Megatron-SWIFT and tensor parallelism enabled, we found a duplicated TP reduction of the input RMSNorm gradients in GDN's combined in_proj. The observed training configuration uses TP4, CP2, sequence parallelism enabled, BF16, and linear_decoupled_in_proj=False.

Component	Version
Megatron-SWIFT (ms-swift)	4.6.0
mcore-bridge	1.6.5
Megatron-Core	0.18.2
Transformer Engine	2.20.2
PyTorch	2.13.0+cu130

The affected parameter is self_attention.in_proj.layer_norm_weight, the gamma of the fused input RMSNorm. Each TP rank already has the complete gradient before TP finalization, but the finalizer sums those complete copies again. With TP4, the resulting gamma gradient is exactly 4 times the correct value. This affects optimizer updates and global gradient clipping.

Bug chain

The following code summaries omit unrelated statements. All links point to the versions listed above.

1. TE marks the gamma parameter as sequence-parallel when the layer is initialized.

Transformer Engine v2.20.2 — transformer_engine/pytorch/module/layernorm_linear.py, L1651–L1655:

setattr(self.layer_norm_weight, "sequence_parallel", self.sequence_parallel)

2. Bridge gathers the sequence externally and disables SP on the projection modules, while leaving the gamma parameter's SP marker unchanged.

mcore-bridge v1.6.5 — src/mcore_bridge/model/modules/gated_delta_net.py, L215–L224:

hidden_states = gather_from_sequence_parallel_region(
    hidden_states, tensor_parallel_output_grad=False, group=tp_group,
)
saved_linear_sp = self._set_linear_sequence_parallel(False)

The helper at L165–L175 only changes module attributes:

for sub in module.modules():
    if hasattr(sub, 'sequence_parallel'):
        saved.append((sub, sub.sequence_parallel))
        sub.sequence_parallel = enabled

3. With module SP disabled, TE all-reduces the projection's input gradient before RMSNorm backward. The resulting gamma gradient is therefore already complete on every TP rank.

Transformer Engine v2.20.2 — layernorm_linear.py, L900–L910:

if ctx.sequence_parallel:
    dgrad, dgrad_work = reduce_scatter_along_first_dim(...)
else:
    dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True)

The same file, L1124–L1159:

if dgrad_work is not None:
    dgrad_work.wait()
# ...
dgrad, dgamma = tex.rmsnorm_bwd(dgrad, inputmat, rsigma, ln_weight, ...)

4. MCore sees the parameter's unchanged SP marker and sums the complete gamma gradients again.

Megatron-Core core_v0.18.2 — megatron/core/distributed/finalize_model_grads.py, L403–L425, simplified to the relevant branch:

elif config.sequence_parallel and getattr(param, "sequence_parallel", False):
    # ...
    grads_sum.append(grad.data)
# grads_sum is subsequently all-reduced with ReduceOp.SUM.

How to Reproduce / 如何复现

Minimal reproduction

This isolates the combined in_proj SP wrapper using Bridge's original SP-toggle methods, a real TE layer, and the real MCore TP finalizer. It does not construct a full GDN layer or require model weights or a dataset. An independent, unsharded PyTorch RMSNorm + Linear serves as the TP1 mathematical reference.

The script runs in FP32 with TF32 disabled to avoid confusing TE GEMM precision differences with the reduction error. Save it as repro_gdn_gamma.py and run:

NVIDIA_TF32_OVERRIDE=0 CUDA_DEVICE_MAX_CONNECTIONS=1 \
  torchrun --standalone --nproc_per_node=4 repro_gdn_gamma.py

"""Isolate Bridge v1.6.5's combined-in_proj SP wrapper; no full GDN/model needed."""
import os
from types import SimpleNamespace

assert os.environ.get("NVIDIA_TF32_OVERRIDE") == "0", "Set this before importing torch"
import torch
import torch.distributed as dist
import torch.nn.functional as F

torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
dist.init_process_group("nccl")

from megatron.core import parallel_state as ps
from megatron.core.distributed.finalize_model_grads import _allreduce_non_tensor_model_parallel_grads
from megatron.core.extensions.transformer_engine import TELayerNormColumnParallelLinear
from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region
from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
from megatron.core.transformer import TransformerConfig
from mcore_bridge.model.modules.gated_delta_net import GatedDeltaNet

tp = dist.get_world_size()
assert tp in (2, 4)
ps.initialize_model_parallel(tensor_model_parallel_size=tp)
model_parallel_cuda_manual_seed(1234)
group, rank = ps.get_tensor_model_parallel_group(), ps.get_tensor_model_parallel_rank()
torch.backends.cuda.matmul.allow_tf32 = False

S, H, O, eps = 128, 128, 512, 1e-6
cfg = TransformerConfig(
    num_layers=1, hidden_size=H, num_attention_heads=4, ffn_hidden_size=256,
    tensor_model_parallel_size=tp, sequence_parallel=True, normalization="RMSNorm",
    layernorm_epsilon=eps, params_dtype=torch.float32, add_bias_linear=False,
    gradient_accumulation_fusion=False, tp_comm_overlap=False,
)
model = torch.nn.Module()
model.ddp_config = SimpleNamespace(use_megatron_fsdp=False)
model.in_proj = TELayerNormColumnParallelLinear(
    H, O, config=cfg, init_method=torch.nn.init.zeros_, gather_output=False,
    bias=False, skip_bias_add=False, is_expert=False, tp_group=group,
)
layer = model.in_proj
x = ((torch.arange(S * H, device="cuda").reshape(S, 1, H) % 97 - 48) / 97).float()
w = ((torch.arange(O * H, device="cuda").reshape(O, H) % 89 - 44) / 89 / 8).float()
gamma = (0.8 + (torch.arange(H, device="cuda") % 17) / 50).float()
dy = ((torch.arange(S * O, device="cuda").reshape(S, 1, O) % 43 - 21) / 43).float()
sl, ol = slice(rank * S // tp, (rank + 1) * S // tp), slice(rank * O // tp, (rank + 1) * O // tp)
with torch.no_grad():
    layer.weight.copy_(w[ol])
    layer.layer_norm_weight.copy_(gamma)

# Independent, unsharded RMSNorm + Linear reference on the same tensors.
xr, wr, gr = [t.clone().requires_grad_() for t in (x, w, gamma)]
yr = F.linear(xr * torch.rsqrt(xr.square().mean(-1, keepdim=True) + eps) * gr, wr)
gx, gw, gg = torch.autograd.grad((yr * dy).sum() / (S * O), (xr, wr, gr))

def rel(a, b):
    a, b = a.detach().double(), b.detach().double()
    return float((a - b).norm() / b.norm().clamp_min(1e-30))

def ratio(a, b):
    a, b = a.detach().double(), b.detach().double()
    return float((a * b).sum() / b.square().sum())

xs = x[sl].clone().requires_grad_()
full = gather_from_sequence_parallel_region(xs, tensor_parallel_output_grad=False, group=group)
saved = GatedDeltaNet._set_linear_sequence_parallel(model, False)
try:
    y, _ = layer(full)
finally:
    GatedDeltaNet._restore_linear_sequence_parallel(model, saved)
assert layer.sequence_parallel and layer.layer_norm_weight.sequence_parallel
(y * dy[..., ol]).sum().div(S * O).backward()
for p in model.parameters():
    p.main_grad = p.grad.detach().float().clone()
pre = layer.layer_norm_weight.main_grad.clone()
_allreduce_non_tensor_model_parallel_grads([model], cfg, tp_group=group)
post = layer.layer_norm_weight.main_grad
expected = tp * gg
errors = [rel(y, yr[..., ol]), rel(xs.grad, gx[sl]), rel(layer.weight.grad, gw[ol]),
          rel(pre, gg), rel(post, expected)]
assert all(e < 3e-4 for e in errors), (rank, errors)
if rank == 0:
    print(f"TP={tp}: gamma_pre/ref={ratio(pre, gg):.6f}, "
          f"gamma_post/ref={ratio(post, gg):.6f}; duplicate SUM reproduced", flush=True)

ps.destroy_model_parallel()
dist.destroy_process_group()

Expected reproduction output for TP4, allowing normal floating-point error:

TP=4: gamma_pre/ref=1.000000, gamma_post/ref=4.000000; duplicate SUM reproduced

The assertions compare complete vectors, not just gradient norms: output, input gradient, and sharded projection-weight gradient match the reference; gamma matches before finalization and becomes TP × reference afterward. With two processes, the expected multiplier is 2.

Expected behavior: the gamma gradient should remain equal to the unsharded reference after TP finalization, without an extra TP-size multiplier.

Validation scope: the original TP4/CP2 mathematical oracle and full Qwen3.8-27B single-node/two-node training checks confirmed the repeated SUM. This shortened TP-only script has passed syntax and source review but has not yet been run separately on GPUs.

Additional Information / 补充信息

No response