Supports TPU context parallel training (#981)

Fix

Fix tests

Author: hanzhi713

axlearn/common/attention.py

@@ -82,6 +82,7 @@
 import einops
 import jax
 from jax import numpy as jnp
+from jax._src.mesh import thread_resources
 
 from axlearn.common import ops, param_init
 from axlearn.common.attention_bias import (
@@ -149,6 +150,7 @@
     save_and_offload_only_these_names_regex,
     shapes,
     split_prng_key,
+    with_sharding_constraint,
 )
 
 
@@ -1419,6 +1421,20 @@ def forward(
         )
         if query_positions is None:
             query_positions = jnp.arange(q_proj.shape[1])[None]
+        # This sharding hint is needed since compiler sometimes will generate large allgather
+        # before the split and then slice, which is not the ideal compilation. Ensure sharding
+        # after the split to ensure allgather is inserted after the split.
+        axis_names = thread_resources.env.physical_mesh.axis_names
+        batch_axes = tuple(x for x in axis_names if x in ("data", "fsdp")) or None
+        spec = PartitionSpec(
+            batch_axes,
+            "seq" if "seq" in axis_names else None,
+            "model" if "model" in axis_names else None,
+            None,
+        )
+        q_proj = with_sharding_constraint(q_proj, spec)
+        k_proj = with_sharding_constraint(k_proj, spec)
+        v_proj = with_sharding_constraint(v_proj, spec)
         return self.Output(
             query=q_proj,
             key=k_proj,

axlearn/common/attention_bias.py

@@ -439,10 +439,16 @@ def partition_spec(
         self, mha_dim_to_partition_spec: dict[str, PartitionSpec]
     ) -> Union[BaseAttentionBias, PartitionSpec]:
         # Segment IDs: [batch_size, seq_len].
-        q_spec = mha_dim_to_partition_spec["btnh"]
-        if q_spec == PartitionSpec(None):
+        # We use the partition spec of KV (which are not sequence sharded) for segment ids. This is
+        # because Splash requires two seq ids, q_seg and kv_seg. Therefore, we pass a not seq
+        # sharded seg ids into the shard map, and manually shard it inside for q_seg and not
+        # shard it for kv_seg.
+        kv_spec = mha_dim_to_partition_spec["bsnh"]
+        if kv_spec == PartitionSpec(None):
             return PartitionSpec(None)
-        return PartitionSpec(q_spec[0], q_spec[1])
+        if kv_spec[1] is not None:
+            raise ValueError("The partition spec of `s` in `bsnh` should be None.")
+        return PartitionSpec(kv_spec[0], kv_spec[1])
 
 
 class MaskFn(Protocol):

axlearn/common/flash_attention/layer.py

@@ -15,10 +15,7 @@
 from axlearn.common.attention import Dropout, GroupedQueryAttention
 from axlearn.common.attention_bias import BaseAttentionBias
 from axlearn.common.config import REQUIRED, ConfigBase, ConfigModifier, Required, config_class
-from axlearn.common.flash_attention.utils import (
-    MultiHeadAttentionImpl,
-    flash_attention_implementation,
-)
+from axlearn.common.flash_attention.utils import flash_attention_implementation
 from axlearn.common.module import Module
 from axlearn.common.utils import Tensor, with_sharding_constraint
 
@@ -167,13 +164,6 @@ def _compute_attention(
 
         attention_logit_biases = attention_logit_biases.astype(q_proj.dtype)
 
-        jit_attn: MultiHeadAttentionImpl = flash_attention_implementation(
-            backend=backend,
-            softmax_scale=1.0,
-            block_size=cfg.tpu_block_size,
-            dropout_rate=cfg.dropout.rate,
-        )
-
         attention_logit_biases_spec = self._logit_biases_spec(attention_logit_biases)
         attention_logit_biases = with_sharding_constraint(
             attention_logit_biases, attention_logit_biases_spec
@@ -188,10 +178,21 @@ def _compute_attention(
         k_proj = with_sharding_constraint(k_proj, cfg.mha_dim_to_partition_spec["bsnh"])
         v_proj = with_sharding_constraint(v_proj, cfg.mha_dim_to_partition_spec["bsnh"])
 
+        shard_map_specs = flash_attention_implementation(
+            query=q_proj,
+            key=k_proj,
+            value=v_proj,
+            bias=attention_logit_biases,
+            backend=backend,
+            softmax_scale=1.0,
+            block_size=cfg.tpu_block_size,
+            dropout_rate=cfg.dropout.rate,
+        )
+
         # We need to manually partition pallas | jax-triton calls.
         # Note: shard_map doesn't support kwargs.
         partitioned_mha = shard_map(
-            jit_attn,
+            shard_map_specs.fn,
             mesh=thread_resources.env.physical_mesh,
             in_specs=(
                 # Q [batch_size, seq_len, num_heads, per_head_dim].
@@ -204,6 +205,7 @@ def _compute_attention(
                 attention_logit_biases_spec,
                 # PRNG Key.
                 PartitionSpec(None),
+                *shard_map_specs.additional_in_specs,
             ),
             # O [batch_size, seq_len, num_heads, per_head_dim].
             out_specs=cfg.mha_dim_to_partition_spec["btnh"],
@@ -221,6 +223,7 @@ def _compute_attention(
                 v_proj,
                 attention_logit_biases,
                 self.dropout.get_prng_key(),
+                *shard_map_specs.additional_args,
             ),
             cfg.output_dim_to_partition_spec["btnh"],
         )
@@ -247,12 +250,13 @@ def default_mha_dim_to_partition_spec(
     Returns:
         A dictionary keyed by MHA tensor dims with partition spec values.
     """
-    batch_axis_names = tuple(el for el in mesh_axis_names if el != "model")
+    batch_axis_names = tuple(el for el in mesh_axis_names if el in ["data", "fsdp"])
     tp_axis_name = "model" if "model" in mesh_axis_names else None
+    sp_axis_name = "seq" if "seq" in mesh_axis_names else None
     return {
-        "btnh": PartitionSpec(batch_axis_names, None, tp_axis_name, None),
+        "btnh": PartitionSpec(batch_axis_names, sp_axis_name, tp_axis_name, None),
         "bsnh": PartitionSpec(batch_axis_names, None, tp_axis_name, None),
-        "bnts": PartitionSpec(batch_axis_names, tp_axis_name, None, None),
+        "bnts": PartitionSpec(batch_axis_names, tp_axis_name, sp_axis_name, None),
     }
 
 
@@ -271,7 +275,7 @@ def default_output_dim_to_partition_spec(
     Returns:
         A dictionary keyed by FlashAttention output tensor dims with partition spec values.
     """
-    batch_axis_names = tuple(el for el in mesh_axis_names if el not in ["seq", "model"])
+    batch_axis_names = tuple(el for el in mesh_axis_names if el in ["data", "fsdp"])
     tp_axis_name = "model" if "model" in mesh_axis_names else None
     sp_axis_name = "seq" if "seq" in mesh_axis_names else None
     return {

axlearn/common/flash_attention/layer_test.py

@@ -451,6 +451,19 @@ def test_backend(
             backend = test_layer._backend()  # pylint: disable=protected-access
             self.assertEqual(backend, "tpu")
 
+    def _maybe_skip_unsupported_context_parallel(
+        self, mesh, mesh_axis_names, use_bias, per_head_dim
+    ):
+        # TODO(hanzhi-zhou): Add GPU support.
+        if jax.default_backend() != "tpu":
+            self.skipTest("Context parallelism is only supported on TPU for now.")
+        for mesh_dim, mesh_name in zip(mesh, mesh_axis_names):
+            if mesh_name == "seq" and mesh_dim > 1 and (use_bias or per_head_dim % 128 != 0):
+                self.skipTest(
+                    "Context parallelism is not supported when need to fallback to legacy TPU"
+                    " flash attention."
+                )
+
     @parameterized.parameters(_TEST_CONFIGS)
     def test_shard_biases(
         self, batch, seq_len, num_heads, num_kv_heads, per_head_dim, mesh, mesh_axis_names
@@ -501,7 +514,7 @@ def as_partition_spec(pytree: CompositeAttentionBias) -> PartitionSpec:
             spec = test_layer._logit_biases_spec(segment_ids)  # pylint: disable=protected-access
             spec = as_partition_spec(spec)
             self.assertIsInstance(spec, PartitionSpec)
-            self.assertEqual(spec, test_layer.config.mha_dim_to_partition_spec["btnh"][:2])
+            self.assertEqual(spec, test_layer.config.mha_dim_to_partition_spec["bsnh"][:2])
 
     @parameterized.product(
         _TEST_CONFIGS,
@@ -534,6 +547,10 @@ def test_forward(
             pytest.skip(reason=f"Unsupported mesh {mesh}.")
         if not causal and left_context is not None:
             pytest.skip(reason="Sliding window attention must be causal.")
+        if query_len_multiplier > 1 and left_context is not None:
+            # When sliding window is enabled and q_len > kv_len, there might be be fully masked
+            # rows.
+            pytest.skip(reason="Sliding window attention does not make sense when q_len > kv_len.")
         if causal and use_bias:
             # TODO(c_lan): Investigate the numerical errors when both causal and bias are used.
             pytest.skip(reason="Only one of causal and use_bias can be True.")
@@ -544,6 +561,7 @@ def test_forward(
             pytest.skip(reason="Unsupported large bias matrix in fp32 format.")
         if dropout_rate > 0.0 and jax.default_backend() == "tpu":
             pytest.skip("Dropout is implemented for GPU only.")
+        self._maybe_skip_unsupported_context_parallel(mesh, mesh_axis_names, use_bias, per_head_dim)
 
         with Mesh(mesh_utils.create_device_mesh(mesh), mesh_axis_names):
             test_layer, ref_layer, params, hidden_dim = _prepare_layers(
@@ -630,6 +648,7 @@ def test_backward(
         if attn_type == "causal" and use_bias:
             # TODO(c_lan): Investigate the numerical errors when both causal and bias are used.
             pytest.skip(reason="Only one of causal and use_bias can be True.")
+        self._maybe_skip_unsupported_context_parallel(mesh, mesh_axis_names, use_bias, per_head_dim)
 
         with Mesh(mesh_utils.create_device_mesh(mesh), mesh_axis_names):
             hidden_dim = num_heads * per_head_dim