[Bug] NaN expert weight gradients at num_groups == 1 with a padded token buffer for SReLU

[GarlGuo]

Describe the bug

At num_groups == 1, the fused grouped MLP produces NaN expert weight gradients whenever the token buffer is padded past sum(split_sizes) and the cuDNN single-group kernel specialization is unavailable. num_groups == 1 selects a family of dense shortcuts that quantize and scale-swizzle the buffer as one tensor over tensor.shape[0] rows and read it back the same way, which is only coherent if the cuDNN kernels also write densely over tensor.shape[0]. That request is already guarded because it can be refused, but its siblings are gated on num_groups == 1 alone, so when the specialization is unavailable the kernels pack their output to sum(split_sizes) rows while TE still reads everything densely. Since the columnwise swizzled MXFP8 scale layout is [k/128][m/128][32][4][4] with the m-tile as an inner stride, that extent mismatch misindexes every k-tile past the first: the head rows pick up the wrong scale factors and high k-tiles read scale memory the producer never wrote. The data is read correctly and only the scales move, so the error is order 1 and usually NaN.

Steps/Code to reproduce bug

Single GPU, no Megatron, no distributed. NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 must be set before the TE import.

import os
os.environ["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] = "1"  # must precede the TE import
import torch
import transformer_engine.pytorch as te
from transformer_engine.pytorch.quantization import MXFP8BlockScaling

HIDDEN, LIVE, PAD = 1024, 512, 256


def wgrad(total_rows, fill):
    torch.manual_seed(0)
    fc1 = te.ops.GroupedLinear(1, HIDDEN, HIDDEN, bias=False, device="cuda",
                               dtype=torch.bfloat16, accumulate_into_main_grad=True)
    fc2 = te.ops.GroupedLinear(1, HIDDEN, HIDDEN, bias=False, device="cuda",
                               dtype=torch.bfloat16, accumulate_into_main_grad=True)
    model = te.ops.Sequential(fc1, te.ops.ScaledSReLU(), fc2)
    for op in (fc1, fc2):
        w = op.get_parameter("weight0")
        w.main_grad = torch.zeros(w.shape, dtype=torch.float32, device="cuda")

    gen = torch.Generator(device="cuda").manual_seed(1)
    live_x = torch.randn(LIVE, HIDDEN, generator=gen, device="cuda", dtype=torch.bfloat16)
    live_dy = torch.randn(LIVE, HIDDEN, generator=gen, device="cuda", dtype=torch.bfloat16)

    x = torch.full((total_rows, HIDDEN), fill, device="cuda", dtype=torch.bfloat16)
    dy = torch.full((total_rows, HIDDEN), fill, device="cuda", dtype=torch.bfloat16)
    x[:LIVE], dy[:LIVE] = live_x, live_dy
    x = x.detach().requires_grad_(True)

    # split_sizes covers the LIVE rows only; the rest of the buffer is padding.
    splits = torch.tensor([LIVE], device="cuda", dtype=torch.int32)
    probs = torch.ones(total_rows, device="cuda", dtype=torch.float32)

    with te.autocast(enabled=True, recipe=MXFP8BlockScaling()):
        out = model(x, splits, probs, splits)
    out.backward(dy)
    return fc1.get_parameter("weight0").main_grad.double()


reference = wgrad(LIVE, 0.0)                      # no padding
for label, fill in (("padding zeroed  ", 0.0), ("padding garbage ", 3.0)):
    got = wgrad(LIVE + PAD, fill)
    rel = ((got - reference).norm() / reference.norm()).item()
    bad = int((~torch.isfinite(got)).sum())
    print(f"{label} rel_err={rel:.4e}  non_finite={bad}/{got.numel()}")

Output on main (ace1873f):

padding zeroed   rel_err=nan  non_finite=411648/1048576
padding garbage  rel_err=nan  non_finite=411648/1048576

The two rows are identical — the corruption does not depend on what the padding holds.

Expected behavior

Rows past sum(split_sizes) are outside every group and must not affect any weight gradient, so both arms should match the unpadded reference:

padding zeroed   rel_err=0.0000e+00  non_finite=0/1048576
padding garbage  rel_err=0.0000e+00  non_finite=0/1048576

num_groups >= 2 already behaves this way at the same geometry, and so does num_groups == 1 with a GLU activation on cuDNN frontend >= 1.27.0 when the padded rows are zero.

Environment overview (please complete the following information)

• Environment location: Bare-metal (Slurm cluster node, RHEL 9.8)

• Method of Transformer Engine install: from source, into a venv, at commit ace1873f ("Fix linter error (Fix linter error #3435)"):

  python -m venv te220 && source te220/bin/activate
  pip install torch --index-url https://download.pytorch.org/whl/cu130
  git clone https://github.com/NVIDIA/TransformerEngine.git && cd TransformerEngine
  git checkout ace1873f && git submodule update --init --recursive
  NVTE_FRAMEWORK=pytorch NVTE_CUDA_ARCHS=103a NVTE_WITH_NCCL_EP=0 pip install --no-build-isolation .

• Docker: not used.

Environment details

• OS version: Red Hat Enterprise Linux 9.8 (Plow), kernel 5.14.0-687.38.1.el9_8.x86_64
• PyTorch version: 2.11.0+cu130
• Python version: 3.12.2
• Transformer Engine version: 2.20.0.dev0 (commit ace1873f)
• CUDA version: 13.0 (torch build); nvcc 13.3, driver 610.57.04
• CUDNN version: 9.19.0.56 (nvidia-cudnn-cu13), frontend nvidia-cudnn-frontend 1.27.0; nvidia-cutlass-dsl 4.7.1

Device details

• GPU model: NVIDIA B300 SXM6, compute capability (10, 3)

Additional context

Which configurations are affected. The trigger is _cudnn_frontend_supports_single_group_runtime_offsets, which returns not issubclass(activation_type, ScaledSReLU) and cudnn-frontend >= 1.27.0. So:

activation	cudnn-frontend	affected
ScaledSReLU	any version	yes
any activation	< 1.27.0	yes
ScaledSwiGLU, ScaledClampedQGeGLU, ScaledSiTUGLU	>= 1.27.0	see below

A second, related exposure on the GLU path. Where the specialization is available, producer and consumer agree and the scale misindex does not occur — but the plain wgrad GEMM in _single_group_wgrad_gemm still contracts over logical_shape[0] rows, so padding past sum(split_sizes) is summed into the weight gradient. That part is value-dependent: a zeroed tail contributes an all-zero outer product and is bit-exact (verified over 10 launches), while a garbage tail gives rel_err around 0.67 with realistic values. Since test_grouped_linear_cuda_graph_safe documents an uninitialized tail as permitted, this is arguably the same bug wearing a different mask, and a complete fix should address both.

Scope. Measured on MXFP8 only. Both gated code paths also accept NVFP4Quantizer, and the mechanism by which only wgrad breaks — wgrad contracts over the token axis, whereas forward and dgrad have M as a free dimension and merely produce extra output rows the caller discards — is quantizer-independent, so NVFP4 is plausibly affected as well. Not measured.

Why existing tests miss it. test_grouped_mlp builds in_shape = (split_sizes.sum(), hidden), so no test pads the input at all; the two tests that do pad (test_grouped_linear_cuda_graph_safe, test_grouped_mlp_cuda_graph_safe_mxfp8) both hard-code group_size = 4 and do not parametrize it. Padding and num_groups == 1 never co-occur.

When it was introduced. 3e7ae6ce ("Single group mxfp8 grouped mlp", #3267) widened both the dense-quantize guard and the single-group wgrad gate from NVFP4Quantizer to (MXFP8Quantizer, NVFP4Quantizer). It is the only commit between the 2.19 and 2.20 version bumps that touches those guards. 91bb9cfe (#3117) is not the cause: it moved the code but left the guard NVFP4-only. TE 2.18.0 measures 0.0000e+00 in every arm of the reproducer above.

Symptoms in a real run. Forward output and dX are bitwise clean, so the loss curve shows nothing while the expert weight gradients are wrong. In an MoE setting this is reachable whenever a rank has one local expert — N experts at EP=N.

Reproducibility. Deterministic. The reproducer gives identical non-finite counts on repeated fresh processes, and reproduces unchanged on stock nvidia-cudnn-frontend 1.27.0 as published.

I have a fix and can open a PR: gate the whole shortcut family on the same predicate the kernels use, so producer and consumer always agree on the row extent. It is host-side (an activation type and a package version), so it costs nothing and stays CUDA-graph capturable, and it leaves the GLU activations byte-for-byte on their existing fast path — which matters, because removing the shortcuts outright measurably slows a dense shared expert (GroupedLinear(num_groups=1), as Megatron's FusedSharedExpertMLP builds).

[GarlGuo]

close as #3485 is merged