Expert-parallel token dispatch: each rank trains on its own part of the batch by qgallouedec · Pull Request #48518 · huggingface/transformers

qgallouedec · 2026-09-04T03:11:38Z

Stacked on #48516 (2-D mesh), itself on #48205. Review this branch against fsdp2-ep-2d-mesh.

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-30B-A3B",
    distributed_config=DistributedConfig(tp_size=8, enable_expert_parallel=True, expert_parallel_dispatch=True),
)

Why

EP today runs the whole batch on every rank: each rank keeps its experts, masks the others, and the group all-reduces the expert outputs after every MoE layer. Everything outside the experts (attention, embeddings, norms) does tp_size times the same work. At 64 GPUs on a 110B model, that is 32 ranks computing the same 2 sequences!

With expert_parallel_dispatch=True, each rank trains on its own part of the batch (the torchtitan / DeepSeek layout). At every MoE layer it routes its tokens, sends each (token, expert) pair to the rank that owns the expert with an all-to-all, runs its local experts, and gets the results back with a second all-to-all. Only the routed tokens travel.

Results

Unique tokens per second per node (under masked EP the tp ranks replay the same batch). Qwen3-30B-A3B full fine-tuning, bf16, 8xH100, seq 2048:

configuration	batch	masked	dispatch	speedup
`tp_size=8`	1	3.4k, 38.6 GB	16.6k, 29.8 GB	4.8x
`tp_size=8`	4	5.8k, 40.1 GB	20.6k, 33.2 GB	3.5x
`tp_size=4, fsdp_size=2`	1	5.8k, 34.2 GB	16.6k, 30.0 GB	2.9x
`tp_size=4, fsdp_size=2`	4	not run	24.9k, 33.5 GB

At scale, GLM-4.5-Air (110B) full fine-tuning on 64xH100 with tp_size=32, fsdp_size=2, same TRL script for both arms:

arm	step time	unique sequences / step	sequences / s
masked	5.9 s	2	0.34
dispatch, batch 1	5.9 s	64	10.8 (32x)
dispatch, batch 4	9.4 s	256	27.3 (80x)

Both arms train (loss 3.9 -> 1.2 masked, 3.4 -> 1.3 dispatch over 20 steps).
Per-rank step time is longer under dispatch (two all-to-alls and a local experts pass per layer instead of one all-reduce), which is why the small-scale speedup is below tp_size; with #48201's gating of the sentinel masking the 8-GPU numbers become 5.3x / 4.1x / 3.1x.

Correctness

Trainer parity, fp32, 8 GPUs, 6 steps with clipping: loss and grad norm within 4.8e-7 / 1.2e-7 of a single-process run on the same total batch, for tp_size=8, 4x2 and 2x4. Saved weights within 3.3e-6 of the single-process save.
Gradient certification on real Qwen3-30B-A3B weights (2 and 4 layers): every parameter within 7.6e-6 of a single GPU processing the same samples, for tp_size=4 and 2x2. The masked path certifies at 2.0e-6 on the same model, so this is the floor.
The tests catch a wrong dispatch: a commit where the style was silently not installed fails test_fsdp2_expert_parallel_2d_vs_ddp[dispatch] at step 0.

Limitations

Two host syncs per MoE layer (the all-to-all split sizes). A capacity-padded, sync-free variant did not pay for its padding traffic.
Map-style training datasets only; dispatch_batches=True and non-random sampling strategies are rejected.
With fsdp_size=1 the experts stay outside FSDP2, so fsdp_mixed_precision and fsdp_cpu_offload do not apply to them.
Evaluation is unchanged: the tp ranks evaluate the same batches.
Same as Enable FSDP2 + expert parallelism via a 2-D (fsdp, tp) device mesh #48516: no resume for models sharded at load time.

HuggingFaceDocBuilderDev · 2026-09-04T03:23:38Z

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.

AmineDiro · 2026-09-04T09:52:26Z

Very cool ! and first all-to-all implementation of EP. I wonder if this works with different expert_forward like sonicMoe that was added by @IlyasMoutawwakil

IlyasMoutawwakil · 2026-09-04T10:11:17Z

awesome work ! let's make sure it works with all experts impls (normally it should ootb since sentinels are the exception) and also think of how it can be extended to more advanced dispatch impls like DeepEP, i'm thinking something like a literal: experts_dispatch=all-reduce/all-to-all/DeepEP(through kernels for example) cc @3outeille

