🚨 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.
- The error spells out the full replacement plan (ep_dispatch_experts on the experts, grouped_gemm on their weights).
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, and its output always all-reduces: the guard that skipped the reduce for replicated experts has no user, since every moe_tp_experts plan shards the expert weights.
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. Its experts still shard with tensor parallelism (tp_size), as before. Moving them to the standard experts interface, so EP can dispatch to them, is a follow-up.
Docs:
- The masking rows and the ep_router plan examples are replaced with ep_implementation="all_reduce".
- EpRouterParallel leaves the parallel styles table.
- A note that Llama 4 shards its experts with TP only.

🚨 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; use tp_size to shard its experts.

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`
After the review round: `tests/expert_parallel`, `tests/tensor_parallel` and the TP/EP tests of Mixtral, Qwen3-MoE and Llama 4 on CPU	60 passed
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

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

gotcha, fair enough we still have tp which is the same then for now (would mention this tho!)

added a comment in the config where the ep plan used to be, a line in the docs, and a line in the breaking changes of the description

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?

vasqu · 2026-11-12T12:16:12Z


-Enable expert parallelism with the [`DistributedConfig`] class and the `ep_size` argument. Most MoE models default to [token dispatch](#token-dispatch), so `ep_size` can be set independently of `tp_size`. A few, such as Llama 4, still default to masking and all-reduce (`"ep_router"` and `"moe_tp_experts"` in `model.ep_plan`). Masking is also available on any model with an `ep_plan` override, and requires `ep_size=tp_size` so every rank in an expert group receives the same tokens.
-
+Enable expert parallelism with the [`DistributedConfig`] class and the `ep_size` argument. MoE models shard their experts with `"ep_dispatch_experts"`. Its implementations pick how tokens reach the experts. [Token dispatch](#token-dispatch) lets `ep_size` differ from `tp_size`. Masking and all-reduce requires `ep_size=tp_size`, so every rank of an expert group holds the same tokens.


Suggested change

Enable expert parallelism with the [`DistributedConfig`] class and the `ep_size` argument. MoE models shard their experts with `"ep_dispatch_experts"`. Its implementations pick how tokens reach the experts. [Token dispatch](#token-dispatch) lets `ep_size` differ from `tp_size`. Masking and all-reduce requires `ep_size=tp_size`, so every rank of an expert group holds the same tokens.

Enable expert parallelism with the [`DistributedConfig`] class and the `ep_size` argument. MoE models shard their experts with `"ep_dispatch_experts"`. Its implementations pick how tokens reach the experts. [Token dispatch](#token-dispatch) lets `ep_size` differ from `tp_size`. A few EP implementations require `ep_size=tp_size`, e.g. `all-reduce`, so every rank of an expert group holds the same tokens.

thinking of things like deep ep, mega moe etc

took it, naming both "all_reduce" and "deepgemm_megamoe" (mega moe also hands out slices of the shared batch)

vasqu · 2026-11-12T12:21:48Z

 Providing a plan does not infer parallel sizes: set `tp_size` and `ep_size` explicitly.

-Most MoE models default to `"ep_dispatch_experts"`, which already uses `"all_reduce"` at inference when `ep_size=tp_size`. To use masking and all-reduce in training too, set `ep_size=tp_size` and override both the router and the expert forward rules (the router module name depends on the model, e.g. `mlp.router` on gpt-oss):
+The default `"auto"` implementation already uses `"all_reduce"` at inference when `ep_size=tp_size`. To use masking and all-reduce in training too, set `ep_implementation="all_reduce"`:


the already sounds a bit off / confusing no?

yeah, reworded: "With ep_size=tp_size, the default "auto" implementation runs "all_reduce" at inference and dispatches tokens in training. To mask and all-reduce in training too, set ep_implementation="all_reduce""

vasqu · 2026-11-12T12:22:38Z

 ## 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 their own meshes, `tp` and `ep`, 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"`.


Suggested change

Tensor and expert parallelism shard the weights across their own meshes, `tp` and `ep`, 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"`.

Tensor and expert parallelism shard the weights across their own meshes, `tp` and `ep`, 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 another mesh dimension with `fsdp_size`. With masking and all-reduce, keep `ep_size=tp_size` and pass `ep_implementation="all_reduce"`.

vasqu · 2026-11-12T12:23:19Z

-        "model.layers.*.mlp.gate": "ep_router",
-        "model.layers.*.mlp.experts": "moe_tp_experts",
-    },
+    ep_implementation="all_reduce",


ig its technically not necessary? auto should be enough

in training auto dispatches tokens (all-to-all), and the Trainer doesn't handle dispatch yet (the warning above). this example is the masked + all-reduce training setup, so it needs the explicit "all_reduce". added a comment on that line saying so

vasqu · 2026-11-12T12:25:31Z

