Expert-parallel implementations for `ep_dispatch_experts`: all-reduce at inference, all-to-all in training by IlyasMoutawwakil · Pull Request #49429 · huggingface/transformers

IlyasMoutawwakil · 2026-10-08T13:09:55Z

"ep_dispatch_experts" reaches its experts through an implementation that DistributedConfig.ep_implementation picks. Inference with ep_size == tp_size no longer pays an all-to-all and a host sync per layer, so decode captures in a CUDA graph.

Implementations

`ep_implementation`	How tokens reach the experts	Needs	CUDA graph
`"all_to_all"`	Token dispatch with two all-to-alls, the previous behaviour.	Any layout.	No, one host sync per layer.
`"all_reduce"`	Each rank masks the routes to other ranks' experts on the shared batch. One all-reduce sums the outputs.	`ep_size == tp_size`	Yes
`"auto"`, the default	`"all_reduce"` at inference on a shared batch, `"all_to_all"` otherwise.		At inference

The implementations live in ALL_EP_IMPLEMENTATIONS. A new one subclasses EPImplementation, implements dispatch and combine, and registers under a name.

Changes

File	Change
`distributed/expert_parallel.py`	`EPImplementation`, `AllToAllEP` (the former dispatch code), `AllReduceEP`, `AutoEP` and `ALL_EP_IMPLEMENTATIONS`. `EpDispatchExpertsParallel` delegates to them and rejects an unknown name.
`distributed/configuration_utils.py`	The `ep_implementation` field, `"auto"` by default.
`distributed/mixin.py`	Passes `ep_implementation` to `apply_expert_parallelism`.
`distributed/parallel_utils.py`	`_same_ranks`, and `_all_reduce_backward`. It runs the backward-only all-reduce only when a gradient can flow, at every TP and EP site that used it: MoE TP experts, MLA rope, router scores and EP inputs. That saves ~1.8 µs of host time per call in eager decode.
`docs/source/en/expert_parallelism.md`	A table of each implementation's advantages and drawbacks, and an "Adding an implementation" section.
`tests/expert_parallel/test_expert_parallel.py`	Each implementation against one process (outputs and input gradients), the `"auto"` choice, gradient checkpointing in both modes, unknown and custom names, and the shared-batch check.

Validation

Test	Result
`tests/expert_parallel/test_expert_parallel.py`, 2 ranks on gloo	7 passed
TP and EP model tests of Mixtral, Qwen3-MoE, GPT-OSS, DeepSeek-V3 and Llama, on CPU	40 passed, 125 skipped
ruff, the docstring and the noisy-comment checks	clean

🤖 Generated with Claude Code

HuggingFaceDocBuilderDev · 2026-10-08T13:16:20Z

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:

en/expert_parallelism

…ers into unify-resolve-plan

This reverts commit 3f1e4ab.

3outeille · 2026-10-15T14:10:55Z

+        ep_group = self.ep_mesh.get_group()
+        # every rank backpropagates through its own experts only: sum the replicated inputs' gradients
+        hidden_states = _AllReduceBackward.apply(hidden_states, ep_group)
+        if torch.is_grad_enabled() and top_k_weights.requires_grad:


