EP implementations are callables; inline the backward all-reduce check

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>

Author: IlyasMoutawwakil

docs/source/en/expert_parallelism.md

@@ -64,13 +64,13 @@ The expert forward rule in `ep_plan` selects how tokens reach the experts:
 
 ### Adding an implementation
 
-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. Its `forward` takes the batch this rank holds, with global expert ids, and 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.
+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 ALL_EP_IMPLEMENTATIONS, EPImplementation
 
 class MyEP(EPImplementation):
-    def forward(self, hidden_states, top_k_index, top_k_weights):
+    def __call__(self, hidden_states, top_k_index, top_k_weights):
         tokens, expert_ids, weights = ...  # bring this rank's experts the rows routed to them
         expert_output = self.run_experts(tokens, expert_ids, weights)
         return ...  # the weighted expert output for each token of the batch this rank holds

src/transformers/distributed/expert_parallel.py

@@ -39,7 +39,7 @@
     from torch.distributed._functional_collectives import all_to_all_single_autograd
     from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor
 
-    from .parallel_utils import _all_reduce_backward, _AllReduceForward, _WaitTensor
+    from .parallel_utils import _AllReduceBackward, _AllReduceForward, _WaitTensor
 
 
 class ExpertParallelLayer:
@@ -173,7 +173,8 @@ def transform_output_post_forward(self, module, output, mesh):
         # Each rank's score gradient covers only its local experts' slots; sum the per-rank partials
         # 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")
-        router_scores = _all_reduce_backward(router_scores, process_group)
+        if torch.is_grad_enabled() and router_scores.requires_grad:
+            router_scores = _AllReduceBackward.apply(router_scores, process_group)
         non_local_mask = (router_indices // num_local_experts) != ep_rank
         router_scores = router_scores.masked_fill(non_local_mask, 0.0)
         router_indices = router_indices.masked_fill(non_local_mask, -1)
@@ -203,9 +204,9 @@ class EPImplementation:
     the plan is applied, so an implementation can keep state across forwards (communication buffers, a layout it
     reuses).
 
-    `forward(hidden_states, top_k_index, top_k_weights)` takes the batch this rank holds (global expert ids) and returns
-    the weighted expert output for each of its tokens. It runs this rank's experts with `run_experts`, on whatever rows,
-    local expert ids and weights it brought them.
+    Calling it with `(hidden_states, top_k_index, top_k_weights)`, the batch this rank holds (global expert ids),
+    returns the weighted expert output for each of its tokens. It runs this rank's experts with `run_experts`, on
+    whatever rows, local expert ids and weights it brought them.
 
     Gradients must reach every rank that holds a token, routing weights included, so a collective that autograd does not
     differentiate has to be wrapped in a `torch.autograd.Function`.
@@ -226,22 +227,18 @@ def __init__(self, experts, ep_mesh, tp_mesh=None):
                 f"{type(self).__name__} needs every rank of the expert group to hold the same batch (`ep_size == tp_size`)."
             )
 
-    def forward(self, hidden_states, top_k_index, top_k_weights):
+    def __call__(self, hidden_states, top_k_index, top_k_weights):
         raise NotImplementedError
 
     def run_experts(self, tokens, expert_ids, weights):
         """This rank's experts on `tokens`, routed to `expert_ids` (local ids) with `weights`."""
-        if tokens.shape[0] == 0:
-            # Keep the empty rows connected to backward, so the reverse exchange and FSDP's reduction still run.
-            return self.experts_forward(
-                tokens.sum(dim=0, keepdim=True), expert_ids.new_zeros((1, 1)), weights.new_ones((1, 1))
-            )[:0]
-        return self.experts_forward(tokens, expert_ids, weights)
-
-    def __call__(self, hidden_states, top_k_index, top_k_weights):
         with _use_local_dtensor_params(self.experts):
-            output = self.forward(hidden_states, top_k_index, top_k_weights)
-        return output.to(hidden_states.dtype)
+            if tokens.shape[0] == 0:
+                # Keep the empty rows connected to backward, so the reverse exchange and FSDP's reduction still run.
+                return self.experts_forward(
+                    tokens.sum(dim=0, keepdim=True), expert_ids.new_zeros((1, 1)), weights.new_ones((1, 1))
+                )[:0]
+            return self.experts_forward(tokens, expert_ids, weights)
 
 
 class _AllToAllState(NamedTuple):
