🚨 Remove the `ep_router` expert-parallel style by IlyasMoutawwakil · Pull Request #49615 · huggingface/transformers

IlyasMoutawwakil · 2026-10-14T19:21:29Z

Depends on #49613.

After #49607 and #49613, no model's EP plan uses the router-masked plan (ep_router on the router, moe_tp_experts on the experts). ep_dispatch_experts covers it: DistributedConfig(ep_implementation="all_reduce") gives the same masking and all-reduce (#49429). So this PR removes the ep_router style.

Changes

distributed/expert_parallel.py: EpRouterParallel and its "ep_router" registration are removed.
distributed/mixin.py:
- The masked-EP branch is removed.
- A plan that names ep_router raises an error that points to ep_dispatch_experts and ep_implementation="all_reduce".
- So does an ep_plan that names moe_tp_experts: without ep_router it would apply nothing and leave the experts unsharded. moe_tp_experts keeps working from tp_plan.
distributed/configuration_utils.py: the ep_size=tp_size check of the masked plan is removed. AllReduceEP checks it.
distributed/tensor_parallel.py: moe_tp_experts loses its is_expert_parallel flag, which no caller passes. It also loses the install_forward override that only passed that flag.
integrations/tensor_parallel.py: the RouterParallel and EpRouterParallel aliases are removed.
Llama 4:
- Its EP plan is removed. It mixed TP styles with an ep_router rule and had no experts rule, so it never applied. With ep_size > 1, attention and the experts stayed replicated.
- Its experts take only hidden_states, so they cannot use ep_dispatch_experts. ep_size > 1 now raises the "no ep_plan" error.
Docs:
- The masking rows and the ep_router plan examples are replaced with ep_implementation="all_reduce".
- EpRouterParallel leaves the parallel styles table.

🚨 Breaking changes

The ep_router style is removed. A plan that names it raises an error with the replacement.
moe_tp_experts in an ep_plan no longer selects masked EP; it raises an error with the replacement. In a tp_plan it works as before.
EpRouterParallel and its RouterParallel alias are removed.
Llama 4 has no EP plan. ep_size > 1 raises instead of silently replicating.

Tests

Test	Result
`tests/expert_parallel` and `tests/tensor_parallel` on CPU, including both legacy plans being refused	32 passed
TP and EP tests of Llama 4, Mixtral, GPT-OSS, Qwen3-MoE, DeepSeek-V3 and Gemma 4 on CPU	432 passed, 1 error: a Gemma 4 processor test needs `tiktoken`
ruff, the noisy-comment, modular and copies checks	clean

🤖 Generated with Claude Code

HuggingFaceDocBuilderDev · 2026-10-14T19:28:04Z

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.

Pages changed in this PR:

vasqu

Not sure whether to fully remove or keep a BC reference at least? and warn when it is used -> remap to dispatch

Every model's experts run on `ep_dispatch_experts`, except Llama 4's, whose EP plan matched no expert-parallel mode and was never applied; its `ep_router` rule goes too. `DistributedConfig(ep_strategy="all_reduce")` gives the same masking and all-reduce, and a plan that still names `ep_router` now says so. `moe_tp_experts` stays for tensor-parallel experts; only its unused expert-parallel flag goes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The class, its registration and the docs describing the router-masked plan were still there. A plan that names `ep_router` now raises an error pointing to `ep_dispatch_experts` and `ep_implementation="all_reduce"` instead of a bare KeyError. `moe_tp_experts` loses its `is_expert_parallel` flag (no caller passes it) and the `install_forward` override that only existed to pass it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

# Conflicts: # docs/source/en/expert_parallelism.md

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

3outeille · 2026-11-04T07:46:45Z

I think it's best to fully remove it otherwise that will create so much confusion

3outeille

small nit but lgtm

3outeille · 2026-11-04T07:47:15Z

-        "layers.*.feed_forward.gate_proj": "colwise",
-        "layers.*.feed_forward.up_proj": "colwise",
-        "layers.*.feed_forward.down_proj": "rowwise",
-        "layers.*.feed_forward.router": "ep_router",


is is the only file that has ep_router left ?