how does that work with gradient checkpointing ? (since it's alternating no_grad and grad for the recompute)

yes interesting case, with our default (use_reentrant=False) the first forward runs with grad enabled, so auto picks all-to-all both times. with use_reentrant=True the first forward runs under no_grad and gets the all-reduce, then the recompute gets the all-to-all: same function and same output shapes, so the gradients are right, only bf16 rounding differs between the two passes. i will add a test just in case

why is this workaround needed in the first place? It seems a tad weird to me ngl

not sur what you mean by workaround here

Why do we need to check for if torch.is_grad_enabled() and top_k_weights.requires_grad:?

bcz otherwise it's unnecessary to do the bwd, #49975 seems related

yeah we can remove it now as the hook will only be attached during training

not exactly, #49975 is about the param grad hooks of ReplicatedWithGradAllReduce. this one is the autograd fn on activations (router scores, ep inputs, mla rope..) and it still runs on every forward, so the check is what makes it free at inference (no autograd fn per layer during decode). so i would keep it

can we at least make a proper helper on when we do need this?

added _all_reduce_backward in parallel_utils, it skips the autograd fn when no grad can flow. all the tp/ep call sites use it now

IlyasMoutawwakil · 2026-11-02T15:48:53Z

@dacorvo you can now just pass DistributedConfig(tp_size=n, ep_size=n, ep_implementation="all_reduce"), it works in training too when ep_size == tp_size, so no all_to_all_single at all. for ep_size > tp_size there is no way around a token exchange

3outeille · 2026-11-04T07:23:23Z

            dist.all_reduce(grad, group=ctx.process_group)
            return grad, None

+    def _all_reduce_backward(x, process_group):


We now have a function and class for _all_reduce_backward. Better call if torch.is_grad_enabled() and x.requires_grad: return _AllReduceBackward.apply(x, process_group) in ep implementation ?

Done: the function is gone, and the 7 call sites (router scores and the all-reduce EP inputs in expert_parallel.py, plus the MoE TP experts, the router weights and MLA rope in tensor_parallel.py) check torch.is_grad_enabled() and x.requires_grad before _AllReduceBackward.apply. Its docstring says why.

I like this more tbh but maybe we can hide it directly under _AllReduceBackward as a compromise, wdyt @3outeille

yes we can do that

3outeille · 2026-11-04T07:28:19Z

+                f"{type(self).__name__} needs every rank of the expert group to hold the same batch (`ep_size == tp_size`)."
+            )
+
+    def dispatch(self, hidden_states, top_k_index, top_k_weights):


My only concern is that AllReduceEP doesn't follow the dispatch + combine pattern which can bring confusion. Maybe we can just turn AllToAllEP and AllReduceEP class into simple functions to avoid having to deal with an explicit dispatch + combine pattern API ?

AllReduceEP did follow it: its dispatch picked the routes this rank computes, and its combine summed the outputs. AutoEP was the one that didn't, since it overrode __call__ around both. To make that uniform, every implementation is now a callable, __call__(hidden_states, top_k_index, top_k_weights), with no dispatch / combine contract. They stay classes, since one can keep state across forwards (DeepEP's buffers, for example). A shared run_experts runs the local experts with their local params and keeps empty rows connected to backward. The docs' "Adding an implementation" example is a single __call__.

`AllReduceEP` had nothing to dispatch and `AutoEP` bypassed both steps, so the two-step contract didn't fit. An `EPImplementation` now implements `forward` and runs the local experts with a shared `run_experts` (local DTensor params, empty rows kept connected to backward). `AllToAllEP` keeps its two exchanges as private `_send` / `_receive` steps. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Every `EPImplementation` defines `__call__` (AutoEP already did), and the base class keeps `run_experts`, now owning the local DTensor params context. `_all_reduce_backward` is gone: its call sites check `torch.is_grad_enabled() and x.requires_grad` before `_AllReduceBackward.apply`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

3outeille

nice work almost there, small changes and it should be good to go

- `EpDispatchExpertsParallel.ep_implementations` maps the three built-in names to their classes, replacing `EPImplementationInterface` / `ALL_EP_IMPLEMENTATIONS`; docs and the config docstring point to it. - `AllToAllEP`'s halves are `_dispatch_tokens` and `_combine_tokens` again. - `EpDispatchExpertsParallel.should_use_local_tensors` is back. - Tests, from review: one parametrized check per implementation and mode (train, inference, both checkpoint modes) against one process, a separate check of AutoEP's choice, and install_forward's refusals. Co-authored-by: Ferdinand Mom <47445085+3outeille@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

`_same_ranks` had one caller; the comparison of the TP and EP groups' ranks now sits where `shares_batch` is set. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

vasqu

Ok overall very nice state, chiming in to give my 2 cents but I think it's close to merge

vasqu · 2026-11-05T14:59:45Z

        # before the mask (each slot has exactly one owning rank, so the sum is exact).
+        process_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp")
        if torch.is_grad_enabled() and router_scores.requires_grad:
-            process_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp")


why did we change the order here? doesnt have any positive effect no?