@@ -281,7 +278,7 @@ class AllToAllEP(EPImplementation):
     and rank 1 does not.
     """
 
-    def forward(self, hidden_states, top_k_index, top_k_weights):
+    def __call__(self, hidden_states, top_k_index, top_k_weights):
         tokens, expert_ids, weights, state = self._send(hidden_states, top_k_index, top_k_weights)
         return self._receive(self.run_experts(tokens, expert_ids, weights), state)
 
@@ -370,12 +367,14 @@ class AllReduceEP(EPImplementation):
 
     needs_shared_batch = True
 
-    def forward(self, hidden_states, top_k_index, top_k_weights):
+    def __call__(self, hidden_states, top_k_index, top_k_weights):
         ep_group = self.ep_mesh.get_group()
         # Each rank backpropagates through its own experts only, so the gradients of the shared inputs it computes are
         # partial: sum them over the group (the forward all-reduce has an identity backward).
-        hidden_states = _all_reduce_backward(hidden_states, ep_group)
-        top_k_weights = _all_reduce_backward(top_k_weights, ep_group)
+        if torch.is_grad_enabled() and hidden_states.requires_grad:
+            hidden_states = _AllReduceBackward.apply(hidden_states, ep_group)
+        if torch.is_grad_enabled() and top_k_weights.requires_grad:
+            top_k_weights = _AllReduceBackward.apply(top_k_weights, ep_group)
         off_rank = (top_k_index // self.num_local_experts) != self.ep_mesh.get_local_rank()
         local_index = torch.where(off_rank, self.num_local_experts, top_k_index % self.num_local_experts)
         expert_output = self.run_experts(hidden_states, local_index, top_k_weights.masked_fill(off_rank, 0.0))
@@ -395,10 +394,10 @@ def __init__(self, experts, ep_mesh, tp_mesh=None):
         # Only built where it can run: an EP group whose ranks hold different batches has no all-reduce path.
         self.all_reduce = AllReduceEP(experts, ep_mesh, tp_mesh) if self.shares_batch else None
 
-    def forward(self, hidden_states, top_k_index, top_k_weights):
+    def __call__(self, hidden_states, top_k_index, top_k_weights):
         if self.all_reduce is not None and not torch.is_grad_enabled():
-            return self.all_reduce.forward(hidden_states, top_k_index, top_k_weights)
-        return self.all_to_all.forward(hidden_states, top_k_index, top_k_weights)
+            return self.all_reduce(hidden_states, top_k_index, top_k_weights)
+        return self.all_to_all(hidden_states, top_k_index, top_k_weights)
 
 
 class EPImplementationInterface(GeneralInterface):
@@ -433,7 +432,7 @@ def ep_forward(hidden_states, top_k_index, top_k_weights):
                 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)
+            return implementation(hidden_states, top_k_index, top_k_weights).to(hidden_states.dtype)
 
         module.forward = ep_forward
         return module

src/transformers/distributed/parallel_utils.py

@@ -122,6 +122,9 @@ class _AllReduceBackward(torch.autograd.Function):
         ranks), but the backward gradient is partial (each rank has 1/tp_size from
         its expert shard). We need to sum the partial gradients without dividing by
         world_size, which is what DTensor's Replicate backward does incorrectly.
+
+        Apply it only when a gradient can flow (`torch.is_grad_enabled() and x.requires_grad`): at inference the
+        per-call overhead of an `autograd.Function` compounds across layers, e.g. in eager decode.
         """
 
         @staticmethod
@@ -135,14 +138,6 @@ def backward(ctx, grad):
             dist.all_reduce(grad, group=ctx.process_group)
             return grad, None
 
-    def _all_reduce_backward(x, process_group):
-        """Sum the gradient of a replicated `x` over the group in backward. Only when a gradient can flow: at
-        inference it returns `x` untouched, so the per-call overhead of an `autograd.Function` does not compound
-        across layers, e.g. in eager decode."""
-        if torch.is_grad_enabled() and x.requires_grad:
-            return _AllReduceBackward.apply(x, process_group)
-        return x
-
     class _WaitTensor(torch.autograd.Function):
         """Wait on a functional collective's result, identity backward (`wait_tensor` has no autograd kernel before
         torch 2.11)."""

src/transformers/distributed/tensor_parallel.py

@@ -47,7 +47,7 @@
     from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor
     from torch.distributed.tensor.placement_types import _StridedShard
 
-    from .parallel_utils import _all_reduce_backward, _AllReduceForward
+    from .parallel_utils import _AllReduceBackward, _AllReduceForward
 
 
 def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None):
