Fix expert parallelism through Trainer by qgallouedec · Pull Request #48208 · huggingface/transformers

qgallouedec · 2026-08-22T02:44:04Z

Stacked on #48200, this PR makes expert parallelism reachable through Trainer, and without the gradient fixes there, the training it unlocks is wrong.

Loading a model with DistributedConfig(tp_size=N, enable_expert_parallel=True) and handing it to Trainer fails on main before the first step completes, in three places:

maybe_distribute_model never assigns model._tp_size (the attribute is declared but dead), so the Trainer builds no ParallelismConfig and accelerate wraps the DTensor model in DDP: ValueError: Your model contains DTensor parameters, which is incompatible with DDP.
With that fixed, gradient clipping crashes: _foreach_norm cannot span a parameter set that mixes DTensors (the experts) and plain tensors (everything else).
Same for the optimizer: the fused/foreach AdamW kernels cannot span the mixed set.

The fix: assign _tp_size where the mesh is recorded; compute the gradient norm (and clip) per-gradient with replication-aware discounting when parameters live on different meshes; fall back to per-parameter AdamW stepping for mixed parameter sets.

Validation

Together with #48205 (which fixes the gradients EP computes), expert-parallel full fine-tuning through Trainer runs end-to-end and tracks a single-GPU control step by step (OLMoE-1B-7B, tp_size=4, 4×H100, identical data):

EP via Trainer:   single GPU control:
12.13             12.13
12.44             12.45
11.52             11.51
11.22             11.26
11.12             11.15
11.04             11.09

Reproduction (each of the three failures appears on main as the previous one is fixed; on this branch it trains)

# torchrun --nproc_per_node 4 repro.py
import torch
from torch.utils.data import Dataset

from transformers import AutoModelForCausalLM, Trainer, TrainingArguments
from transformers.distributed import DistributedConfig


class RandomTokens(Dataset):
    def __len__(self):
        return 512

    def __getitem__(self, i):
        g = torch.Generator().manual_seed(i)
        ids = torch.randint(0, 50000, (512,), generator=g)
        return {"input_ids": ids, "labels": ids.clone()}


model = AutoModelForCausalLM.from_pretrained(
    "allenai/OLMoE-1B-7B-0924",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=4, enable_expert_parallel=True),
)

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="/tmp/ep-trainer",
        per_device_train_batch_size=1,
        max_steps=8,
        logging_steps=1,
        report_to=[],
        save_strategy="no",
    ),
    train_dataset=RandomTokens(),
)
trainer.train()

HuggingFaceDocBuilderDev · 2026-08-22T02:59:49Z

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.

ArthurZucker

LGTM for TP, for trainer, @SunMarc can you approve?

ArthurZucker · 2026-09-01T06:39:08Z

+    def _has_mixed_mesh_grads(self, model) -> bool:
+        from torch.distributed.tensor import DTensor
+
+        grads = [p.grad for p in model.parameters() if p.grad is not None]


detecting ep plan or local flags form the tp plan might be better / faster / cashable to waste less?

+1 on this. Or we could try to cache de result and store the result somewhere, either on the model or on the trainer.

Cached it on the trainer in d0e4af5 (resolved at the first clip, sharding is static for the run).
Went with the cache rather than reading the tp plan so the Trainer stays out of plan semantics; the optimizer-side check runs once at construction so it needed nothing.

SunMarc

Thanks ! Left a few nits

SunMarc · 2026-09-01T08:34:22Z

    return AdamW, ctx.optimizer_kwargs


+def _has_mixed_dtensor_params(model) -> bool:


This is basically the same function as the mixed_mesh one no ?

Correct, deduped in d0e4af5

SunMarc · 2026-09-01T08:35:54Z

+    def _has_mixed_mesh_grads(self, model) -> bool:
+        from torch.distributed.tensor import DTensor
+
+        grads = [p.grad for p in model.parameters() if p.grad is not None]


+1 on this. Or we could try to cache de result and store the result somewhere, either on the model or on the trainer.