+                "`moe_tp_experts` is a tensor-parallel style: keep it in `tp_plan`, or name the experts "
+                '`ep_dispatch_experts` in `ep_plan`. `DistributedConfig(ep_implementation="all_reduce")` gives the '
+                "masking and all-reduce the old `ep_router` + `moe_tp_experts` plan did."


Should we nudge towards a full plan example? grouped_gemm is technically also one slight difference

merged the two errors into one that spells out the full plan: experts -> ep_dispatch_experts, gate_up_proj / down_proj -> grouped_gemm, and says most moe models default to it so ep_plan can usually be left out

vasqu · 2026-11-12T12:26:44Z

-        return module
-
-    def transform_output_post_forward(self, module, output, mesh, source=None):
+    def transform_output_post_forward(self, module, output, mesh):


do we even need this override anymore?

the all-reduce itself is still needed: the base transform_output_post_forward returns the output as is, and with tp each rank only computes its shard of every expert, so the outputs are partial sums. but the "any sharded param" guard in it (skip the reduce when the experts are replicated) has no user anymore: every built-in moe_tp_experts plan also has rules that shard gate_up_proj / down_proj, and every tp style shards through distribute_tensor. so i dropped the guard and it always reduces

vasqu · 2026-11-12T12:27:42Z

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


gotcha, fair enough we still have tp which is the same then for now (would mention this tho!)

vasqu · 2026-11-12T12:28:23Z

        self.assertEqual(model.tp_plan, {"model.layers.*.self_attn.q_proj": "colwise"})

+    @parameterized.expand([("ep_router", "model.layers.*.mlp.gate"), ("moe_tp_experts", "model.layers.*.mlp.experts")])
+    def test_legacy_masked_ep_plan_is_refused(self, style, module):


moved it to the end of TestExpertParallelLayer in tests/expert_parallel/test_expert_parallel.py

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

github-actions · 2026-11-12T16:47:19Z

[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

Merged

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

Merged

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 changed the title ~~Remove the legacy ep_router expert-parallel plan~~ 🚨 Remove the ep_router expert-parallel style Nov 2, 2026

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

vasqu approved these changes Nov 12, 2026

View reviewed changes

IlyasMoutawwakil enabled stacked auto-merge November 12, 2026 15:02

IlyasMoutawwakil added this pull request to the merge queue Nov 12, 2026

github-merge-queue Bot removed this pull request from the merge queue because a pull request earlier in the stack was removed Nov 12, 2026

IlyasMoutawwakil enabled stacked auto-merge November 12, 2026 16:01

IlyasMoutawwakil added this pull request to the merge queue Nov 12, 2026

github-merge-queue Bot removed this pull request from the merge queue because a pull request earlier in the stack was unmerged and was not in the queue Nov 12, 2026

Base automatically changed from ep-mxfp4-gpt-oss to main November 12, 2026 16:43

🚨 Remove the ep_router expert-parallel style

78e10dc

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

IlyasMoutawwakil force-pushed the ep-remove-router branch from e7e8f75 to 78e10dc Compare November 12, 2026 16:45

vasqu enabled stacked auto-merge November 12, 2026 16:46

vasqu added this pull request to the merge queue Nov 12, 2026

Merged via the queue into main with commit 506ac61 Nov 12, 2026
124 checks passed

vasqu deleted the ep-remove-router branch November 12, 2026 17:11

hmellor mentioned this pull request Nov 13, 2026

[CI] Bump Transformers version to 5.23.0 vllm-project/vllm#67502

Merged

6 tasks done


		Enable expert parallelism with the [`DistributedConfig`] class and the `ep_size` argument. Most MoE models default to [token dispatch](#token-dispatch), so `ep_size` can be set independently of `tp_size`. A few, such as Llama 4, still default to masking and all-reduce (`"ep_router"` and `"moe_tp_experts"` in `model.ep_plan`). Masking is also available on any model with an `ep_plan` override, and requires `ep_size=tp_size` so every rank in an expert group receives the same tokens.

		Enable expert parallelism with the [`DistributedConfig`] class and the `ep_size` argument. MoE models shard their experts with `"ep_dispatch_experts"`. Its implementations pick how tokens reach the experts. [Token dispatch](#token-dispatch) lets `ep_size` differ from `tp_size`. Masking and all-reduce requires `ep_size=tp_size`, so every rank of an expert group holds the same tokens.

	Tensor and expert parallelism shard the weights across their own meshes, `tp` and `ep`, 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"`.
	Tensor and expert parallelism shard the weights across their own meshes, `tp` and `ep`, 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 another mesh dimension with `fsdp_size`. With masking and all-reduce, keep `ep_size=tp_size` and pass `ep_implementation="all_reduce"`.

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!

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!

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

IlyasMoutawwakil commented Oct 14, 2026 •

edited

Loading

HuggingFaceDocBuilderDev commented Oct 14, 2026 •

edited

Loading

IlyasMoutawwakil Nov 5, 2026 •

edited

Loading