GroupedLinear(single_grouped_weight=True) produces no weight gradient (wgrad)

[wujingyue]

Description

With single_grouped_weight=True (fused single-parameter GroupedLinear, enabled via NVTE_GROUPED_LINEAR_SINGLE_PARAM=1), the backward pass runs and produces the input gradient (dgrad), but the weight gradient (wgrad) is never emitted anywhere:

• fuse_wgrad_accumulation=False → weight.grad stays None
• fuse_wgrad_accumulation=True → a preallocated weight.main_grad is left all-zero (never written)

This happens in both bf16 and fp32. The per-expert path (single_grouped_weight=False) emits wgrad normally, which isolates the failure to the fused single-group weight. Because the fused weight receives no gradient, it cannot be trained (the optimizer sees no update) nor sharded/reduced by a data-parallel wrapper.

Environment

• TransformerEngine: 2.17.0+2e559f06
• PyTorch: 2.13.0a0+8145d630e8.nv26.06
• Container: nvcr.io/nvidia/pytorch:26.06-py3
• GPU: NVIDIA H100

Minimal reproduction

import os
os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = "1"
import torch
import transformer_engine.pytorch as te

layer = te.GroupedLinear(
    num_gemms=2, in_features=16, out_features=32, bias=False,
    single_grouped_weight=True, params_dtype=torch.bfloat16, device="cuda",
)
x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True)
layer(x, [4, 4]).sum().backward()

print("input dgrad present:", x.grad is not None)   # True  -> backward ran
print("weight wgrad:", layer.weight.grad)            # None  -> BUG (expected a gradient)

Expected behavior

layer.weight.grad is populated after backward() (as with nn.Linear and with single_grouped_weight=False); or, under fuse_wgrad_accumulation=True, the wgrad is accumulated into weight.main_grad.

Actual behavior

The backward runs (input dgrad is always produced), but the fused weight receives no gradient. Full matrix:

single_grouped_weight	fuse_wgrad_accumulation	dtype	dgrad (x.grad)	wgrad
False	False	bf16	set	✅ weight0/1.grad set
False	True	bf16	set	✅ present
True	False	bf16	set	❌ no gradient anywhere
True	True	bf16	set	❌ main_grad all-zero
False	False	fp32	set	✅ set
False	True	fp32	set	✅ present
True	False	fp32	set	❌ no gradient anywhere
True	True	fp32	set	❌ main_grad all-zero

Self-contained script covering all eight cases

import os
os.environ.setdefault("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")

import torch
import transformer_engine
import transformer_engine.pytorch as te

NUM_GEMMS, IN_FEATURES, OUT_FEATURES, TOKENS_PER_EXPERT = 2, 16, 32, 4


def weight_grad_state(weight):
    if weight.grad is not None:
        return f".grad set {tuple(weight.grad.shape)}"
    main_grad = getattr(weight, "main_grad", None)
    if main_grad is not None:
        return f".main_grad {'nonzero' if bool(main_grad.abs().sum() > 0) else 'ALL-ZERO (never written)'}"
    return "NO GRAD ANYWHERE"


def run(single_grouped_weight, fuse_wgrad_accumulation, dtype):
    torch.manual_seed(0)
    layer = te.GroupedLinear(
        num_gemms=NUM_GEMMS, in_features=IN_FEATURES, out_features=OUT_FEATURES, bias=False,
        single_grouped_weight=single_grouped_weight, params_dtype=dtype, device="cuda",
        fuse_wgrad_accumulation=fuse_wgrad_accumulation,
    )
    weights = [layer.weight] if single_grouped_weight else [layer.weight0, layer.weight1]
    if fuse_wgrad_accumulation:
        for weight in weights:
            weight.main_grad = torch.zeros(weight.shape, device="cuda", dtype=torch.float32)
            weight.grad_added_to_main_grad = False
    x = torch.randn(NUM_GEMMS * TOKENS_PER_EXPERT, IN_FEATURES, device="cuda", dtype=dtype, requires_grad=True)
    layer(x, [TOKENS_PER_EXPERT] * NUM_GEMMS).sum().backward()
    states = ", ".join(weight_grad_state(w) for w in weights)
    print(f"single_grouped_weight={single_grouped_weight!s:5}  fuse_wgrad={fuse_wgrad_accumulation!s:5}  "
          f"dtype={str(dtype).split('.')[-1]:8} | dgrad={'set' if x.grad is not None else 'None'} | wgrad -> [{states}]")