SunMarc

Thanksss !

SunMarc · 2026-09-03T15:02:53Z

@bot /style

github-actions · 2026-09-03T15:03:34Z

Style fix bot fixed some files and pushed the changes.

Loading a model with DistributedConfig(tp_size=N, enable_expert_parallel=True) and handing it to Trainer fails on main in three places: 1. maybe_distribute_model never assigns model._tp_size, so the Trainer builds no ParallelismConfig and accelerate wraps the DTensor model in DDP: ValueError: Your model contains DTensor parameters, which is incompatible with DDP. 2. _get_grad_norm calls clip_grad_norm_ over the full parameter set, and _foreach_norm cannot span a mix of DTensor (experts) and plain parameters. 3. The fused/foreach AdamW kernels cannot span that mix either. Set _tp_size where the mesh is recorded, compute the gradient norm (and clip) per-gradient with replication-aware discounting when parameters live on different meshes, and fall back to per-parameter AdamW stepping for mixed parameter sets. With this and #48205, expert-parallel full fine-tuning through Trainer runs end-to-end and tracks a single-GPU control step by step (OLMoE-1B-7B, tp=4: 12.13 -> 12.45 -> 11.52 -> ... -> 11.04 vs 12.13 -> 12.45 -> 11.51 -> ... -> 11.09).

…he trainer

github-actions · 2026-09-03T18:22:33Z

CI recap

Dashboard: View test results in Grafana
Latest run: 33787745619:1
Result: success | Jobs: 16 | Tests: 89,822 | Failures: 0 | Duration: 7h 23m