AmineDiro · 2026-09-07T12:47:41Z

+            raise OSError("Expert-parallel token dispatch requires `torch>=2.7`.")
+        # The DTensor parameters are the expert-parallel experts (the expert parallel plan shards only them).
+        expert_modules = [
+            module for module in model.modules() if any(is_dtensor(p) for p in module.parameters(recurse=False))


Is this always sufficient ? i.e. all the Dtensors are by definition expert tensors ?

after a3e218c it's guaranteed. maybe_distribute_model rewrites _ep_plan to keep only the expert styles before apply_tensor_parallelism runs, so the only DTensor params by the time we get here are the experts.

But it used to be enforced by a raise, which broke DeepseekV4: its EP plan also shards the lightning indexer, so tp test died at from_pretrained.
Those entries are now dropped with a warning and the modules stay replicated, which is what dispatch wants anyway since every rank holds a different batch. Both tests pass now.

qgallouedec · 2026-09-07T21:42:54Z

@IlyasMoutawwakil

let's make sure it works with all experts impls (normally it should ootb since sentinels are the exception)

should be fine for all four.
The dispatch calls the experts module's own forward as a top-1 routing with unit weights, so deepgemm, batched_mm, grouped_mm and sonicmoe all get a normal routing tensor.
About the sentinels: yes, with dispatch every rank receives the tokens its own experts own, so no sentinel rows are ever built and the impls that handle them just never see them.

Caveat: the EP mixin tests are CPU-only, so what I actually exercised is the default path. deepgemm and sonicmoe need a GPU test.

and also think of how it can be extended to more advanced dispatch impls like DeepEP, i'm thinking something like a literal: experts_dispatch=all-reduce/all-to-all/DeepEP(through kernels for example)

I like it (although I'm not really into it aha). Mega MoE is already the precedent, MoeTensorParalellMegaMoeExperts does its own dispatch, combine and token sharding inside the kernel, it just happens to be inference-only so it does not collide with this path today. DeepEP would be the same shape.

VI-Arthur

Will checkout and fix myself!

VI-Arthur · 2026-09-08T09:19:45Z

+        if expert_mesh is not None:
+            for module in expert_modules:
+                fully_shard(module, mesh=expert_mesh, reshard_after_forward=True, **fsdp_policy_kwargs)
+        else:
+            fsdp_policy_kwargs["ignored_params"] = {p for module in expert_modules for p in module.parameters()}


No, the way to say "we want only these FSDP" is here: https://github.com/VI-Arthur/transformers/blob/fix-ep-eager-waste/src/transformers/distributed/mixin.py#L202-L202

This needs a small update but:

fsdp_plan = dict(getattr(model, "_fsdp_plan", None) or {})

setting a custom plan will make sure onlyl these are used

Should of course be in the distributed config

Yes, the DTensor sniff is fragile: it happens to be exact only because the block above rewrites _ep_plan to expert-only styles, so it depends on something 40 lines away.
_fsdp_plan is already resolved at line 202 in this function, so it's the natural handle.

You said you'd take this one, so I'll leave it to you 👍

VI-Arthur · 2026-09-08T09:22:48Z

+                    replicated = sorted(
+                        name
+                        for name, style in model.tp_plan.items()
+                        if style not in ("ep_router", "grouped_gemm", "moe_tp_experts")
+                    )


pretty sure default is replicate, so if you pass the ep plan hard coded, it won't follow the one form the config. You don't need any of that in that case no?

it's not setting a default, it's dropping entries.
tp_plan reads _ep_plan under EP, and dispatch can only shard the experts, so the rewrite keeps grouped_gemm/moe_tp_experts and drops everything else the plan would otherwise shard.
Without it a plan that shards e.g. attention would still shard it, and each rank is training on its own tokens, so that would be wrong rather than just wasteful.
The warning lists whatever got dropped.

Agreed it's the wrong place to hardcode the style names though. If it moves behind _fsdp_plan in the distributed config as you suggested, this block should go with it. 👌

VI-Arthur · 2026-09-08T09:23:32Z

+                hidden_states = hidden_states.to_local()
+            with self.context_around_forward(module, mesh):
+                # The sharding leaves the module with its local expert count.
+                return dispatch_experts_forward(


if this func is only used here, declare it here

Can do, it's ~60 lines and standalone, no moe.py helpers.
Only hesitation is cohesion: it's the fourth experts-forward variant and the other three live in moe.py, so tensor_parallel.py would start importing MoE kernels.

…he batch DistributedConfig(expert_parallel_dispatch=True) sends every selected (token, expert) pair to the rank that owns the expert with an all-to-all, runs the local experts, and sends the results back, instead of running the whole batch on every expert-parallel rank and all-reducing the outputs. Each rank trains on its own part of the batch: the parameters outside the experts are sharded with FSDP2 across every rank (the flattened (fsdp, tp) mesh) so FSDP2 owns their gradient reduction, the experts stay sharded across tp and, if set, fsdp, and the all-to-all backward accumulates their gradients across the group (scaled by 1/ep_size to match the data-parallel average). The Trainer gives each rank its own batches and counts tokens across all of them.

…rded_data_parallelism

…on, as the masked path does

…oken dispatch accelerate only shards and seeds the sampler when it sees more than one data-parallel rank, so with tp_size alone every rank drew its own random full batch.

… forward for the local experts, ids from the count exchange Also reject expert parallel plans that shard anything but the experts, since the ranks no longer see the same batch.

…h in the Trainer ep_dispatch_router and ep_dispatch_experts are substituted into the expert parallel plan when expert_parallel_dispatch is set, so no style carries a flag it ignores and nothing is stored on the modules. get_tp_size() is 1 under dispatch, which the loss scale, the token count and the total batch size all go through; non-random sampling strategies are rejected explicitly.

…dispatch_batches

…patch styles

…rts implementation

…bine passes cost 15-25% of throughput

The dispatch kernel no longer assumes the grouped-GEMM parameter layout: the experts module's forward is called as a top-1 routing with unit weights, so any experts implementation (and any model-specific step inside it) works under dispatch.

… check the batches at init

…the experts implementation check

…default plan

`expert_parallel_dispatch=True` raised when the expert parallel plan sharded anything but the experts. DeepseekV4's plan does: it shards the lightning indexer colwise and all-reduces the scorer, so `test_ep_forward_2` and `test_ep_backward_1` died in `from_pretrained`. Drop those entries with a warning instead. Under dispatch every rank trains on its own part of the batch, so the trunk is replicated and FSDP2-sharded across the whole mesh anyway, which is what dropping them gives.

Co-authored-by: VI-Arthur <arthurzucker1@gmail.com>

stevhliu

thanks for the docs!

i think this section may contradict the "Combining with FSDP2" section below it a bit? this was how i interpreted it:

Token dispatch → dense weights are sharded across everyone
Combining with FSDP2 → dense weights are replicated on every rank

so it feels like the Combining section "undoes" the first one. it'd be easier to follow i think if we clarified this with a sentence in the Combining section below

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Dropping it left the parameter with no gradient reduction at all: FSDP2 treats the experts module as expert-owned and skips it, and every rank saw different tokens. fp32 certification on a tiny Muse-shaped model: post_expert_norm.weight at rel 0.80 and 0.50 against a 1-GPU reference under dispatch, 0.75 and 0.59 under dispatch2d; masked and 2-D at 2e-6. A ~0.75 error is the local quarter of the sum that never happened.

A missing or doubled gradient reduction leaves the forward and the loss untouched, so the loss-only check passed on both. Gather each sharded gradient back to the full parameter and compare every parameter's gradient against the single-process reference, and require the same set of parameters to have received one.

github-actions · 2026-09-09T01:21:35Z

CI recap

Dashboard: View test results in Grafana
Latest run: 34297557595:1
Result: failure | Jobs: 16 | Tests: 187,304 | Failures: 2 | Duration: 17h 59m

qgallouedec force-pushed the fsdp2-ep-2d-mesh branch from 393af53 to d226b02 Compare September 5, 2026 03:10

qgallouedec force-pushed the ep-token-dispatch branch from 4a9e39d to 1329a33 Compare September 5, 2026 03:10

AmineDiro reviewed Sep 7, 2026

View reviewed changes

qgallouedec force-pushed the fsdp2-ep-2d-mesh branch from d226b02 to 680e1ff Compare September 7, 2026 20:00

qgallouedec force-pushed the ep-token-dispatch branch from 1329a33 to 5341f4b Compare September 7, 2026 20:01

qgallouedec force-pushed the fsdp2-ep-2d-mesh branch from 680e1ff to d9394c9 Compare September 7, 2026 21:43

qgallouedec force-pushed the ep-token-dispatch branch from a3e218c to 18ed512 Compare September 7, 2026 21:43

qgallouedec mentioned this pull request Sep 7, 2026

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

Draft

VI-Arthur reviewed Sep 8, 2026

View reviewed changes

qgallouedec added 17 commits September 8, 2026 13:09

Make the dispatch handling of the experts explicit in apply_fully_sha…

7063c45

…rded_data_parallelism

Combine the dispatched expert outputs in the routing weights' precisi…

dc1fdee

…on, as the masked path does

Split the batches across every rank with a DistributedSampler under t…

f2a3c26

…oken dispatch accelerate only shards and seeds the sampler when it sees more than one data-parallel rank, so with tp_size alone every rank drew its own random full batch.

Simplify the dispatch: torch's differentiable all-to-all, the grouped…

3ab94ff

… forward for the local experts, ids from the count exchange Also reject expert parallel plans that shard anything but the experts, since the ranks no longer see the same batch.

Select the dispatch styles in the mixin, guard the torch version and …

b9bf58b

…dispatch_batches

Rewrite the expert parallel plan the model actually reads for the dis…

39cdce5

…patch styles

Seed the dispatch sampler with data_seed; require the grouped_mm expe…

9f86230

…rts implementation

Run the local experts directly: the grouped forward's masking and com…

930c146

…bine passes cost 15-25% of throughput

Dispatch: one host sync per layer, comments say what the code does

51f281d

Dispatch: drop the router style, reuse the named mesh and is_dtensor,…

101bfea

… check the batches at init

Dispatch: skip the mixin tests for models without an FSDP plan, drop …

2d9046b

…the experts implementation check

Fix the _has_fsdp_plan docstring

42606af

Drop the FSDP plan skip from the mixin tests: every config carries a …

d8eb1fc

…default plan

Format

23920e3

qgallouedec added 3 commits September 8, 2026 13:09

Keep the base install_forward signature on the dispatch experts style

e149869

Docs: shorter token dispatch overview

e139e8a

qgallouedec force-pushed the fsdp2-ep-2d-mesh branch from d9394c9 to 3a8da11 Compare September 8, 2026 17:09

qgallouedec force-pushed the ep-token-dispatch branch from 18ed512 to cf74929 Compare September 8, 2026 17:09

qgallouedec and others added 2 commits September 8, 2026 14:02

Merge branch 'fsdp2-ep-2d-mesh' into ep-token-dispatch

18144ab

Update docs/source/en/expert_parallelism.md

c9a4216

Co-authored-by: VI-Arthur <arthurzucker1@gmail.com>

stevhliu reviewed Sep 8, 2026

View reviewed changes

Comment thread docs/source/en/expert_parallelism.md Outdated

Comment thread docs/source/en/expert_parallelism.md

Comment thread docs/source/en/expert_parallelism.md Outdated

Comment thread docs/source/en/expert_parallelism.md Outdated

qgallouedec and others added 5 commits September 8, 2026 15:48

Update docs/source/en/expert_parallelism.md

8f2f9cb

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Update docs/source/en/expert_parallelism.md

3c8f600

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Update docs/source/en/expert_parallelism.md

e7be30a

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Update docs/source/en/expert_parallelism.md

8dbccad

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Merge branch 'fsdp2-ep-2d-mesh' into ep-token-dispatch

79fe11d

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

Drop the HF_EP_DISPATCH prototype call sites superseded by #48518

b89b0e6

qgallouedec added 2 commits September 9, 2026 00:44

AmineDiro mentioned this pull request Sep 9, 2026

New example: [REDACTED] huggingface/trl#7121

Open

Uh oh!

Conversation

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

Uh oh!

Why

Results

Correctness

Limitations

Uh oh!

HuggingFaceDocBuilderDev commented Sep 4, 2026

Uh oh!

AmineDiro commented Sep 4, 2026

Uh oh!

IlyasMoutawwakil commented Sep 4, 2026

Uh oh!

AmineDiro Sep 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

qgallouedec commented Sep 7, 2026

Uh oh!

VI-Arthur left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

stevhliu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Sep 9, 2026

CI recap

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

qgallouedec commented Sep 4, 2026 •

edited by github-actions Bot

Loading

AmineDiro Sep 7, 2026 •

edited

Loading