[distributed] Context parallelism with expert parallelism

src/transformers/distributed/configuration_utils.py

@@ -193,10 +193,10 @@ def _validate_mesh_config(self):
         if self.cp_size > 1:
             if not is_torch_greater_or_equal("2.10"):
                 raise OSError("Context parallelism requires `torch>=2.10` (`DeviceMesh._unflatten`).")
-            if unsupported := [name for name in ("pp_size", "ep_size") if getattr(self, name) > 1]:
+            if self.pp_size > 1:
                 raise ValueError(
-                    f"Context parallelism (`cp_size={self.cp_size}`) is not supported yet with `{unsupported[0]}`. "
-                    "Use `DistributedConfig(fsdp_size=N, tp_size=T, cp_size=M)`."
+                    f"Context parallelism (`cp_size={self.cp_size}`) is not supported yet with pipeline parallelism "
+                    f"(`pp_size={self.pp_size}`)."
                 )
 
         if self.dp_size > 1 and self.pp_size > 1:

src/transformers/distributed/fsdp.py

@@ -202,10 +202,10 @@ def apply_fully_sharded_data_parallelism(
         for module_name, module in model.named_modules():
             if ep_plan.get(module_name) == "ep_dispatch_experts":
                 fully_shard(module, mesh=expert_mesh, reshard_after_forward=True, **fsdp_policy_kwargs)
-                # an EP group holds ep_size / tp_size distinct batches
+                # an EP group holds ep_size / tp_size distinct batches, or cp slices of them under context parallelism
                 # the efsdp reduce the sums on efsdp_size copies of that expert.
-                # The expert gradient therefore covers ep_size / tp_size * efsdp_size = fsdp_size batches
-                # not efsdp_size batches (the relationship is efsdp_size = fsdp_size * tp_size / ep_size)
+                # The expert gradient therefore covers ep_size / tp_size * efsdp_size = fsdp_size * cp_size batches
+                # not efsdp_size batches (the relationship is efsdp_size = fsdp_size * cp_size * tp_size / ep_size)
                 # With HSDP, the all-reduce across replicas adds dp_replicate_size times more: dp_size batches.
                 module.set_gradient_divide_factor(float(distributed_config.dp_size))
                 # FSDP2 applies a divide factor other than the group size through PREMUL_SUM, which non-NCCL backends

tests/context_parallel/test_context_parallel.py

@@ -33,6 +33,7 @@
     LlamaForSequenceClassification,
     LlamaModel,
     MistralConfig,
+    Qwen3MoeConfig,
 )
 from transformers.distributed import DistributedConfig
 from transformers.distributed.context_parallel import (
@@ -576,6 +577,9 @@ def test_a_registered_attention_style_is_a_context_parallel_style(self):
 
 CP_RTOL = 1e-5
 CP_ATOL = 1e-5
+# The all-reduce plan of `Qwen3Moe`, which overrides the token-dispatch rules of its own `ep_plan`: the router masks the
+# experts of the other ranks and the partial expert outputs are all-reduced.
+ALL_REDUCE_EP_PLAN = {"model.layers.*.mlp.gate": "ep_router", "model.layers.*.mlp.experts": "moe_tp_experts"}
 
 
 def _tiny_batch(config, batch_size, seq_len, packed=False, seed=0):
@@ -675,6 +679,7 @@ def _compare_with_unsharded(init_dir, batch, distributed_config, keep=None):
             continue
         grad = _average_over_group(grad, groups, group_size)
         torch.testing.assert_close(grad, expected, rtol=CP_RTOL, atol=CP_ATOL, msg=f"Gradient mismatch for {name}")
+    return model
 
 
 def _test_matches_unsharded_with_tp(rank, tp_size, cp_size, fsdp_size, packed):
@@ -754,6 +759,64 @@ def _test_padded(rank, cp_size, config=None, packed=False):
         dist.barrier()
 
 
+def _tiny_moe_config(**kwargs):
+    """A `Qwen3MoeConfig` small enough for CPU tests, whose default `ep_plan` dispatches the tokens to the experts."""
+    return Qwen3MoeConfig(
+        **{
+            "hidden_size": 32,
+            "moe_intermediate_size": 16,
+            "num_hidden_layers": 2,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 4,
+            "vocab_size": 99,
+            "num_experts": 8,
+            "num_experts_per_tok": 2,
+            "attn_implementation": "sdpa",
+            **kwargs,
+        }
+    )
+
+
+def _test_matches_unsharded_with_ep(rank, cp_size, fsdp_size, ep_size, tp_size=1, dp_replicate_size=1):
+    """
+    Token dispatch with context parallelism: every rank sends the tokens of its own slice to the owners of the experts
+    they selected, and the experts' gradients reduce over the ranks that hold a copy of them.
+    """
+    config = _tiny_moe_config()
+    batch_size, seq_len = 2, 8
+    with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
+        rows = batch_size * fsdp_size * dp_replicate_size
+        input_ids, _ = _tiny_batch(config, rows, seq_len)
+        distributed_config = DistributedConfig(
+            cp_size=cp_size,
+            fsdp_size=fsdp_size,
+            ep_size=ep_size,
+            tp_size=tp_size,
+            dp_replicate_size=dp_replicate_size,
+        )
+        batch = {"input_ids": input_ids, "labels": input_ids.clone(), "use_cache": False}
+        model = _compare_with_unsharded(init_dir, batch, distributed_config)
+        assert model.model.layers[0].mlp.experts.num_experts == config.num_experts // ep_size
+        dist.barrier()
+
+
+def _test_matches_unsharded_with_all_reduce_ep(rank, cp_size, ep_size, tp_size):
+    """
+    All-reduce experts need every rank of an ep group to see the same tokens: with `ep_size == tp_size` that group is
+    the tp group, whose ranks hold the same cp slice.
+    """
+    config = _tiny_moe_config()
+    batch_size, seq_len = 2, 8
+    with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
+        input_ids, _ = _tiny_batch(config, batch_size, seq_len)
+        distributed_config = DistributedConfig(
+            cp_size=cp_size, ep_size=ep_size, tp_size=tp_size, ep_plan=ALL_REDUCE_EP_PLAN
+        )
+        batch = {"input_ids": input_ids, "labels": input_ids.clone(), "use_cache": False}
+        _compare_with_unsharded(init_dir, batch, distributed_config)
+        dist.barrier()
+
+
 def _test_inputs_embeds(rank, cp_size):
     """`inputs_embeds` is sharded along its sequence dim, not along its last, hidden one."""
     config = _tiny_config()
@@ -847,7 +910,7 @@ def _test_save_load_distributed_checkpoint(rank, cp_size):
 @unittest.skipUnless(is_torch_greater_or_equal("2.10"), "Context parallelism requires torch>=2.10")
 class TestContextParallelConfig(unittest.TestCase):
     def test_refuses_pipeline_parallelism(self):
-        with self.assertRaisesRegex(ValueError, "not supported yet with `pp_size`"):
+        with self.assertRaisesRegex(ValueError, "not supported yet with pipeline parallelism"):
             DistributedConfig(cp_size=2, pp_size=2)
 
     def test_refuses_a_model_without_a_cp_plan(self):
@@ -945,6 +1008,29 @@ def test_matches_unsharded_with_tensor_parallelism(self):
     def test_tensor_parallelism_and_context_parallelism_compose(self):
         run_distributed(world_size=4, backend="gloo")(_test_tp_and_cp_compose)(2, 2)
 
+    def test_matches_unsharded_with_expert_parallelism(self):
+        run_distributed(world_size=2, backend="gloo")(_test_matches_unsharded_with_ep)(2, 1, 2)
+
+    def test_matches_unsharded_with_fsdp_and_expert_parallelism(self):
+        # ep_size = fsdp_size * cp_size, so efsdp_size is 1 and every expert lives whole on a single rank
+        run_distributed(world_size=4, backend="gloo")(_test_matches_unsharded_with_ep)(2, 2, 4)
+
+    def test_matches_unsharded_with_expert_fsdp(self):
+        # efsdp_size = fsdp_size * cp_size * tp_size / ep_size = 2: pairs of ranks hold the same experts, sharded
+        # between them
+        run_distributed(world_size=4, backend="gloo")(_test_matches_unsharded_with_ep)(2, 2, 2)
+
+    def test_matches_unsharded_with_tensor_parallelism_and_expert_parallelism(self):
+        # ep_size = 2 * tp_size puts both cp slices of the batch, two ranks each, in the same ep group
+        run_distributed(world_size=4, backend="gloo")(_test_matches_unsharded_with_ep)(2, 1, 4, 2)
+
+    def test_matches_unsharded_with_hsdp_and_expert_parallelism(self):
+        # the expert gradient reduces over efsdp and then over the two replicas: dp_size = 2 * 1 * 2 ranks
+        run_distributed(world_size=4, backend="gloo")(_test_matches_unsharded_with_ep)(2, 1, 2, 1, 2)
+
+    def test_matches_unsharded_with_all_reduce_expert_parallelism(self):
+        run_distributed(world_size=4, backend="gloo")(_test_matches_unsharded_with_all_reduce_ep)(2, 2, 2)
+
     def test_inputs_embeds(self):
         run_distributed(world_size=2, backend="gloo")(_test_inputs_embeds)(2)