…ributions (huggingface#48205) * Fix NaN gradients in expert-parallel training: mask uninitialized grouped_mm rows Under EP, sentinel token-expert slots sit beyond offsets[-1] and torch._grouped_mm leaves those output rows (fwd output and bwd d_input) uninitialized. The forward relied on a single post-mask plus a single pre-mask, letting NaN/Inf from uninitialized memory transit the activation and down-projection backward. The gate product's backward (act_fn(gate) * up) turns 0 x Inf into NaN (torch.autograd anomaly mode names this exact Mul), and it escapes into finite gradients: full fine-tuning of any EP-sharded MoE produced nan grad_norm on the second step (the first step survives only because freshly-allocated CUDA memory happens to be zeroed) and the loss collapsed to 0. Mask the sentinel-tail rows after each grouped GEMM instead. Full fine-tuning of OLMoE-1B-7B under ep=4 now matches the single-GPU loss trajectory. * Fix wrong gradients for all non-expert parameters in expert-parallel training Under EP the router hook zeroes the routing scores of non-local experts, so in backward each rank's score gradient covers only the slots of its local experts, and nothing sums the per-rank partial gradients: the gate weights and, through the gate's input, every parameter upstream of each MoE block receive gradients missing the contributions that flow through remote experts. The existing _AllReduceBackward on the experts' hidden input covers the dispatch branch, and the top_k_weights branch is explicitly skipped when is_expert_parallel -- but under EP it is exactly as partial as under TP-MoE. Measured against a single-GPU reference (OLMoE-1B-7B, one batch, fp32 so rounding noise vanishes): before the fix, 3/179 parameters agree (relative max-abs errors 0.3-2.5 on attention, norms, embeddings and router gates, 10-100x above the run-to-run noise floor; only the last layer's experts and the final norm -- the parameters backward reaches before crossing an expert block -- are correct). After the fix: 179/179 agree, max relative error 2.7e-5. Fix: allreduce-sum the score gradient in the EP router hook, before the non-local mask (each slot has exactly one owning rank, so the sum is exact). * Trim comments * Gate the router-score backward all-reduce on grad mode; drop the post-mask superseded by the per-mm masks * Fix expert parallelism through Trainer (huggingface#48208)

qgallouedec changed the base branch from main to fix-ep-training-gradients August 22, 2026 02:44

qgallouedec mentioned this pull request Aug 23, 2026

[Umbrella] Enable FSDP2 + expert parallelism via a 2-D (fsdp, tp) device mesh #48204

Draft

qgallouedec force-pushed the fix-ep-training-gradients branch from e79ef7a to ab27829 Compare August 25, 2026 16:55

qgallouedec force-pushed the fix-ep-through-trainer branch from f3e0be8 to 894690a Compare August 25, 2026 16:55

qgallouedec requested a review from 3outeille August 27, 2026 22:45

qgallouedec force-pushed the fix-ep-training-gradients branch from ab27829 to 1b718f6 Compare August 27, 2026 22:45

qgallouedec force-pushed the fix-ep-through-trainer branch from 894690a to a2bfd50 Compare August 27, 2026 22:45

ArthurZucker approved these changes Sep 1, 2026

View reviewed changes

SunMarc approved these changes Sep 1, 2026

View reviewed changes

qgallouedec force-pushed the fix-ep-training-gradients branch from 1b718f6 to 0f3bd58 Compare September 3, 2026 01:25

qgallouedec force-pushed the fix-ep-through-trainer branch from a2bfd50 to ba4a3c3 Compare September 3, 2026 01:25

qgallouedec force-pushed the fix-ep-training-gradients branch from 7398af0 to 1a958cd Compare September 3, 2026 02:40

qgallouedec force-pushed the fix-ep-through-trainer branch from ba4a3c3 to 3019094 Compare September 3, 2026 02:40

SunMarc approved these changes Sep 3, 2026

View reviewed changes

qgallouedec force-pushed the fix-ep-training-gradients branch from 1a958cd to b52614e Compare September 3, 2026 17:58

qgallouedec force-pushed the fix-ep-through-trainer branch from 352c82f to 967df85 Compare September 3, 2026 17:58

qgallouedec and others added 4 commits September 3, 2026 14:07

Remove comment

a2b284e

Share the mixed-DTensor predicate and cache the clip-time result on t…

a19ebb8

…he trainer

Apply repo consistency fixes

f464638

qgallouedec force-pushed the fix-ep-training-gradients branch from b52614e to e0b2e71 Compare September 3, 2026 18:07

qgallouedec force-pushed the fix-ep-through-trainer branch from 967df85 to f464638 Compare September 3, 2026 18:07

qgallouedec merged commit caf370a into fix-ep-training-gradients Sep 3, 2026
113 checks passed

qgallouedec deleted the fix-ep-through-trainer branch September 3, 2026 18:22

This was referenced Sep 3, 2026

Enable FSDP2 + expert parallelism via a 2-D (fsdp, tp) device mesh #48516

Open

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

Merged

qgallouedec added a commit that referenced this pull request Sep 5, 2026

Fix expert parallelism through Trainer (#48208)

f290f32

qgallouedec added a commit that referenced this pull request Sep 7, 2026

Fix expert parallelism through Trainer (#48208)

835d4e6

qgallouedec added a commit that referenced this pull request Sep 7, 2026

Fix expert parallelism through Trainer (#48208)

0e5d1eb

SunMarc mentioned this pull request Sep 8, 2026

Clip grad norm support for dtensors huggingface/accelerate#4219

Merged

qgallouedec added a commit that referenced this pull request Sep 8, 2026

Fix expert parallelism through Trainer (#48208)

96769d9

		return AdamW, ctx.optimizer_kwargs


		def _has_mixed_dtensor_params(model) -> bool:

Uh oh!

Conversation

qgallouedec commented Aug 22, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Validation

Uh oh!

HuggingFaceDocBuilderDev commented Aug 22, 2026

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ArthurZucker Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

SunMarc Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

qgallouedec Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

SunMarc left a comment

Choose a reason for hiding this comment

Uh oh!

SunMarc Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

qgallouedec Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

SunMarc Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

SunMarc left a comment

Choose a reason for hiding this comment

Uh oh!

SunMarc commented Sep 3, 2026

Uh oh!

github-actions Bot commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Sep 3, 2026

CI recap

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

qgallouedec commented Aug 22, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Sep 3, 2026 •

edited

Loading