[distributed] Context parallelism with tensor parallelism

src/transformers/distributed/configuration_utils.py

@@ -61,11 +61,13 @@ class DistributedConfig:
             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. Requires `torch>=2.10`.
+            batch, and each runs on a contiguous slice of every sequence. The ranks of a tp group get the same batch
+            and the same slice. 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`.
+            `"cp_ulysses"` needs `cp_size * tp_size` to divide the number of attention and key/value heads.
         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*):
@@ -191,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", "tp_size", "ep_size") if getattr(self, name) > 1]:
+            if unsupported := [name for name in ("pp_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)`."
+                    "Use `DistributedConfig(fsdp_size=N, tp_size=T, cp_size=M)`."
                 )
 
         if self.dp_size > 1 and self.pp_size > 1:

src/transformers/distributed/context_parallel.py

@@ -159,7 +159,10 @@ def install_forward(self, module: nn.Module, cp_mesh: DeviceMesh):
 
 
 class UlyssesAttention(ContextParallelAttention):
-    """All-to-all to the whole sequence with `num_heads / cp_size` heads, run the attention, all-to-all back."""
+    """
+    All-to-all to the whole sequence with `1 / cp_size` of the heads this rank holds, run the attention, all-to-all
+    back. Under tensor parallelism this rank holds `num_heads / tp_size` heads.
+    """
 
     def check(self, module: nn.Module, cp_mesh: DeviceMesh) -> None:
         from ..modeling_utils import ALL_ATTENTION_FUNCTIONS
@@ -176,12 +179,14 @@ def check(self, module: nn.Module, cp_mesh: DeviceMesh) -> None:
         config = module.config.get_text_config()
         num_kv_heads = getattr(config, "num_key_value_heads", None) or config.num_attention_heads
         cp_size = cp_mesh.size()
-        if config.num_attention_heads % cp_size or num_kv_heads % cp_size:
+        distributed_config = getattr(module.config, "distributed_config", None)
+        tp_size = 1 if distributed_config is None else distributed_config.tp_size
+        if config.num_attention_heads % (cp_size * tp_size) or num_kv_heads % (cp_size * tp_size):
             raise ValueError(
-                f"Ulysses context parallelism splits the attention heads across the {cp_size} ranks of the `cp` mesh, "
-                f"so `cp_size` must divide both the number of attention heads ({config.num_attention_heads}) and the "
-                f"number of key/value heads ({num_kv_heads}) of {module.__class__.__name__}. Use another `cp_size`, or "
-                "disable context parallelism."
+                f"Ulysses context parallelism splits the attention heads across the {cp_size} ranks of the `cp` mesh "
+                f"and the {tp_size} ranks of the `tp` mesh, so `cp_size * tp_size` must divide both the number of "
+                f"attention heads ({config.num_attention_heads}) and the number of key/value heads ({num_kv_heads}) "
+                f"of {module.__class__.__name__}. Use smaller `cp_size` or `tp_size`, or disable context parallelism."
             )
 
     def attention(self, cp_mesh: DeviceMesh, attention_function, module, query, key, value, attention_mask, **kwargs):

tests/context_parallel/test_context_parallel.py