print(f"TransformerEngine {transformer_engine.__version__}, torch {torch.__version__}\n")
for dtype in (torch.bfloat16, torch.float32):
    for sgw in (False, True):
        for fuse in (False, True):
            run(sgw, fuse, dtype)
    print()

[wujingyue]

For contrast, using the op-fuser API (transformer_engine.pytorch.ops), the single-group weight gradient is produced. Repro (TE 2.17.0, H100):

import os
os.environ.setdefault("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")

import torch
import transformer_engine.pytorch.ops as te_ops

NUM_GEMMS, HIDDEN, FFN, TPE = 2, 16, 32, 4


def wstate(w):
    g = ".grad=None" if w.grad is None else f".grad={'nonzero' if w.grad.abs().sum()>0 else 'zero'}"
    mg = getattr(w, "main_grad", None)
    m = ".main_grad=absent" if mg is None else f".main_grad={'nonzero' if mg.abs().sum()>0 else 'ALL-ZERO'}"
    return f"{g}, {m}"


def gl(inf, outf, accum, dtype):
    op = te_ops.GroupedLinear(NUM_GEMMS, inf, outf, bias=False, device="cuda", dtype=dtype,
                              accumulate_into_main_grad=accum, single_grouped_weight=True)
    if accum:
        op.weight.main_grad = torch.zeros(op.weight.shape, device="cuda", dtype=torch.float32)
        op.weight.grad_added_to_main_grad = False
    return op


def run(accum, dtype):
    torch.manual_seed(0)
    fc1, fc2 = gl(HIDDEN, 2 * FFN, accum, dtype), gl(FFN, HIDDEN, accum, dtype)
    ops = te_ops.Sequential(fc1, te_ops.ScaledSwiGLU(), fc2)
    n = NUM_GEMMS * TPE
    x = torch.randn(n, HIDDEN, device="cuda", dtype=dtype, requires_grad=True)
    tpe = torch.tensor([TPE] * NUM_GEMMS, dtype=torch.int, device="cuda")
    ops(x, tpe, torch.ones(n, device="cuda", dtype=dtype), tpe).sum().backward()
    print(f"accumulate_into_main_grad={accum!s:5} dtype={str(dtype).split('.')[-1]:8} | "
          f"fc1[{wstate(fc1.weight)}] fc2[{wstate(fc2.weight)}]")


for dt in (torch.bfloat16, torch.float32):
    for acc in (False, True):
        run(acc, dt)

Output:

accumulate_into_main_grad=False dtype=bfloat16 | fc1[.grad=nonzero, .main_grad=absent] fc2[.grad=nonzero, .main_grad=absent]
accumulate_into_main_grad=True  dtype=bfloat16 | fc1[.grad=zero, .main_grad=nonzero] fc2[.grad=zero, .main_grad=nonzero]
accumulate_into_main_grad=False dtype=float32  | fc1[.grad=nonzero, .main_grad=absent] fc2[.grad=nonzero, .main_grad=absent]
accumulate_into_main_grad=True  dtype=float32  | fc1[.grad=nonzero, .main_grad=nonzero] fc2[.grad=nonzero, .main_grad=nonzero]

[zhongbozhu]

Yes it's a known issue.

For the use_transformer_engine_op_fuser path, you will need this PR to Mcore to land: NVIDIA/Megatron-LM#5487

For regular grouped linear, you will need the following two changes:

• TE (WIP): Improve device-init grouped linear module with single grouped weight support  #3224
• Mcore (WIP): https://github.com/zhongbozhu/Megatron-LM/tree/main_support_device_group_linear

With these two changes, user need to switch to grouped_tensor API to use single weight feature as well.

    # I need this convince Mcore people to merge this :) 

    moe_grouped_gemm_backend: Literal["legacy", "grouped_tensor"] = "legacy"
    """Grouped GEMM execution backend. ``legacy`` uses per-expert weights and host split
    metadata. ``grouped_tensor`` uses Transformer Engine GroupedTensor kernels with CUDA split
    metadata. The TE operation fuser selects ``grouped_tensor`` for grouped MoE MLPs.
    """

[wujingyue]

Thanks for working on these and sharing the update, @zhongbozhu!

Could you self-assign this issue and update it once it's done? That way I'll get notified automatically instead of having to remember to follow up—I'm not great at chasing these things.

For now, I'll treat single-grouped weights as premature and keep using separate expert weights in MFSDP.