[distributed] Context parallelism with expert parallelism

Each rank of a cp group dispatches the tokens of its own slice only, so no token reaches an expert
twice and `efsdp * ep` still covers `fsdp * cp * tp`. The expert gradient divide factor already
counts the cp ranks through `dp_size`, so only the explanatory comment and the mesh docstring
change; `cp_size` with `ep_size` is no longer refused.

Author: 3outeille

src/transformers/distributed/configuration_utils.py

@@ -196,10 +196,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