@@ -22,6 +22,7 @@
 import torch.distributed as dist
 from parameterized import parameterized
 from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.tensor import DTensor
 
 from transformers import (
     AutoModelForCausalLM,
@@ -120,14 +121,16 @@ def _init_cp_mesh(device_type):
 def _tiny_config(**kwargs):
     """A `LlamaConfig` small enough for CPU tests, in float32 with `sdpa` attention."""
     return LlamaConfig(
-        hidden_size=32,
-        intermediate_size=64,
-        num_hidden_layers=2,
-        num_attention_heads=4,
-        num_key_value_heads=2,
-        vocab_size=99,
-        attn_implementation="sdpa",
-        **kwargs,
+        **{
+            "hidden_size": 32,
+            "intermediate_size": 64,
+            "num_hidden_layers": 2,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "vocab_size": 99,
+            "attn_implementation": "sdpa",
+            **kwargs,
+        }
     )
 
 
@@ -477,6 +480,26 @@ def test_refuses_indivisible_heads(self):
         with self.assertRaisesRegex(ValueError, "must divide both the number of attention heads"):
             self._check(config, cp_size=4)
 
+    def test_head_divisibility_covers_tensor_parallelism(self):
+        # tensor parallelism keeps `num_heads / tp_size` heads on this rank, and those are what cp splits
+        def config_with_heads(num_heads, num_kv_heads):
+            config = LlamaConfig(
+                hidden_size=32,
+                num_attention_heads=num_heads,
+                num_key_value_heads=num_kv_heads,
+                attn_implementation="sdpa",
+            )
+            config.distributed_config = DistributedConfig(tp_size=2, cp_size=2)
+            return config
+
+        self._check(config_with_heads(8, 4), cp_size=2)
+        # 1 key/value head left on each rank cannot be split across 2 cp ranks
+        with self.assertRaisesRegex(ValueError, r"`cp_size \* tp_size` must divide"):
+            self._check(config_with_heads(8, 2), cp_size=2)
+        # so few attention heads left on each rank that 2 cp ranks cannot split them either
+        with self.assertRaisesRegex(ValueError, r"`cp_size \* tp_size` must divide"):
+            self._check(config_with_heads(8, 4), cp_size=4)
+
     def test_accepts_supported_config(self):
         self._check(LlamaConfig(attn_implementation="sdpa", num_attention_heads=8, num_key_value_heads=4))
         # the mask the model builds over the whole sequence already has the sliding window in it
@@ -654,6 +677,44 @@ def _compare_with_unsharded(init_dir, batch, distributed_config, keep=None):
         torch.testing.assert_close(grad, expected, rtol=CP_RTOL, atol=CP_ATOL, msg=f"Gradient mismatch for {name}")
 
 
+def _test_matches_unsharded_with_tp(rank, tp_size, cp_size, fsdp_size, packed):
+    """The tp ranks hold the same batch and the same cp slice, and split the heads before cp splits them again."""
+    # 4 attention heads and 4 key/value heads, both divisible by `tp_size * cp_size` = 4
+    config = _tiny_config(num_key_value_heads=4)
+    batch_size, seq_len = 2, 8
+    with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
+        input_ids, position_ids = _tiny_batch(config, batch_size * fsdp_size, seq_len, packed)
+        distributed_config = DistributedConfig(tp_size=tp_size, cp_size=cp_size, fsdp_size=fsdp_size)
+        batch = {"input_ids": input_ids, "labels": input_ids.clone(), "use_cache": False}
+        if position_ids is not None:
+            batch["position_ids"] = position_ids
+        _compare_with_unsharded(init_dir, batch, distributed_config)
+        dist.barrier()
+
+
+def _test_tp_and_cp_compose(rank, tp_size, cp_size):
+    """The tp plan shards the projections of an attention layer, the cp plan registers its attention call."""
+    config = _tiny_config(num_key_value_heads=4)
+    with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
+        distributed_config = DistributedConfig(tp_size=tp_size, cp_size=cp_size)
+        model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config)
+        attention = model.model.layers[0].self_attn
+        # tensor parallelism sharded the projections of the very module that context parallelism registered
+        assert isinstance(attention.q_proj.weight, DTensor), type(attention.q_proj.weight)
+        assert attention.cp_attention.func == ALL_PARALLEL_STYLES["cp_ulysses"].attention
+
+        # and one forward goes through both: the logits of this rank's slice match the unsharded run
+        input_ids, _ = _tiny_batch(config, 2, 8)
+        with torch.no_grad():
+            logits = model(input_ids=input_ids).logits
+        reference = AutoModelForCausalLM.from_pretrained(init_dir)
+        with torch.no_grad():
+            expected = reference(input_ids=input_ids).logits
+        cp_rank = model._mesh_manager.get_mesh("cp").get_local_rank()
+        torch.testing.assert_close(logits, expected.chunk(cp_size, dim=1)[cp_rank], rtol=CP_RTOL, atol=CP_ATOL)
+        dist.barrier()
+
+
 def _test_matches_unsharded(rank, cp_size, fsdp_size, dp_replicate_size, packed):
     config = _tiny_config()
     batch_size, seq_len = 2, 8
@@ -742,15 +803,16 @@ def _test_custom_attention_style(rank, cp_size):
         dist.barrier()
 
 
-def _test_save_pretrained(rank, cp_size):
+def _test_save_pretrained(rank, cp_size, tp_size=1):
     """Weights saved from a context parallel model reload without a distributed config, unchanged."""
-    config = _tiny_config()
+    config = _tiny_config(num_key_value_heads=4 if tp_size > 1 else 2)
     with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
         reference = AutoModelForCausalLM.from_pretrained(init_dir)
         expected = {name: value.clone() for name, value in reference.state_dict().items()}
         del reference
 
-        model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=DistributedConfig(cp_size=cp_size))
+        distributed_config = DistributedConfig(cp_size=cp_size, tp_size=tp_size)
+        model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config)
         with _distributed_tmpdir(rank) as save_dir:
             model.save_pretrained(save_dir, is_main_process=(rank == 0))
             dist.barrier()
@@ -784,10 +846,9 @@ 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_other_parallelisms(self):
-        for other in ({"pp_size": 2}, {"tp_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`"):
+            DistributedConfig(cp_size=2, pp_size=2)
 
     def test_refuses_a_model_without_a_cp_plan(self):
         config = _tiny_config()
@@ -876,6 +937,14 @@ def test_sliding_window_and_full_layers(self):
         )
         run_distributed(world_size=2, backend="gloo")(_test_padded)(2, config, False)
 
+    def test_matches_unsharded_with_tensor_parallelism(self):
+        for packed in (False, True):
+            with self.subTest(packed=packed):
+                run_distributed(world_size=4, backend="gloo")(_test_matches_unsharded_with_tp)(2, 2, 1, packed)
+
+    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_inputs_embeds(self):
         run_distributed(world_size=2, backend="gloo")(_test_inputs_embeds)(2)
 
@@ -885,5 +954,8 @@ def test_custom_attention_style(self):
     def test_save_pretrained(self):
         run_distributed(world_size=2, backend="gloo")(_test_save_pretrained)(2)
 
+    def test_save_pretrained_with_tensor_parallelism(self):
+        run_distributed(world_size=4, backend="gloo")(_test_save_pretrained)(2, 2)
+
     def test_save_load_distributed_checkpoint(self):
         run_distributed(world_size=2, backend="gloo")(_test_save_load_distributed_checkpoint)(2)

tests/distributed_testing_utils.py

@@ -102,6 +102,10 @@ def _worker_main(rank, world_size, master_port, backend, pipe):
         for group in list(_world.pg_map):
             if group is not GroupMember.WORLD:
                 dist.destroy_process_group(group)
+        # Torch caches DTensor dispatch results by mesh, and a mesh rebuilt from the same ranks compares equal to the
+        # destroyed one: drop the entries that still point at the destroyed groups.
+        if clear_mesh_cache := getattr(torch._C, "_clear_DTensor_sharding_propagator_cache", None):
+            clear_mesh_cache()
         set_seed(42)
         try:
             func(rank, *args, **kwargs)