[distributed] DistributedConfig(cp_size=N)

src/transformers/configuration_utils.py

@@ -320,6 +320,9 @@ class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin, Heterogeneous
     - **base_model_fsdp_plan** (`dict[Any, str]`) -- A dict that maps sub-modules of a base model to an FSDP2
       sharding strategy (e.g. `"free_full_weight"` / `"keep_full_weight"`). Keys can be wildcard module paths
       (e.g. `"layers.*"`) or tuples of paths (grouped into a single `fully_shard` call).
+    - **base_model_cp_plan** (`dict[str, str]`) -- A dict that maps sub-modules of a base model to a context parallel
+      style, e.g. `"layers.*": "cp_split_sequence"` to keep this rank's slice of the sequence in the decoder layers
+      and `"layers.*.self_attn": "cp_ulysses"` to give their attention the whole sequence with Ulysses.
     - **base_model_pp_plan** (`dict[str, tuple[list[str]]]`) -- A dict that maps child-modules of a base model to a
       pipeline parallel plan that enables users to place the child-module on the appropriate device.
 
@@ -400,6 +403,10 @@ class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin, Heterogeneous
         "layers.*": "free_full_weight",
         "norm": "keep_full_weight",
     }
+    base_model_cp_plan: ClassVar[dict[str, str]] = {
+        "layers.*": "cp_split_sequence",
+        "layers.*.self_attn": "cp_ulysses",
+    }
     base_model_pp_plan: ClassVar[dict[str, Sequence[list[str]]] | None] = None
     base_model_ep_plan: ClassVar[dict[str, Sequence[list[str]]] | None] = None
     _auto_class: ClassVar[str | None] = None
@@ -1426,6 +1433,7 @@ def _remove_keys_not_serialized(self, d: dict[str, Any]) -> None:
             "base_model_tp_plan",
             "base_model_pp_plan",
             "base_model_fsdp_plan",
+            "base_model_cp_plan",
             "distributed_config",
         ]:
             d.pop(key_to_remove, None)

src/transformers/distributed/configuration_utils.py

@@ -25,7 +25,7 @@
 @dataclass
 class DistributedConfig:
     """
-    Configuration for native distributed inference and training with tensor, pipeline, or FSDP2 parallelism.
+    Configuration for native distributed inference and training with tensor, pipeline, context, or FSDP2 parallelism.
 
     Args:
         tp_size (`int`, *optional*):
@@ -59,6 +59,14 @@ class DistributedConfig:
             Whether to enable CPU offloading for FSDP2.
         fsdp_mixed_precision (`bool`, *optional*, defaults to `False`):
             Whether to enable mixed precision for FSDP2.
+        cp_size (`int`, *optional*):
+            Number of devices that split each sequence (context parallelism). The ranks of a cp group get the same
+            batch, and each runs on a contiguous slice of every sequence, padded at its end to a multiple of `cp_size`.
+            Requires `torch>=2.10`.
+        cp_plan (`dict[str, str]`, *optional*):
+            Context parallel plan. Leave as `None` to use the model's predefined `base_model_cp_plan` and class-level
+            `_cp_plan`. Pass a dictionary to override individual rules; unspecified rules are kept, so
+            `{"model.layers.*.self_attn": "cp_ring"}` only changes the attention style. Applied only when `cp_size > 1`.
         pp_size (`int`, *optional*):
             Number of devices for pipeline parallelism. If `None` and another parallel mode is set, defaults to 1.
         ep_size (`int`, *optional*):
@@ -81,6 +89,8 @@ class DistributedConfig:
     fsdp_plan: dict[str, str] | None = None
     fsdp_cpu_offload: bool = False
     fsdp_mixed_precision: bool = False
+    cp_size: int | None = None
+    cp_plan: dict[str, str] | None = None
     pp_size: int | None = None
     ep_size: int | None = None
     ep_plan: dict[str, str] | None = None
@@ -91,7 +101,7 @@ def efsdp_size(self) -> int:
         The number of ranks that own the same experts and FSDP-shard them between each other.
         Dense and expert parameters are laid out over the same world size:
 
-            pp x fsdp x tp == pp x efsdp x ep   =>   efsdp = fsdp x tp // ep
+            pp x fsdp x cp x tp == pp x efsdp x ep   =>   efsdp = fsdp x cp x tp // ep
 
         Experts are not tensor-parallel, so EP and expert-FSDP together cover all the ranks that
         dense parameters split between FSDP and TP. Example with 16 ranks, ep=8:
@@ -101,41 +111,46 @@ def efsdp_size(self) -> int:
             efsdp groups : {0,8} {1,9} ... {7,15}  each pair holds the same experts, FSDP-sharded
                                                    on dim 0 (the expert dim) and all-gathered for compute
 
-        - ep == tp       : efsdp == fsdp
-        - ep == fsdp * tp: efsdp == 1, every expert lives whole on a single rank
+        - ep == tp            : efsdp == fsdp * cp
+        - ep == fsdp * cp * tp: efsdp == 1, every expert lives whole on a single rank
 
         Sharding over `efsdp` is applied by the EP token-dispatch path (`ep_dispatch_experts`).
         """