No: ep_router was in nine model configs (Beam, DiffusionGemma, Gemma 4, HY-V4, Inkling, Llama 4, Mellum, MiMo-V2-Flash, Qwen3-MoE), and this PR removes it from all of them, modular files included. Llama 4 loses its whole base_model_ep_plan (4bad170). After this PR the only ep_router left in the repo is the error mixin.py raises for a user plan that still names it.

can we not dispatch on llama4's moe instead?

not without refactoring llama4's moe: its experts take every token already scaled by its router score (hidden_states only), not (hidden_states, top_k_index, top_k_weights), so there's nothing for dispatch to route. moving it to the standard experts interface is a follow-up

The base inlined the router's backward all-reduce check; this branch removes the `ep_router` style, so the class stays gone. In `MoeExpertsParallel`, the routing weights' backward all-reduce is the inline check, without the expert-parallel condition this branch removes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

vasqu

Mostly agree here, it's better to remove this indeed otherwise we keep some weird references overall

vasqu · 2026-11-06T14:45:45Z

 | Configuration | Result |
 | :--- | :--- |
-| `DistributedConfig(tp_size=4, fsdp_size=2, ep_size=4)` | Each TP group of four shares one batch. Training dispatches tokens, each rank taking a quarter of the batch; inference uses `"all_reduce"`. A masked `ep_plan` uses masking and all-reduce throughout. |
+| `DistributedConfig(tp_size=4, fsdp_size=2, ep_size=4)` | Each TP group of four shares one batch. Training dispatches tokens, each rank taking a quarter of the batch; inference uses `"all_reduce"`; `ep_implementation="all_reduce"` uses it throughout. |


setting ...

vasqu · 2026-11-06T14:46:16Z

 | `DistributedConfig(tp_size=1, fsdp_size=8, ep_size=4)` | Dispatch with an independent batch on each rank, no slicing, and experts FSDP-sharded across pairs of ranks. |
 | `DistributedConfig(tp_size=2, fsdp_size=4, ep_size=4)` | Dispatch with a TP pair per batch, each pair slicing its batch in two; two batches per expert group. |
-| `DistributedConfig(tp_size=8, ep_size=8)` | Every rank shares one batch. Training dispatches tokens, each rank taking an eighth of the batch; inference uses `"all_reduce"`. A masked `ep_plan` uses masking and all-reduce throughout. |
+| `DistributedConfig(tp_size=8, ep_size=8)` | Every rank shares one batch. Training dispatches tokens, each rank taking an eighth of the batch; inference uses `"all_reduce"`; `ep_implementation="all_reduce"` uses it throughout. |


vasqu · 2026-11-06T14:47:05Z

 ## Combining with FSDP2

