Fix expert-parallel training: NaN gradients and missing gradient contributions

[qgallouedec]

Two bugs in expert-parallel training (enable_expert_parallel=True), found while validating FSDP2 × EP (#48204): full fine-tuning of an EP-sharded MoE trains on wrong gradients, then collapses to NaN on step 2. +14 lines total.

1. NaN from uninitialized memory. torch._grouped_mm leaves the sentinel-tail rows of its output and backward d_input uninitialized, and the NaN escapes through the act_fn(gate) * up backward (0 × Inf). Step 1 survives only because fresh CUDA memory is zeroed; step 2 reuses dirty memory → grad_norm=nan, loss → 0. Fix: mask the sentinel rows after each grouped GEMM.

2. Missing gradient contributions for every non-expert parameter. The router hook zeroes non-local score slots, so each rank's score gradient covers only its local experts, and the per-rank partials are never summed: everything upstream of each MoE block loses the gradient flowing through remote experts. Fix: allreduce-sum the score gradient before the mask (each slot has one owning rank, so the sum is exact).

Certification

fp32 gradient diff vs a single-GPU reference (OLMoE-1B-7B, one batch, tp_size=4) — median relative error per class:

class	main	this PR
attention (64)	3.0e-1	7.8e-7
norms (65)	3.0e-1	7.7e-7
router gates (16)	1.0e+0	1.4e-6
embed / lm_head (2)	2.3e-1	1.5e-6
experts (32)	3.6e-1	1.3e-6

With this PR all 179 parameters match the reference (max rel err 2.5e-5). Same certification on real Qwen3-30B-A3B weights (first 4 layers): 47/47, max 2.7e-6; and on gpt-oss-20b (expert biases + sigmoid gate): 71/71, max 9.5e-3. In bf16, EP full fine-tuning goes 12.13 → nan/0 on main; with this PR it tracks a single-GPU control step by step (12.13 → 12.45 → 11.52 → … → 11.04 vs 12.13 → 12.45 → 11.51 → … → 11.09). Also fixes NaN under fused losses consuming EP outputs (e.g. TRL's chunked cross-entropy).

Certification script (raw forward/backward, EP through Trainer is separately broken on main, see #48204)

# 1) python cert.py single ref.pt        (1 GPU)
# 2) torchrun --nproc_per_node 4 cert.py ep ref.pt
import sys

import torch
from torch.distributed.tensor import DTensor

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

mode, ref_path = sys.argv[1], sys.argv[2]

kwargs = {}
if mode == "ep":
    kwargs["distributed_config"] = DistributedConfig(tp_size=4, fsdp_size=1, enable_expert_parallel=True)
model = AutoModelForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924", dtype=torch.float32, **kwargs)
if mode == "single":
    model.cuda()
model.train()

g = torch.Generator().manual_seed(0)
ids = torch.randint(0, 50000, (1, 512), generator=g).cuda()
out = model(input_ids=ids, labels=ids.clone())
out.loss.backward()

rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
if mode == "single":
    grads = {n: p.grad.detach().float().cpu() for n, p in model.named_parameters() if p.grad is not None}
    torch.save(grads, ref_path)
    print(f"saved {len(grads)} grads")
else:
    ref = torch.load(ref_path, weights_only=False) if rank == 0 else None
    bad = 0
    for name, p in model.named_parameters():
        if p.grad is None:
            continue
        gr = p.grad
        if isinstance(gr, DTensor):
            gr = gr.full_tensor()
        if rank == 0:
            ga = ref[name].cuda()
            rel = ((ga - gr.detach().float()).abs().max() / (ga.abs().max() + 1e-12)).item()
            bad += rel > 2e-2
    if rank == 0:
        print(f"{bad} parameters with relative error > 2e-2")
    torch.distributed.barrier()
    torch.distributed.destroy_process_group()

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[IlyasMoutawwakil]

do we need a pre and post masking 🥲 ? i feel like it's redundant, the sample_weights_g are already zero in sentinel rows

[qgallouedec]

The post-mask on weighted_out was indeed redundant once this PR zeroes `proj_out per mm, dropped it in 7398af0.

The pre-mask is important though: the danger isn't the values (where weights=0 would save us) but uninitialized kernel memory, and the grouped-mm backward writes garbage d_input rows for the skipped sentinel rows regardless of upstream grads.

With the allocator pre-poisoned with NaN buffers, masks on gives finite grads; ablating just the pre-mask gives NaN grads (tested locally).

[ArthurZucker]

Yep, this is great. @3outeille your tests need to catch this, can you make sure they now do?

[qgallouedec]

Thanks for the reviews guys.

@ArthurZucker I can also turn the harness into a slow 8-GPU test (EP fwd+bwd grads vs single-GPU reference on a tiny MoE) and add it to this PR, say the word and I'll write it in the repo's test style.

In any case, here there are:

from transformers.testing_utils import TestCasePlus, is_tensor_parallel_test
EP_TRAINING_WORKER = """
import torch
from transformers import AutoModelForCausalLM, Qwen3MoeConfig, Qwen3MoeForCausalLM
from transformers.distributed import DistributedConfig

config = Qwen3MoeConfig(
    vocab_size=128, hidden_size=64, intermediate_size=128, moe_intermediate_size=32,
    num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=4, num_experts=16,
    num_experts_per_tok=4, decoder_sparse_step=1, head_dim=16,
)
torch.manual_seed(0)
reference = Qwen3MoeForCausalLM(config).float()

import tempfile, os
tmp = os.environ["EP_TEST_CKPT"]
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
if rank == 0:
    reference.save_pretrained(tmp)
torch.distributed.init_process_group(backend="nccl")
torch.cuda.set_device(rank)
torch.distributed.barrier()

torch.manual_seed(1)
batch = torch.randint(0, 128, (4, 32))

# Single-model reference gradients (every rank computes its own copy, fp32)
reference = reference.cuda()
ref_out = reference(input_ids=batch.cuda(), labels=batch.cuda())
ref_out.loss.backward()

# Expert-parallel model over the whole world
model = AutoModelForCausalLM.from_pretrained(
    tmp, dtype=torch.float32,
    distributed_config=DistributedConfig(tp_size=world_size, fsdp_size=1, enable_expert_parallel=True),
)
out = model(input_ids=batch.to(model.device), labels=batch.to(model.device))
out.loss.backward()

torch.testing.assert_close(out.loss, ref_out.loss, rtol=1e-5, atol=1e-5)
ref_params = dict(reference.named_parameters())
checked = 0
for name, param in model.named_parameters():
    if param.grad is None:
        continue
    assert param.grad.isfinite().all(), f"non-finite grad in {name}"
    grad = param.grad
    ref_grad = ref_params[name].grad
    if hasattr(grad, "_local_tensor"):  # EP-sharded experts: compare this rank's slice
        local = grad._local_tensor
        n_local = local.shape[0]
        ref_grad = ref_grad[rank * n_local : (rank + 1) * n_local]
        grad = local
    torch.testing.assert_close(grad, ref_grad, rtol=1e-4, atol=1e-5, msg=name)
    checked += 1
assert checked > 0
if rank == 0:
    print(f"EP_GRADIENTS_MATCH ({checked} params)")
"""


@is_tensor_parallel_test
class TestExpertParallelTrainingGradients(TestCasePlus):
    def test_ep_training_gradients_match_single_model(self):
        """Expert-parallel fwd+bwd must reproduce the single-model loss and gradients (router
        included): guards the sentinel-row masking and the router-score gradient all-reduce."""
        import tempfile

        from transformers.testing_utils import backend_device_count, torch_device, torchrun

        nproc = min(backend_device_count(torch_device), 8)
        if nproc < 2:
            self.skipTest("needs at least 2 accelerators")
        with tempfile.TemporaryDirectory() as tmp:
            import os

            env = os.environ.copy()
            env["EP_TEST_CKPT"] = tmp
            torchrun(EP_TRAINING_WORKER, nproc, env=env)

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 33788592511:1
Result: success | Jobs: 16 | Tests: 186,276 | Failures: 0 | Duration: 14h 47m