-        return self.fsdp_size * self.tp_size // self.ep_size
+        return self.fsdp_size * self.cp_size * self.tp_size // self.ep_size
 
     @property
     def dp_size(self) -> int:
-        """The number of ranks that each see a distinct batch: `dp_replicate_size * fsdp_size`."""
-        return self.dp_replicate_size * self.fsdp_size
+        """
+        The number of ranks FSDP2 reduces gradients over: `dp_replicate_size * fsdp_size * cp_size`. The ranks of a
+        cp group share one batch, but each processes different tokens of it.
+        """
+        return self.dp_replicate_size * self.fsdp_size * self.cp_size
 
     def __post_init__(self):
         self._resolve_parallelism()
         self._validate_mesh_config()
 
     def _resolve_parallelism(self):
         """Resolve parallel sizes and legacy EP settings."""
-        for value in (self.tp_size, self.fsdp_size, self.dp_replicate_size, self.pp_size, self.ep_size):
+        for value in (self.tp_size, self.fsdp_size, self.dp_replicate_size, self.cp_size, self.pp_size, self.ep_size):
             if value is not None and value < 1:
                 raise ValueError(f"Parallelism sizes must be >= 1, got {value}.")
 
         if self.fsdp_size is None:
             self.fsdp_size = 1
         if self.dp_replicate_size is None:
             self.dp_replicate_size = 1
+        if self.cp_size is None:
+            self.cp_size = 1
         if self.pp_size is None:
             self.pp_size = 1
         if self.tp_size is None and self.tp_plan is not None:
             world_size = int(os.environ.get("WORLD_SIZE", 1))
-            other_parallel_size = self.dp_replicate_size * self.fsdp_size * self.pp_size
+            other_parallel_size = self.dp_replicate_size * self.fsdp_size * self.cp_size * self.pp_size
             if world_size % other_parallel_size != 0:
                 raise ValueError(
-                    f"WORLD_SIZE ({world_size}) must be divisible by dp_replicate_size * fsdp_size * pp_size "
-                    f"({other_parallel_size}) to derive tp_size."
+                    f"WORLD_SIZE ({world_size}) must be divisible by dp_replicate_size * fsdp_size * cp_size * "
+                    f"pp_size ({other_parallel_size}) to derive tp_size."
                 )
             self.tp_size = world_size // other_parallel_size
         elif self.tp_size is None:
@@ -165,12 +180,23 @@ def _validate_mesh_config(self):
             raise ValueError("`ep_plan` must be a dictionary or None.")
         if self.fsdp_plan is not None and not isinstance(self.fsdp_plan, dict):
             raise ValueError("`fsdp_plan` must be a dictionary or None.")
+        if self.cp_plan is not None and not isinstance(self.cp_plan, dict):
+            raise ValueError("`cp_plan` must be a dictionary or None.")
 
         if self.ep_size > 1:
             if self.ep_size % self.tp_size:
                 raise ValueError("`ep_size` must be a multiple of `tp_size`.")