-Tensor and expert parallelism shard the weights across `tp`, but the optimizer state and the modules without a rule are still replicated on every rank of the group, which limits how large a model you can train. Add [FSDP2](./fsdp) on a second mesh dimension with `fsdp_size`. With masking and all-reduce, keep `ep_size=tp_size` and pass `ep_plan={"layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts": "moe_tp_experts"}``.
+Tensor and expert parallelism shard the weights across `tp`, but the optimizer state and the modules without a rule are still replicated on every rank of the group, which limits how large a model you can train. Add [FSDP2](./fsdp) on a second mesh dimension with `fsdp_size`. With masking and all-reduce, keep `ep_size=tp_size` and pass `ep_implementation="all_reduce"`.


is this still correct we shard across their own meshes tp/ep no?

right, fixed: they shard across their own meshes, tp and ep

vasqu · 2026-11-06T14:50:03Z

-            if {"ep_router", "moe_tp_experts"}.issubset(ep_plan.values()):
-                # Legacy masked EP: the EP group is the TP group, every rank keeps every token.
-                model = apply_tensor_parallelism(model, tp_mesh, ep_plan)


this is a weird artifect ig if we keep moe tp experts on tp plan it still works as expected only adding it to ep plan is properly ignored right?

this is somewhat breaking as well so best to mention this change more explicitly

yes, in the tp plan it works as before. in the ep plan it only meant something next to ep_router (the masked path), and alone it would now apply nothing and leave the experts unsharded, so it raises and points to ep_dispatch_experts / ep_implementation="all_reduce". added it to the breaking changes in the description, with a test

vasqu · 2026-11-06T14:50:57Z

-                if torch.is_grad_enabled() and top_k_weights.requires_grad:
-                    top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group)
+            if torch.is_grad_enabled() and top_k_weights.requires_grad:
+                top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group)


we should use our helper here then depending on what we do on 1 in the stack

done, it calls _AllReduceBackward.apply unconditionally now, the grad check lives inside it since #49429

vasqu · 2026-11-06T14:51:59Z

-        "layers.*.feed_forward.gate_proj": "colwise",
-        "layers.*.feed_forward.up_proj": "colwise",
-        "layers.*.feed_forward.down_proj": "rowwise",
-        "layers.*.feed_forward.router": "ep_router",


can we not dispatch on llama4's moe instead?

# Conflicts: # src/transformers/distributed/expert_parallel.py # src/transformers/distributed/tensor_parallel.py

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

github-actions · 2026-11-12T10:47:50Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: llama4

IlyasMoutawwakil mentioned this pull request Oct 14, 2026

🚨 Run MXFP4 GPT-OSS experts through ep_dispatch_experts #49613

Open

IlyasMoutawwakil added this pull request to stack #49611 October 14, 2026 19:22

This was referenced Oct 14, 2026

Expert-parallel implementations for ep_dispatch_experts: all-reduce at inference, all-to-all in training #49429

Merged

Move the remaining router-masked MoE plans to ep_dispatch_experts #49607

Open

This was referenced Oct 15, 2026

Port Llama 4's MoE to the shared experts interface #49432

Closed

Add the deepep expert-parallel implementation #49638

Open

IlyasMoutawwakil force-pushed the ep-remove-router branch from 7d13b07 to 1cbdb8c Compare October 16, 2026 06:56

IlyasMoutawwakil changed the title ~~Remove the ep_router expert-parallel style~~ Remove the legacy ep_router expert-parallel plan Oct 16, 2026

vasqu reviewed Oct 19, 2026

View reviewed changes

IlyasMoutawwakil force-pushed the ep-remove-router branch from 1cbdb8c to 9845688 Compare October 27, 2026 21:49

IlyasMoutawwakil and others added 4 commits October 27, 2026 15:42

Merge the updated MXFP4 GPT-OSS PR

90570c8

# Conflicts: # docs/source/en/expert_parallelism.md

Drop Llama 4's EP plan, which never applied

4bad170

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Merge the updated MXFP4 GPT-OSS PR

2cbebd4

IlyasMoutawwakil changed the title ~~Remove the legacy ep_router expert-parallel plan~~ 🚨 Remove the ep_router expert-parallel style Nov 2, 2026

Merge the Mega MoE docs fix

3a37bd2

3outeille approved these changes Nov 4, 2026

View reviewed changes

qgallouedec mentioned this pull request Nov 4, 2026

remove ep_router #49587

Open

vasqu approved these changes Nov 6, 2026

View reviewed changes

IlyasMoutawwakil and others added 5 commits November 12, 2026 11:25

Merge branch 'ep-mxfp4-gpt-oss' into ep-remove-router

8020209

# Conflicts: # src/transformers/distributed/expert_parallel.py # src/transformers/distributed/tensor_parallel.py

Merge branch 'ep-mxfp4-gpt-oss' into ep-remove-router

37cfed9

Refuse a legacy moe_tp_experts rule in ep_plan; docs wording

2309c2f

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Merge branch 'ep-mxfp4-gpt-oss' into ep-remove-router

6149bb5

Merge branch 'ep-mxfp4-gpt-oss' into ep-remove-router

81da3fe

Trim comments

bb5ccea

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

IlyasMoutawwakil added 2 commits November 12, 2026 11:58

Merge branch 'ep-mxfp4-gpt-oss' into ep-remove-router

5622cd0

Merge branch 'ep-mxfp4-gpt-oss' into ep-remove-router

d08be28

Conversation

IlyasMoutawwakil commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

🚨 Breaking changes

Tests

Uh oh!

HuggingFaceDocBuilderDev commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

3outeille commented Nov 4, 2026

Uh oh!

3outeille left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

IlyasMoutawwakil Nov 5, 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!

Choose a reason for hiding this comment

Uh oh!

vasqu left a comment

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!

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!

github-actions Bot commented Nov 12, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

IlyasMoutawwakil commented Oct 14, 2026 •

edited

Loading

HuggingFaceDocBuilderDev commented Oct 14, 2026 •

edited

Loading

IlyasMoutawwakil Nov 5, 2026 •

edited

Loading