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

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

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)

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

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

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

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

"""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()