i made this "if grad: all reduce backward" a helper function before but we inlined it after, this is just the result of that revert, now it just aligns with the pattern without the process_group line in the middle:

if torch.is_grad_enabled() and x.requires_grad: x = _AllReduceBackward.apply(x, process_group)

Is it expensive to check the process group? Ig its negligible, just wanna be sure it doesnt cause an unnecessary overhead on inference

vasqu · 2026-11-05T15:00:38Z

+    """
+
+    # Set by implementations that only run when every rank of the EP group holds the same batch, e.g. `AllReduceEP`.
+    needs_shared_batch = False


ig deepgeemm and deep ep will also need it?

no, only implementations that skip the token exchange need it. deep ep dispatches tokens like all_to_all so it works whatever batch each rank holds, and megamoe is its own style (megamoe_experts), not an EPImplementation

vasqu · 2026-11-05T15:02:26Z

+    def __init__(self, experts, ep_mesh, tp_mesh=None):
+        self.experts = experts
+        self.experts_forward = experts.forward  # the experts' own forward, on local expert ids
+        self.num_local_experts = experts.num_experts


are we sure num experts is safe for all models? do we need something like in fp8 checking for attrs?

yes, the experts behind @use_experts_implementation all have num_experts (the sentinel masking reads it), and grouped_gemm sharding already rewrites it to the local count

vasqu · 2026-11-05T15:02:39Z

+    needs_shared_batch = False
+
+    def __init__(self, experts, ep_mesh, tp_mesh=None):
+        self.experts = experts


why are we saving the experts?

run_experts needs the module to swap its dtensor params for local ones (_use_local_dtensor_params(self.experts)), and an implementation with its own kernel would read the weights from it

vasqu · 2026-11-05T15:04:38Z

    """
-    Dispatch disjoint TP token slices to the experts' owners, then replicate the combined output on TP.
+    The ranks of a TP group keep disjoint slices of the batch they share, send each (token, expert) pair to the
+    expert's rank with an all-to-all and get the results back with a second one, then gather the slices back over TP.


so we always have a TP mesh no matter what, just to be sure I follow

yeah it's a proper EP and it needs the tp_size to avoid computing tp_size time the same token. Cf the explanation here: #48865

yep always there, size 1 without tp. made tp_mesh required and dropped the None branches

vasqu · 2026-11-05T15:13:00Z

+        def ep_forward(hidden_states, top_k_index, top_k_weights):
+            if isinstance(hidden_states, DTensor):
+                hidden_states = hidden_states.to_local()
+            if isinstance(top_k_weights, DTensor):
+                top_k_weights = top_k_weights.to_local()
+            return implementation(hidden_states, top_k_index, top_k_weights).to(hidden_states.dtype)


feels a bit weird, maybe we can reuse the local_params style within tp?

these are the inputs not the params, params already go through _use_local_dtensor_params in run_experts (same as local_params). the input unwrap is carried over from #48865, MoeExpertsParallel does the same in its pre-forward. can move both to a shared helper in a follow up

yea that would be nice, indeed it's the input mb

any call to parallel function should in general / priority be done here

done, the dtensor unwrapping and the local params both live in ep_forward now, the implementations only move tokens and call the experts

vasqu · 2026-11-05T15:13:24Z

            # Dispatch hooks need both meshes to redistribute tokens between TP and EP ranks.
            if style_name == "ep_dispatch_experts":
-                style.install_forward(module, ep_mesh=ep_mesh, tp_mesh=tp_mesh)
+                style.install_forward(module, ep_mesh, tp_mesh=tp_mesh, ep_implementation=ep_implementation)


lets keep kwargs styöe on ep mesh?

vasqu · 2026-11-05T15:13:59Z

yea looks like we should do a proper helper

done, _all_reduce_backward

vasqu · 2026-11-05T15:16:10Z

+        assert (all_reduce_call.called, all_to_all_call.called) == (all_reduce, not all_reduce), (tp_mesh, grad)
+
+
+class _ScaledExperts(torch.nn.Module):


why an extra one and not include it as optional kwarg within dummy experts?

merged them, _DummyExperts(scales=None)

vasqu · 2026-11-05T15:17:00Z