-            if (self.fsdp_size * self.tp_size) % self.ep_size:
-                raise ValueError("`ep_size` must divide `fsdp_size * tp_size`.")
+            if (self.fsdp_size * self.cp_size * self.tp_size) % self.ep_size:
+                raise ValueError("`ep_size` must divide `fsdp_size * cp_size * tp_size`.")
+
+        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", "tp_size", "ep_size") if getattr(self, name) > 1]:
+                raise ValueError(
+                    f"Context parallelism (`cp_size={self.cp_size}`) is not supported yet with `{unsupported[0]}`. "
+                    "Use `DistributedConfig(fsdp_size=N, cp_size=M)`."
+                )
 
         if self.dp_size > 1 and self.pp_size > 1:
             raise ValueError(

src/transformers/distributed/fsdp.py

@@ -180,16 +180,18 @@ def apply_fully_sharded_data_parallelism(
     ep_plan: dict[str, str],
 ) -> nn.Module:
     """
-    Apply FSDP2 (fully_shard) to a model: dispatched experts on `efsdp` and the rest of the modules on `fsdp`.
+    Apply FSDP2 (fully_shard) to a model: dispatched experts on `efsdp` and the rest of the modules on `fsdp`, which
+    spans `cp` too when context parallelism splits the sequences.
 
     Torch availability, distributed initialization and the version requirement
     are asserted upstream by `initialize_distributed_mesh`.
     """
     distributed_config = getattr(model.config, "distributed_config", None)
+    fsdp_dim = "fsdp_cp" if distributed_config is not None and distributed_config.cp_size > 1 else "fsdp"
     if distributed_config is not None and distributed_config.dp_replicate_size > 1:
-        fsdp_mesh = mesh_manager.get_mesh(("dp_replicate", "fsdp"))
+        fsdp_mesh = mesh_manager.get_mesh(("dp_replicate", fsdp_dim))
     else:
-        fsdp_mesh = mesh_manager.get_mesh("fsdp")
+        fsdp_mesh = mesh_manager.get_mesh(fsdp_dim)
     fsdp_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config)
 
     if distributed_config is not None and distributed_config.ep_size > 1:

src/transformers/distributed/mixin.py

@@ -32,6 +32,12 @@
     consolidate_distributed_checkpoint,
 )
 from .configuration_utils import DistributedConfig
+from .context_parallel import (
+    ALL_CP_STYLES,
+    ContextParallelAttention,
+    SplitSequenceParallel,
+    apply_context_parallelism,
+)
 from .expert_parallel import apply_expert_parallelism
 from .fsdp import ALL_FSDP_STYLES, _resolve_tied_embed_lm_head_plan, apply_fully_sharded_data_parallelism
 from .parallel_utils import ALL_PARALLEL_STYLES, _get_parameter_plan, replace_layer_number_by_wildcard
@@ -67,6 +73,7 @@ class DistributedMixin:
     _fsdp_size = None
     _pp_plan: dict[str, tuple[str, str]] | None = None
     _fsdp_plan: dict[str, str] | None = None
+    _cp_plan: dict[str, str] | None = None
 
     @property
     def is_distributed_loading_by_transformers(self) -> bool:
@@ -83,12 +90,14 @@ def init_parallel_plans(self) -> None:
         self._ep_plan = dict(getattr(model_cls, "_ep_plan", None) or {})
         self._pp_plan = dict(getattr(model_cls, "_pp_plan", None) or {})
         self._fsdp_plan = dict(getattr(model_cls, "_fsdp_plan", None) or {})
+        self._cp_plan = dict(getattr(model_cls, "_cp_plan", None) or {})
 
         if self.base_model is self:
             self._pp_plan.update(self.config.base_model_pp_plan or {})
             self._tp_plan.update(self.config.base_model_tp_plan or {})
             self._ep_plan.update(self.config.base_model_ep_plan or {})
             self._fsdp_plan.update(self.config.base_model_fsdp_plan or {})
+            self._cp_plan.update(self.config.base_model_cp_plan or {})
             # Shard the input embeddings on the vocabulary unless the plan already names them.
             if self._tp_plan and isinstance(getattr(self, "embed_tokens", None), torch.nn.Embedding):
                 self._tp_plan = {"embed_tokens": "embedding_rowwise", **self._tp_plan}
@@ -102,6 +111,8 @@ def init_parallel_plans(self) -> None:
                 self._pp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()})
             if plan := getattr(module, "_fsdp_plan", None):
                 self._fsdp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()})
+            if plan := getattr(module, "_cp_plan", None):
+                self._cp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()})
 
     @property
     def tp_plan(self) -> dict[str, str]:
@@ -117,6 +128,11 @@ def ep_plan(self) -> dict[str, str]:
     def fsdp_plan(self) -> dict[str, str]:
         return self._fsdp_plan
 
+    @property
+    def cp_plan(self) -> dict[str, str]:
+        """The full context parallel plan for the model's modules."""
+        return self._cp_plan
+
     @property
     def pp_plan(self) -> dict[str, tuple[str, str]]:
         return self._pp_plan
@@ -154,6 +170,17 @@ def fsdp_plan(self, plan: dict[str, str] | None):
         self._validate_parallel_plan(plan, ALL_FSDP_STYLES, "fsdp_plan")
         self._fsdp_plan = plan
 
+    @cp_plan.setter
+    def cp_plan(self, plan: dict[str, str] | None):
+        if plan is None:
+            self._cp_plan = {}
+            return
+        if not isinstance(plan, dict):
+            raise ValueError("Can only set a dictionary as `cp_plan`")
+
+        self._validate_parallel_plan(plan, ALL_CP_STYLES, "cp_plan")
+        self._cp_plan = plan
+
     @pp_plan.setter
     def pp_plan(self, plan: dict[str, tuple[str, str]] | None):
         if plan is None:
@@ -215,8 +242,8 @@ def _expand_parallel_plan(self, plan: dict[str, str]) -> dict[str, str]:
         for module_name, module in self.named_modules():
             module_style = _get_parameter_plan(parameter_name=module_name, plan=plan, is_weight=False)
 
-            if module_style in ALL_FSDP_STYLES:
-                # FSDP: list the module to wrap.
+            if module_style in ALL_FSDP_STYLES or module_style in ALL_CP_STYLES:
+                # FSDP2 and context parallelism wrap the module itself.
                 expanded_plan[module_name] = module_style
             else:
                 # TP/EP: list the module's own parameters (to shard) before the module (to hook).
