[distributed] Test context parallelism with expert parallelism

A tiny Qwen3 MoE against the same weights on full sequences: loss, this rank's logits slice and every gradient, with token dispatch (cp2 x ep2, cp2 x fsdp2 x ep4, cp2 x fsdp2 x ep2, cp2 x tp2 x ep4, cp2 x hsdp2 x ep2) and with the all-reduce plan (tp2 x cp2 x ep2).

Author: 3outeille

tests/context_parallel/test_context_parallel.py

@@ -28,6 +28,7 @@
     LlamaForSequenceClassification,
     LlamaModel,
     MistralConfig,
+    Qwen3MoeConfig,
 )
 from transformers.distributed import DistributedConfig, context_parallel
 from transformers.distributed.context_parallel import (
@@ -446,6 +447,9 @@ def test_packed_flash_attention(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):
@@ -531,6 +535,7 @@ def _compare_with_unsharded(init_dir, batch, distributed_config, model_class=Aut
             continue
         grad = _average_over_data_parallel(grad, dp_group, fsdp_size * dp_replicate_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(rank, cp_size, fsdp_size, dp_replicate_size, packed):
@@ -565,6 +570,64 @@ def _test_matches_unsharded_with_tp(rank, tp_size, cp_size, fsdp_size, packed):
         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()
@@ -726,17 +789,26 @@ def setUp(self):
         if not is_torch_greater_or_equal("2.10"):
             self.skipTest("Context parallelism requires torch>=2.10")
 
-    def test_refuses_other_parallelisms(self):
-        for other in ({"pp_size": 2}, {"ep_size": 2}):
-            with self.subTest(**other), self.assertRaisesRegex(ValueError, "not supported yet with"):
-                DistributedConfig(cp_size=2, **other)
+    def test_refuses_pipeline_parallelism(self):
+        with self.assertRaisesRegex(ValueError, "not supported yet with `pp_size=2`"):
+            DistributedConfig(cp_size=2, pp_size=2)
 
     def test_accepts_tensor_parallelism(self):
         distributed_config = DistributedConfig(tp_size=2, cp_size=2)
         self.assertEqual((distributed_config.tp_size, distributed_config.cp_size), (2, 2))
         # the tp ranks hold the same batch: they are not data parallel ranks
         self.assertEqual(distributed_config.dp_size, 2)
 
+    def test_accepts_expert_parallelism(self):
+        distributed_config = DistributedConfig(cp_size=2, ep_size=2)
+        self.assertEqual((distributed_config.cp_size, distributed_config.ep_size), (2, 2))
+        # the cp ranks are data parallel ranks, and the experts shard over the axis ep leaves of fsdp * cp * tp
+        self.assertEqual(distributed_config.dp_size, 2)
+        self.assertEqual(distributed_config.efsdp_size, 1)
+        self.assertEqual(DistributedConfig(cp_size=2, fsdp_size=2, ep_size=2).efsdp_size, 2)
+        with self.assertRaisesRegex(ValueError, r"must divide `fsdp_size \* cp_size \* tp_size`"):
+            DistributedConfig(cp_size=2, ep_size=4)
+
     def test_refuses_size_below_one(self):
         with self.assertRaisesRegex(ValueError, "must be >= 1"):
             DistributedConfig(cp_size=0)
@@ -770,6 +842,29 @@ def test_matches_unsharded_with_tensor_parallelism(self):
     def test_matches_unsharded_with_tensor_parallelism_and_fsdp(self):
         run_distributed(world_size=8, backend="gloo")(_test_matches_unsharded_with_tp)(2, 2, 2, False)
 
+    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)