+        output = checkpoint(forward, experts, *inputs, use_reentrant=False)
+    elif mode == "reentrant_checkpoint":
+        output = checkpoint(forward, experts, *inputs, use_reentrant=True)


imo reentrant true behaves very weird and is actively discouraged for the most part; do we really want to test it?

it is exactly the gradient checkpointing case from above: auto runs all_reduce on the no_grad first pass and all_to_all on the recompute, so it is the one place both implementations get mixed. users can still pass use_reentrant=True so i would keep it, it is cheap

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

…xperts class in tests Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

vasqu

Ok I think this is ready, some last comments but nothing major anymore

vasqu · 2026-11-06T12:27:48Z

+| :--- | :--- | :--- |
+| `"all_to_all"`, see [token dispatch](#token-dispatch) | Works for any layout, including `ep_size > tp_size`. Only the routed tokens move. | Two all-to-all collectives and one host sync per layer, so it does not capture in a CUDA graph. |
+| `"all_reduce"` | No tokens move and no host sync, so decode captures in a CUDA graph. | Requires `ep_size=tp_size`. Every rank runs its experts over the whole batch, so its cost grows with the batch. |
+| `"auto"`, the default | Runs `"all_reduce"` at inference when the ranks of an expert group share one batch, and `"all_to_all"` otherwise. | |


Disadvantage might be that this is not always ideal and your workload can benefit from custom implementations at times

vasqu · 2026-11-06T13:26:09Z

+| `"all_reduce"` | No tokens move and no host sync, so decode captures in a CUDA graph. | Requires `ep_size=tp_size`. Every rank runs its experts over the whole batch, so its cost grows with the batch. |
+| `"auto"`, the default | Runs `"all_reduce"` at inference when the ranks of an expert group share one batch, and `"all_to_all"` otherwise. | |
+
+`ep_implementation` only chooses within the parallel sizes set in [`DistributedConfig`]. When the ranks of an expert group hold different batches, for example `tp_size=1` with `fsdp_size=ep_size`, `"all_to_all"` is the only implementation that applies.


Maybe good to wrap as note

vasqu · 2026-11-06T13:44:48Z

+        expert_output = self.run_experts(tokens, expert_ids, weights)
+        return ...  # the weighted expert output for each token of the batch this rank holds
+
+EpDispatchExpertsParallel.ep_implementations["my_ep"] = MyEP


do we not have a proper function to register?

vasqu · 2026-11-06T13:45:58Z

+distributed_config = DistributedConfig(tp_size=1, fsdp_size=8, ep_size=8, ep_implementation="my_ep")
+```
+
+Gradients must reach every rank that holds a token, routing weights included, so wrap a collective that autograd does not differentiate in a `torch.autograd.Function`.


Also more appropriate as a note imo

vasqu · 2026-11-06T13:52:04Z

            Number of devices owning distinct expert shards. Defaults to 1. Set it explicitly to enable EP. Must be
            a multiple of `tp_size` and divide `fsdp_size * tp_size`. All-reduce expert plans require
-            `ep_size=tp_size`; token dispatch (`"ep_dispatch_experts"`) also allows `ep_size > tp_size`.
+            `ep_size=tp_size`; `"ep_dispatch_experts"` also allows `ep_size > tp_size`, with all-to-all token dispatch.


Imo this whole last sentence needs a rewrite no? ep dispatch experts is non telling with this change as the underlying implementation is important and id rather we go over that no?

vasqu · 2026-11-06T13:56:55Z

            dist.all_reduce(grad, group=ctx.process_group)
            return grad, None

+    def _all_reduce_backward(x, process_group):


I like this more tbh but maybe we can hide it directly under _AllReduceBackward as a compromise, wdyt @3outeille

vasqu · 2026-11-06T13:58:43Z

+            return self.experts_forward(tokens, expert_ids, weights)
+
+
+class _AllToAllState(NamedTuple):


prolly unnecessary worry but is it safe for older torch versions under compile?

vasqu · 2026-11-06T13:59:40Z

+        """Each (token, expert) pair to the expert's rank: the rows this rank's experts run on, their local expert ids
+        and weights, and what `_combine_tokens` needs to bring the results back."""
+        num_tokens = hidden_states.size(0)
+        if self.tp_mesh.size() > 1:


not this PR necessarily but I do feel like we can make the input handling under here more pretty but no worries for now

3outeille

Hiding _all_reduce_backward inside _AllReduceBackward

3outeille · 2026-11-06T15:52:59Z

+
+        Apply it through `_all_reduce_backward`, which skips it when no gradient can flow.
        """



Suggested change

@classmethod

def apply(cls, x, process_group):

if torch.is_grad_enabled() and x.requires_grad:

return super().apply(x, process_group)

return x

3outeille · 2026-11-06T15:53:33Z

            dist.all_reduce(grad, group=ctx.process_group)
            return grad, None

+    def _all_reduce_backward(x, process_group):


Suggested change

def _all_reduce_backward(x, process_group):

…-batch

…face registry, doc notes Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ArthurZucker

one abstraction not super sure it makes sense to register fully or just inherit for now but otherwise yes, ty

ArthurZucker · 2026-11-10T05:59:07Z

+An implementation subclasses `EPImplementation` and registers under the name `ep_implementation` selects. It is built once per experts module, so it can keep state such as communication buffers. Calling it with the batch this rank holds, with global expert ids, returns the weighted expert output for each of its tokens. It runs this rank's experts with `run_experts`, on the rows, local expert ids and weights it brought them.
+
+```py
+from transformers.distributed.expert_parallel import EPImplementation, EPImplementationInterface


Suggested change

from transformers.distributed.expert_parallel import EPImplementation, EPImplementationInterface

from transformers.distributed.expert_parallel import ExpertImplementation, ExpertImplementationInterface

let's use good naming

then it has to be expert prallel implementation, because expert implementation is already a thing (grouped/batched/sonicmoe/etc..)

on second thought kept EPImplementation: everything EP here is ep_ / EP (ep_size, ep_plan, ep_implementation, AllToAllEP, EpDispatchExpertsParallel), and it stays clear of experts_implementation

ArthurZucker · 2026-11-10T05:59:28Z

+        expert_output = self.run_experts(tokens, expert_ids, weights)
+        return ...  # the weighted expert output for each token of the batch this rank holds
+
+EPImplementationInterface.register("my_ep", MyEP)


probably not needed if post_init form super registers with STR of name

i'd rather keep it explicit: registering on subclass definition means just defining a class mutates a global registry (even in tests). EPImplementationInterface is a GeneralInterface, same as AttentionInterface / ExpertsInterface, so registering an ep implementation looks like registering an attention or experts one

ArthurZucker · 2026-11-10T05:59:46Z

            `ep_size > 1`, and its rules take precedence over `tp_plan` rules for the same modules. An
-            `"ep_dispatch_experts"` rule selects all-to-all token dispatch instead of router masking and all-reduce.
+            `"ep_dispatch_experts"` rule lets `ep_implementation` pick how tokens reach the experts.
+        ep_implementation (`str`, *optional*, defaults to `"auto"`):


expert_implementation

kept ep_implementation, same reason as the class name: expert_implementation would be one letter off experts_implementation (grouped_mm, batched_mm, ...) which is a different thing, and ep_ matches ep_size / ep_plan

ArthurZucker · 2026-11-10T06:00:37Z



-class EpDispatchExpertsParallel(ExpertParallelLayer):
+class EPImplementation:


Suggested change

class EPImplementation:

class ExpertImplementation:

kept EPImplementation, see the docs thread: ExpertImplementation would read as an experts_implementation (grouped_mm etc.), and the rest of the EP code uses the EP prefix

ArthurZucker · 2026-11-10T06:01:29Z

+    # Set by implementations that only run when every rank of the EP group holds the same batch, e.g. `AllReduceEP`.
+    needs_shared_batch = False
+
+    def __init__(self, experts, ep_mesh, tp_mesh):


type hint on the expert at least !

ArthurZucker · 2026-11-10T06:08:27Z

+    _global_mapping = {"auto": AutoEP, "all_to_all": AllToAllEP, "all_reduce": AllReduceEP}
+
+
+class EpDispatchExpertsParallel(ExpertParallelLayer):


not super sure we need AutoEP, EPImplementationInterfac.
There's not much shared between both, and some mixing of parallel stuff

kept both: auto is the default path, as its own small class it's readable and tested on its own, instead of branching inside install_forward. the interface is for the same reason as the other thread (same pattern as ExpertsInterface). what did move is the local params context, see below

ArthurZucker · 2026-11-10T06:09:59Z

+
+    def run_experts(self, tokens, expert_ids, weights):
+        """This rank's experts on `tokens`, routed to `expert_ids` (local ids) with `weights`."""
+        with _use_local_dtensor_params(self.experts):


well, it's sure is weird to define this whole interface just to call _use_local_dtensor_params.

I mean, this should be done by the ExpertParallel.forward ?

yep moved it, ep_forward now runs the implementation inside the style's context_around_forward (local params), and run_experts is gone, implementations call experts_forward directly

ArthurZucker · 2026-11-10T06:10:34Z

+        """This rank's experts on `tokens`, routed to `expert_ids` (local ids) with `weights`."""
+        with _use_local_dtensor_params(self.experts):
+            if tokens.shape[0] == 0:
+                # Keep the empty rows connected to backward, so the reverse exchange and FSDP's reduction still run.


correct me if I am wrong, only a single class uses it as off today

right, only all-to-all can receive zero rows, moved it into AllToAllEP

ArthurZucker · 2026-11-10T06:11:24Z

+        def ep_forward(hidden_states, top_k_index, top_k_weights):
+            if isinstance(hidden_states, DTensor):
+                hidden_states = hidden_states.to_local()
+            if isinstance(top_k_weights, DTensor):
+                top_k_weights = top_k_weights.to_local()
+            return implementation(hidden_states, top_k_index, top_k_weights).to(hidden_states.dtype)


any call to parallel function should in general / priority be done here

…y-rows case in `AllToAllEP` Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

3outeille and others added 10 commits October 12, 2026 13:03

build and expand once the tp/ep_plan

6440b72

unify

79b5fdb

clean fsdp

bb63397

reorganize to avoid breaking changes

f9158f8

better distinction

64eb317

Merge branch 'main' into unify-resolve-plan

61a4954

revert

5fbf21c

Merge branch 'unify-resolve-plan' of github.com:huggingface/transform…

ff70527

…ers into unify-resolve-plan

add tp plan to deepseek v4 (but attention stays replicated)

3f1e4ab

Revert "add tp plan to deepseek v4 (but attention stays replicated)"

ff3629c

This reverts commit 3f1e4ab.

IlyasMoutawwakil marked this pull request as ready for review October 14, 2026 11:22

github-actions Bot requested review from ArthurZucker and Rocketknight1 October 14, 2026 11:23

IlyasMoutawwakil changed the title ~~Expert-parallel dispatch runs the experts masked when the EP group shares one batch~~ Expert-parallel strategies: masked all-reduce for small shared batches, all-to-all above, picked per forward Oct 14, 2026

vasqu mentioned this pull request Oct 14, 2026

remove ep_router #49587

Open

IlyasMoutawwakil mentioned this pull request Oct 14, 2026

Expert-parallel dispatch hands the experts the received rows, not the all-to-all's pending result #49421

Merged

6 tasks

IlyasMoutawwakil requested review from 3outeille and qgallouedec October 14, 2026 14:52

IlyasMoutawwakil changed the title ~~Expert-parallel strategies: masked all-reduce for small shared batches, all-to-all above, picked per forward~~ Expert-parallel strategies for ep_dispatch_experts: all-reduce at inference, all-to-all in training Oct 14, 2026

IlyasMoutawwakil commented Oct 14, 2026

View reviewed changes

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

This was referenced Oct 14, 2026

🚨 deepgemm_megamoe EP implementation for DeepGEMM Mega MoE #49606

Merged

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

Merged

IlyasMoutawwakil added this pull request to stack #49611 October 14, 2026 16:49

This was referenced Oct 14, 2026

🚨 Run MXFP4 GPT-OSS experts through ep_dispatch_experts #49613

Merged

🚨 Remove the ep_router expert-parallel style #49615

Merged

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

Closed

3outeille reviewed Oct 15, 2026

View reviewed changes

Comment thread tests/expert_parallel/test_expert_parallel.py Resolved

Comment thread tests/expert_parallel/test_expert_parallel.py Resolved

Comment thread src/transformers/distributed/tensor_parallel.py Resolved

Comment thread src/transformers/distributed/utils.py Resolved

3outeille reviewed Oct 15, 2026

View reviewed changes

Comment thread src/transformers/distributed/expert_parallel.py Resolved

IlyasMoutawwakil requested review from 3outeille and vasqu and removed request for 3outeille and ArthurZucker November 2, 2026 13:56

IlyasMoutawwakil requested a review from ArthurZucker November 2, 2026 16:41

3outeille reviewed Nov 4, 2026

View reviewed changes

IlyasMoutawwakil and others added 2 commits November 4, 2026 12:19

IlyasMoutawwakil requested a review from 3outeille November 4, 2026 12:28

3outeille reviewed Nov 5, 2026

View reviewed changes

IlyasMoutawwakil force-pushed the ep-masked-when-shared-batch branch from 7c1410e to d876395 Compare November 5, 2026 11:04

EP: inline the shared-batch check into EPImplementation

6eae29f

`_same_ranks` had one caller; the comparison of the TP and EP groups' ranks now sits where `shares_batch` is set. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

IlyasMoutawwakil requested a review from 3outeille November 5, 2026 13:19

vasqu reviewed Nov 5, 2026

View reviewed changes

IlyasMoutawwakil and others added 2 commits November 6, 2026 10:37

EP: always take the TP mesh in EP implementations

389f942

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

EP: helper for the gradient-only all-reduce, kwarg ep mesh, one toy e…

4a1cd28

…xperts class in tests Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

vasqu approved these changes Nov 6, 2026

View reviewed changes

3outeille requested changes Nov 6, 2026

View reviewed changes

3outeille approved these changes Nov 6, 2026

View reviewed changes

IlyasMoutawwakil and others added 2 commits November 9, 2026 18:41

Merge remote-tracking branch 'origin/main' into ep-masked-when-shared…

87b97e3

…-batch

EP: grad check inside _AllReduceBackward.apply, EPImplementationInter…

b8c0ec3

…face registry, doc notes Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

3outeille mentioned this pull request Nov 10, 2026

[distributed] Context parallel module, with Ulysses #50493

Draft

ArthurZucker reviewed Nov 10, 2026

View reviewed changes

EP: local params around the implementation call in ep_forward; empt…

2bd83bd

…y-rows case in `AllToAllEP` Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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

Merged via the queue into main with commit b7394bc Nov 12, 2026
142 checks passed

IlyasMoutawwakil deleted the ep-masked-when-shared-batch branch November 12, 2026 10:33

		assert (all_reduce_call.called, all_to_all_call.called) == (all_reduce, not all_reduce), (tp_mesh, grad)


		class _ScaledExperts(torch.nn.Module):

		return self.experts_forward(tokens, expert_ids, weights)


		class _AllToAllState(NamedTuple):


		Apply it through `_all_reduce_backward`, which skips it when no gradient can flow.
		"""

+	 @classmethod
+	  def apply(cls, x, process_group):
+	       if torch.is_grad_enabled() and x.requires_grad:
+	           return super().apply(x, process_group)
+	       return x

	from transformers.distributed.expert_parallel import EPImplementation, EPImplementationInterface
	from transformers.distributed.expert_parallel import ExpertImplementation, ExpertImplementationInterface



		class EpDispatchExpertsParallel(ExpertParallelLayer):
		class EPImplementation:

		_global_mapping = {"auto": AutoEP, "all_to_all": AllToAllEP, "all_reduce": AllReduceEP}


		class EpDispatchExpertsParallel(ExpertParallelLayer):

Conversation

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

Uh oh!

Implementations

Changes

Validation

Uh oh!

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

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

IlyasMoutawwakil commented Nov 2, 2026

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!

3outeille left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

IlyasMoutawwakil commented Oct 8, 2026 •

edited

Loading

HuggingFaceDocBuilderDev commented Oct 8, 2026 •

edited

Loading