@@ -232,16 +259,16 @@ def _expand_parallel_plan(self, plan: dict[str, str]) -> dict[str, str]:
 
     def _resolve_parallel_plans(
         self, distributed_config: DistributedConfig
-    ) -> tuple[dict[str, str], dict[str, str], dict[str, str]]:
-        """Build the final TP, EP and FSDP plans: apply the config overrides, then replace `*` with the layer numbers.
+    ) -> tuple[dict[str, str], dict[str, str], dict[str, str], dict[str, str]]:
+        """Build the final TP, EP, FSDP and CP plans: apply the config overrides, then replace `*` by the layer numbers.
 
         Each returned plan maps one concrete module or parameter name to a style, e.g. `model.layers.3.mlp.up_proj`
         instead of `model.layers.*.mlp.up_proj`. A module that has an EP rule keeps no TP rule.
         """
         module_names = {replace_layer_number_by_wildcard(n) for n, _ in self.named_modules()}
         param_names = {replace_layer_number_by_wildcard(n) for n, _ in self.named_parameters()}
         names = module_names | param_names
-        for plan_name in ("tp_plan", "ep_plan", "fsdp_plan"):
+        for plan_name in ("tp_plan", "ep_plan", "fsdp_plan", "cp_plan"):
             override = getattr(distributed_config, plan_name)
             if isinstance(override, dict):
                 plan = getattr(self, plan_name)
@@ -253,8 +280,9 @@ def _resolve_parallel_plans(
         ep_plan = dict(self.ep_plan) if distributed_config.ep_size > 1 else {}
         # HSDP shards each replica with the FSDP plan.
         fsdp_plan = dict(self.fsdp_plan) if distributed_config.dp_size > 1 else {}
+        cp_plan = dict(self.cp_plan) if distributed_config.cp_size > 1 else {}
 
-        for kind, plan in (("tp", tp_plan), ("ep", ep_plan), ("fsdp", fsdp_plan)):
+        for kind, plan in (("tp", tp_plan), ("ep", ep_plan), ("fsdp", fsdp_plan), ("cp", cp_plan)):
             size = distributed_config.dp_size if kind == "fsdp" else getattr(distributed_config, f"{kind}_size")
             if size > 1 and not plan:
                 raise ValueError(
@@ -275,6 +303,7 @@ def _resolve_parallel_plans(
         self._validate_parallel_plan(tp_plan, ALL_PARALLEL_STYLES, "tp_plan", names)
         self._validate_parallel_plan(ep_plan, ALL_PARALLEL_STYLES, "ep_plan", names)
         self._validate_parallel_plan(fsdp_plan, ALL_FSDP_STYLES, "fsdp_plan", names)
+        self._validate_parallel_plan(cp_plan, ALL_CP_STYLES, "cp_plan", names)
         distributed_config._validate_resolved_ep_plan(ep_plan)
 
         # Tied embed/lm_head weights are wrapped once: fold the head's FSDP rule into the embedding's.
@@ -289,10 +318,24 @@ def _resolve_parallel_plans(
         tp_plan = self._expand_parallel_plan(tp_plan)
         ep_plan = self._expand_parallel_plan(ep_plan)
         fsdp_plan = self._expand_parallel_plan(fsdp_plan)
+        cp_plan = self._expand_parallel_plan(cp_plan)
         # FSDP2 needs a submodule sharded before the module containing it, e.g. `mlp.gate` before its layer.
         fsdp_plan = dict(sorted(fsdp_plan.items(), key=lambda item: item[0].count("."), reverse=True))
 
-        return tp_plan, ep_plan, fsdp_plan
+        if distributed_config.cp_size > 1:
+            styles = [ALL_CP_STYLES[style] for style in cp_plan.values()]
+            # a split style on the decoder layers, and an attention style (Ulysses, or another registered one)
+            for name, kind in (("cp_split_sequence", SplitSequenceParallel), ("cp_ulysses", ContextParallelAttention)):
+                if not any(isinstance(style, kind) for style in styles):
+                    raise ValueError(
+                        f"`cp_plan` has no `{name}` rule matching a module of {type(self).__name__}. Context "
+                        "parallelism needs `cp_split_sequence` on the decoder layers, to keep this rank's slice of the "
+                        "sequence, and an attention style such as `cp_ulysses` on their attention modules, to give "
+                        "attention the whole sequence. Add them to `base_model_cp_plan` of the model's config, pass "
+                        "`cp_plan`, or set `cp_size=1`."
+                    )
+
+        return tp_plan, ep_plan, fsdp_plan, cp_plan
 
     @classmethod
     def maybe_distribute_model(
@@ -301,7 +344,9 @@ def maybe_distribute_model(
         distributed_config: DistributedConfig | None,
         mesh_manager: TransformersDeviceMesh | None,
     ):
-        """Apply pipeline, tensor and expert parallelism, then FSDP2, after model init and before weight loading."""
+        """
+        Apply pipeline, tensor, expert and context parallelism, then FSDP2, after model init and before weight loading.
+        """
         if mesh_manager is None:
             return model
 
@@ -312,7 +357,7 @@ def maybe_distribute_model(
         model._tp_size = distributed_config.tp_size
         model._fsdp_size = distributed_config.fsdp_size
 
-        tp_plan, ep_plan, fsdp_plan = model._resolve_parallel_plans(distributed_config)
+        tp_plan, ep_plan, fsdp_plan, cp_plan = model._resolve_parallel_plans(distributed_config)
 
         if distributed_config.pp_size > 1:
             model = apply_pipeline_parallelism(model, mesh_manager.get_mesh("pp"))
@@ -333,6 +378,9 @@ def maybe_distribute_model(
                 # we will have to slice the batch here in order to avoid computing tp_size times the same batch.
                 model = apply_expert_parallelism(model, ep_mesh, tp_mesh, ep_plan)
 
+        if distributed_config.cp_size > 1:
+            model = apply_context_parallelism(model, mesh_manager.get_mesh("cp"), cp_plan)
+
         if distributed_config.dp_size > 1:
             model = apply_fully_sharded_data_parallelism(model, mesh_manager, fsdp_plan, ep_plan)
 

src/transformers/distributed/tensor_parallel.py

@@ -664,7 +664,7 @@ def resolve_parallel_plans(
         "`resolve_parallel_plans` is deprecated and will be removed in v5.23: the model resolves its plans itself "
         "in `_resolve_parallel_plans`, which also returns the FSDP plan."
     )
-    tp_plan, ep_plan, _ = model._resolve_parallel_plans(distributed_config)
+    tp_plan, ep_plan, _, _ = model._resolve_parallel_plans(distributed_config)
     return tp_plan, ep_plan