[PyTorch] LayerNormLinear drops gradients from gathered LayerNorm output under TP/SP

# Fork of the repro in https://github.com/NVIDIA-NeMo/Megatron-Bridge/issues/6453.

import math
import os

# Set before importing torch/TE so TE's cuBLAS calls also avoid TF32.
os.environ["NVIDIA_TF32_OVERRIDE"] = "0"

import torch
import torch.distributed as dist
import transformer_engine
import transformer_engine.pytorch as te


S, BATCH, H, O = 64, 2, 256, 512
TP = 2
EPS = 1e-5
TOL = 1e-3  # Same threshold used to flag the input-gradient error in the issue.
COLUMNS = ("dx", "dy", "da", "db")


def all_gather(tensor, dim):
    parts = [torch.empty_like(tensor) for _ in range(TP)]
    dist.all_gather(parts, tensor.contiguous())
    return torch.cat(parts, dim=dim)


def rel(got, ref):
    return ((got - ref).abs().max() / ref.abs().max()).item()


class TwoProjections(torch.nn.Module):
    def __init__(
        self, *, tensor_parallel, sequence_parallel, return_layernorm_output_gathered, device
    ):
        super().__init__()
        self.sequence_parallel = sequence_parallel
        self.return_layernorm_output_gathered = return_layernorm_output_gathered
        common = dict(
            tp_group=dist.group.WORLD if tensor_parallel else None,
            tp_size=TP if tensor_parallel else 1,
            parallel_mode="column" if tensor_parallel else None,
            sequence_parallel=sequence_parallel,
            params_dtype=torch.float32,
            fuse_wgrad_accumulation=False,
            bias=False,
            device=device,
        )
        self.fused_layernorm_linear = te.LayerNormLinear(
            H,
            O,
            eps=EPS,
            normalization="LayerNorm",
            return_layernorm_output=True,
            return_layernorm_output_gathered=return_layernorm_output_gathered,
            **common,
        )
        self.fc1 = te.Linear(H, O, **common)

    def forward(self, x):
        base, ln = self.fused_layernorm_linear(x)
        if self.sequence_parallel and self.return_layernorm_output_gathered:
            # [S,B,H] -> [S/TP,B,H]. Linear(SP=True) will gather this once.
            ln = ln.chunk(TP, dim=0)[dist.get_rank()].contiguous()
        return base + self.fc1(ln)


@torch.no_grad()
def initialize_and_copy(reference, model, tensor_parallel, branch_scale, device):
    """Initialize full tensors once, then copy exact TP slices into the test."""
    rng = torch.Generator(device=device).manual_seed(7)
    rank = dist.get_rank()
    target = dict(model.named_parameters())
    for name, param in reference.named_parameters():
        if name.endswith("layer_norm_weight"):
            param.normal_(mean=1.0, std=0.1, generator=rng)
        elif name == "fc1.weight":
            # Vary only the second projection's magnitude, like the issue.
            param.normal_(mean=0.0, std=1.0, generator=rng).mul_(branch_scale)
        else:
            param.normal_(mean=0.0, std=0.1, generator=rng)
        dist.broadcast(param, src=0)
        if tensor_parallel and name in ("fused_layernorm_linear.weight", "fc1.weight"):
            # Column parallelism partitions weight [O,H] along O.
            source = param.chunk(TP, dim=0)[rank]
        else:
            # LayerNorm gamma and beta are replicated across TP ranks.
            source = param
        target[name].copy_(source)


def check(tensor_parallel, sequence_parallel, return_layernorm_output_gathered, branch_scale, device):
    rank = dist.get_rank()
    # Always return local LN in the TP=1 reference. This avoids the gathered
    # flag's backward path in TE versions affected by the reported issue.
    reference = TwoProjections(
        tensor_parallel=False,
        sequence_parallel=False,
        return_layernorm_output_gathered=False,
        device=device,
    )
    model = TwoProjections(
        tensor_parallel=tensor_parallel,
        sequence_parallel=sequence_parallel,
        return_layernorm_output_gathered=return_layernorm_output_gathered,
        device=device,
    )
    initialize_and_copy(reference, model, tensor_parallel, branch_scale, device)

    rng = torch.Generator(device=device).manual_seed(42)
    x_full = torch.randn(S, BATCH, H, device=device, generator=rng)
    gout_full = torch.randn(S, BATCH, O, device=device, generator=rng)
    dist.broadcast(x_full, src=0)
    dist.broadcast(gout_full, src=0)

    # SP shards the sequence; column TP shards output features.
    x_local = x_full.chunk(TP, dim=0)[rank] if sequence_parallel else x_full
    gout = gout_full.chunk(TP, dim=-1)[rank].contiguous() if tensor_parallel else gout_full
    x = x_local.clone().requires_grad_(True)
    y = model(x)
    y.retain_grad()
    y.backward(gout)

    # All ranks participate in collectives before rank 0 compares results.
    # a = fused_layernorm_linear.weight; b = fc1.weight.
    # dy is the supplied upstream gradient, so it checks output-gradient slicing.
    dx = all_gather(x.grad, dim=0) if sequence_parallel else x.grad
    dy = all_gather(y.grad, dim=-1) if tensor_parallel else y.grad
    da = model.fused_layernorm_linear.weight.grad
    db = model.fc1.weight.grad
    if tensor_parallel:
        da = all_gather(da, dim=0)
        db = all_gather(db, dim=0)

    if rank == 0:
        xr = x_full.clone().requires_grad_(True)
        yr = reference(xr)
        yr.retain_grad()
        yr.backward(gout_full)
        errors = {
            "dx": rel(dx, xr.grad),
            "dy": rel(dy, yr.grad),
            "da": rel(da, reference.fused_layernorm_linear.weight.grad),
            "db": rel(db, reference.fc1.weight.grad),
        }
        bad = [k for k in COLUMNS if not math.isfinite(errors[k]) or errors[k] > TOL]
        status = "  <-- mismatch: " + ", ".join(bad) if bad else "  OK"
        print(
            f"{TP if tensor_parallel else 1:>2} {str(sequence_parallel):>5} "
            f"{str(return_layernorm_output_gathered):>32} {branch_scale:>8g} | "
            + " ".join(f"{errors[k]:10.2e}" for k in COLUMNS)
            + status,
            flush=True,
        )
        return errors


