[distributed] Allow context parallelism together with expert parallelism

Under cp the ranks of an ep group hold the slices of several batches, which token dispatch already handles: each rank sends the tokens of its own slice. The expert gradient divides by dp_size, which includes cp_size, so the reduced expert gradient is still the gradient of the full-sequence loss.

Author: 3outeille

src/transformers/distributed/configuration_utils.py

@@ -63,8 +63,9 @@ class DistributedConfig:
             Number of devices that split each sequence (context parallelism). The ranks of a cp group get the same
             batch, and each keeps a contiguous slice of every sequence, so the sequence length must be divisible by
             `cp_size`. With `cp_size > 1`, FSDP2 shards over `dp_size` ranks: the ranks of a cp group hold parts of
-            the same batch, so their gradients are partial sums of the same batch. Can be combined with `tp_size`, in
-            which case the ranks of a tp group hold the same batch and the same slice. Defaults to 1.
+            the same batch, so their gradients are partial sums of the same batch. Can be combined with `tp_size` and
+            `ep_size`, in which case the ranks of a tp group hold the same batch and the same slice, while each rank
+            of an ep group dispatches the tokens of its own slice. Defaults to 1.
         cp_implementation (`str`, *optional*, defaults to `"ulysses"`):
             How attention sees the whole sequence while each rank holds a slice of it. A name in
             `transformers.distributed.context_parallel.cp_implementations`. `"ulysses"`: an all-to-all gives each
@@ -194,12 +195,11 @@ 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`).")
-            unsupported = [name for name in ("pp_size", "ep_size") if getattr(self, name) > 1]
-            if unsupported:
+            if self.pp_size > 1:
                 raise ValueError(
                     f"Context parallelism (`cp_size={self.cp_size}`) is not supported yet with "
-                    f"{', '.join(f'`{name}={getattr(self, name)}`' for name in unsupported)}. Use "
-                    "`DistributedConfig(tp_size=N, cp_size=M)` or `DistributedConfig(fsdp_size=N, cp_size=M)`."
+                    f"`pp_size={self.pp_size}`. Use `DistributedConfig(tp_size=N, cp_size=M)`, "
+                    "`DistributedConfig(fsdp_size=N, cp_size=M)` or `DistributedConfig(ep_size=N, cp_size=M)`."
                 )
 
         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

src/transformers/distributed/utils.py

@@ -205,7 +205,8 @@ class TransformersDeviceMesh:
     parameters over `fsdp` and `cp` together (`fsdp_cp`): the ranks of a cp group hold parts of the same batch, so their
     gradients are partial sums of that batch. `fsdp_cp` is an axis of the dense root mesh because FSDP2 composes the
     mesh it shards over with the mesh the parameters are already sharded on (`tp`), and torch only composes meshes
-    built from the same root.
+    built from the same root. On the expert mesh a cp slice plays the role of a batch: each rank dispatches the tokens
+    of its own slice only, so no token is sent twice and `efsdp * ep` still covers `fsdp * cp * tp`.
     """
 
     def __init__(self, dense_mesh: DeviceMesh, expert_mesh: DeviceMesh, dense_root_mesh: DeviceMesh | None = None):