@@ -184,7 +184,8 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
         # input-gradient contributions explicitly.
         if is_local_input and getattr(module, "_hf_quantized_needs_local_tp", False):
             process_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp")
-            x = _all_reduce_backward(x, process_group)
+            if torch.is_grad_enabled() and x.requires_grad:
+                x = _AllReduceBackward.apply(x, process_group)
             return (x,) + args[1:], kwargs
 
         # DTensor path: handle regular training and layout redistribution.
@@ -438,7 +439,8 @@ def transform_output_post_forward(self, module, output, mesh):
             )
         rope_dim = module.config.qk_rope_head_dim
         pass_output, rope_output = output.split([output.shape[-1] - rope_dim, rope_dim], dim=-1)
-        rope_output = _all_reduce_backward(rope_output, mesh.get_group())
+        if torch.is_grad_enabled() and rope_output.requires_grad:
+            rope_output = _AllReduceBackward.apply(rope_output, mesh.get_group())
         return torch.cat([pass_output, rope_output], dim=-1)
 
 
@@ -558,14 +560,16 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_
         tp_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp")
         if isinstance(hidden_states, DTensor):
             hidden_states = hidden_states.to_local()
-        hidden_states = _all_reduce_backward(hidden_states, tp_group)
+        if torch.is_grad_enabled() and hidden_states.requires_grad:
+            hidden_states = _AllReduceBackward.apply(hidden_states, tp_group)
 
         if len(routing_args) >= 2:
             top_k_index, top_k_weights, *extra_args = routing_args
             if isinstance(top_k_weights, DTensor):
                 top_k_weights = top_k_weights.to_local()
             if not is_expert_parallel:
-                top_k_weights = _all_reduce_backward(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)
             routing_args = [top_k_index, top_k_weights, *extra_args]
 
         return (hidden_states, *routing_args), kwargs

tests/expert_parallel/test_expert_parallel.py

@@ -105,7 +105,7 @@ class Custom(EPImplementation):
             def __init__(self, experts, ep_mesh, tp_mesh=None):
                 self.experts = experts
 
-            def forward(self, hidden_states, top_k_index, top_k_weights):
+            def __call__(self, hidden_states, top_k_index, top_k_weights):
                 return hidden_states + 1
 
         with patch.dict(EPImplementationInterface._global_mapping, {"custom": Custom}):
@@ -175,15 +175,15 @@ def _ep_implementations_match_one_process(rank):
         raise AssertionError("AllReduceEP accepted an EP group without a shared batch")
 
     exchanged = []
-    all_to_all_forward = AllToAllEP.forward
+    all_to_all_call = AllToAllEP.__call__
     for ep_implementation, training, exchanges in (
         ("auto", True, True),
         ("auto", False, False),
         ("all_reduce", True, False),
         ("all_to_all", False, True),
     ):
         exchanged.clear()
-        with patch.object(AllToAllEP, "forward", lambda *args: exchanged.append(True) or all_to_all_forward(*args)):
+        with patch.object(AllToAllEP, "__call__", lambda *args: exchanged.append(True) or all_to_all_call(*args)):
             module = _ep_experts(rank, ep_implementation)
             h, w = _HIDDEN_STATES.clone().requires_grad_(), _TOP_K_WEIGHTS.clone().requires_grad_()
             with torch.set_grad_enabled(training):
@@ -204,19 +204,19 @@ def _ep_gradient_checkpointing_matches_reference(rank, use_reentrant):
     reference_output.backward(_GRAD_OUTPUT)
 
     all_to_all_calls, all_reduce_calls = [], []
-    all_to_all_forward, all_reduce_forward = AllToAllEP.forward, AllReduceEP.forward
+    all_to_all_call, all_reduce_call = AllToAllEP.__call__, AllReduceEP.__call__
 
     def record_all_to_all(self, *args):
         all_to_all_calls.append(True)
-        return all_to_all_forward(self, *args)
+        return all_to_all_call(self, *args)
 
     def record_all_reduce(self, *args):
         all_reduce_calls.append(True)
-        return all_reduce_forward(self, *args)
+        return all_reduce_call(self, *args)
 
     with (
-        patch.object(AllToAllEP, "forward", record_all_to_all),
-        patch.object(AllReduceEP, "forward", record_all_reduce),
+        patch.object(AllToAllEP, "__call__", record_all_to_all),
+        patch.object(AllReduceEP, "__call__", record_all_reduce),
     ):
         module = _ep_experts(rank, "auto")
         h, w = _HIDDEN_STATES.clone().requires_grad_(), _TOP_K_WEIGHTS.clone().requires_grad_()