def main():
    if int(os.environ.get("WORLD_SIZE", "1")) != TP:
        raise SystemExit("Run with torchrun --standalone --nproc_per_node=2 <script.py>")
    if not torch.cuda.is_available():
        raise SystemExit("This reproduction requires two CUDA GPUs and CUDA-enabled TE.")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False
    dist.init_process_group("nccl")
    try:
        rank = dist.get_rank()
        if rank == 0:
            print(
                f"PyTorch {torch.__version__}; TE {transformer_engine.__version__}; "
                f"GPU {torch.cuda.get_device_name(device)}",
                flush=True,
            )
            print("Baseline: TP=1, SP=False; FP32")
            print("Relative max error = max(abs(got-ref)) / max(abs(ref))")
            print(
                f"{'TP':>2} {'SP':>5} {'return_layernorm_output_gathered':>32} {'W scale':>8} | "
                + " ".join(f"{name:>10}" for name in COLUMNS),
                flush=True,
            )
        for tensor_parallel, sequence_parallel, return_layernorm_output_gathered, branch_scale in [
            (False, False, False, 0.1),  # Unsharded control.
            (True, True, False, 0.1),   # TP + SP, local returned LN.
            (True, True, True, 0.1),    # TP + SP, gathered returned LN.
            (True, True, True, 0.01),
            (True, True, True, 0.001),
            (True, True, True, 0.0),    # No input-gradient contribution from fc1.
        ]:
            check(
                tensor_parallel, sequence_parallel, return_layernorm_output_gathered,
                branch_scale, device,
            )
    finally:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()

diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -685,6 +685,24 @@
             )
             nvtx_range_pop(f"{nvtx_label}.fsdp_gather")
 
+            grad_ln_out = grad_outputs[1] if ctx.return_layernorm_output else None
+            ln_grad_needs_reduce_scatter = (
+                grad_ln_out is not None
+                and ctx.return_layernorm_output_gathered
+                and ctx.parallel_mode == "column"
+                and ctx.sequence_parallel
+                and ctx.tp_size > 1
+            )
+            if ln_grad_needs_reduce_scatter:
+                # Combine both full-sequence gradients before reducing them. The
+                # overlap kernels can reduce dgrad inside GEMM, before this sum
+                # can be formed, so use the explicit communication path here.
+                ctx.ub_overlap_ag = False
+                ctx.ub_overlap_rs_dgrad = False
+                ctx.ub_bulk_dgrad = False
+                ctx.ub_bulk_wgrad = False
+                ctx.grad_input_quantizer = None
+
             # Configure Userbuffers communication (comm+GEMM overlap)
             ctx.ub_obj_gradout = None
             ub_obj_dgrad = None
@@ -901,6 +919,8 @@
                 nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad")
                 dgrad = gemm_out
                 if ctx.sequence_parallel:
+                    if ln_grad_needs_reduce_scatter:
+                        dgrad = dgrad + grad_ln_out.reshape_as(dgrad)
                     dgrad, dgrad_work = reduce_scatter_along_first_dim(
                         dgrad,
                         ctx.tp_group,
@@ -1131,8 +1151,10 @@
 
             # Residual gradient
             dgrad = dgrad.view(inputmat.shape)
-            if ctx.return_layernorm_output and not ctx.return_layernorm_output_gathered:
-                dgrad = dgrad + grad_outputs[1].view_as(dgrad)
+            if grad_ln_out is not None and not ln_grad_needs_reduce_scatter:
+                # A local returned output already has a local gradient. This
+                # also covers the gathered flag when no gather actually ran.
+                dgrad = dgrad + grad_ln_out.reshape_as(dgrad)
 
             # Norm gradient
             dgamma = None