Allow qkvo partition specs

* Allow qkvo partition specs

* Add unit test

* Add missing skip

GitOrigin-RevId: 143aaef

Author: hanzhi713

axlearn/common/attention.py

@@ -162,6 +162,7 @@
     save_and_offload_only_these_names_regex,
     shapes,
     split_prng_key,
+    with_sharding_constraint,
 )
 
 
@@ -1556,6 +1557,20 @@ class Config(BaseLayer.Config):
         # If true, use learnable logit sinks.
         logit_sink: Optional[bool] = None
 
+        # Partition spec for query ([batch, seq, q_heads, head_dim]) after input projections.
+        q_partition_spec: Optional[PartitionSpec] = None
+
+        # Partition spec for key ([batch, seq, kv_heads, head_dim]) after input projections.
+        # Follows `q_partition_spec` if None.
+        k_partition_spec: Optional[PartitionSpec] = None
+
+        # Partition spec for value ([batch, seq, kv_heads, head_dim]) after input projections.
+        # Follows `q_partition_spec` if None.
+        v_partition_spec: Optional[PartitionSpec] = None
+
+        # Partition spec for output ([batch, seq, hidden_dim]) after output projections.
+        o_partition_spec: Optional[PartitionSpec] = None
+
     def __init__(self, cfg: Config, *, parent: Module):
         super().__init__(cfg, parent=parent)
         cfg = self.config
@@ -1719,6 +1734,12 @@ def _forward_for_mode(
             time_step = cached_states["time_step"]
             query_positions = query_positions + time_step[:, None]  # [batch, steps]
         q_proj, k_proj, v_proj = self.i_proj(query, query_positions=query_positions, **kv_kwargs)
+        if cfg.q_partition_spec:
+            q_proj = with_sharding_constraint(q_proj, cfg.q_partition_spec)
+        if cfg.q_partition_spec or cfg.k_partition_spec:
+            k_proj = with_sharding_constraint(k_proj, cfg.k_partition_spec or cfg.q_partition_spec)
+        if cfg.q_partition_spec or cfg.v_partition_spec:
+            v_proj = with_sharding_constraint(v_proj, cfg.v_partition_spec or cfg.q_partition_spec)
 
         if cfg.scale_kv_before_cache_update:
             if has_external_kv_state:
@@ -1821,6 +1842,8 @@ def _forward_for_mode(
 
         # [batch, target_length, output_dim].
         o_proj = self.o_proj(context)
+        if cfg.o_partition_spec:
+            o_proj = with_sharding_constraint(o_proj, cfg.o_partition_spec)
         outputs = self._remat_name(o_proj, "o_proj")
         self._add_tensor_stats("o_proj_outputs", outputs)
         return_aux = return_aux or set()

axlearn/common/attention_test.py

@@ -129,6 +129,7 @@
     TestCase,
     assert_allclose,
     dummy_segments_positions,
+    is_supported_mesh_shape,
     set_threefry_partitionable,
 )
 from axlearn.common.torch_utils import parameters_from_torch_layer
@@ -2466,6 +2467,82 @@ def test_gqa_forward(
         )
         self.assertNestedAllClose(base_outputs, test_outputs)
 
+    @parameterized.product(kv_part=[None, PartitionSpec("fsdp", None, "model", None)])
+    @pytest.mark.d8
+    def test_qkvo_partition_spec(self, kv_part):
+        """Tests that QKVO partition spec are applied correctly when specified."""
+        mesh_shape = (2, 2, 2)
+        if not is_supported_mesh_shape(mesh_shape):
+            self.skipTest(f"Unsupported mesh shape {mesh_shape}")
+        model_dim = 16
+        num_heads = 4
+        mesh = jax.make_mesh(mesh_shape, axis_names=("fsdp", "seq", "model"))
+        q_part = PartitionSpec("fsdp", "seq", "model", None)
+        o_part = PartitionSpec("fsdp", "seq", None)
+
+        layer_kwargs = dict(
+            query_dim=model_dim,
+            key_dim=model_dim,
+            value_dim=model_dim,
+            num_heads=num_heads,
+            dtype=jnp.float32,
+            q_partition_spec=q_part,
+            o_partition_spec=o_part,
+            k_partition_spec=kv_part,
+            v_partition_spec=kv_part,
+        )
+        init_key = jax.random.PRNGKey(123)
+        base_cfg = attention.MultiheadAttention.default_config().set(**layer_kwargs)
+        base_layer = base_cfg.set(name="base").instantiate(parent=None)
+        base_state = base_layer.initialize_parameters_recursively(prng_key=init_key)
+
+        # Dummy inputs.
+        batch_size, tgt_len = 2, 6
+        base_inputs = dict(
+            query=jax.random.normal(
+                jax.random.PRNGKey(124),
+                [batch_size, tgt_len, model_dim],
+                dtype=jnp.float32,
+            ),
+            key=None,
+            value=None,
+        )
+        forward_key = jax.random.PRNGKey(456)
+
+        def patched_remat_name(_, tensor, name):
+            def callback(sharding):
+                # pylint: disable-next=protected-access
+                normalize_spec = sharding.spec._normalized_spec_for_aval(len(tensor.shape))
+                if name == "q_proj":
+                    self.assertEqual(normalize_spec, q_part)
+                elif name == "o_proj":
+                    self.assertEqual(normalize_spec, o_part)
+                elif name in ["k_proj", "v_proj"]:
+                    if kv_part is None:
+                        self.assertEqual(normalize_spec, q_part)
+                    else:
+                        self.assertEqual(normalize_spec, kv_part)
+
+            jax.debug.inspect_array_sharding(tensor, callback=callback)
+            return tensor
+
+        with mesh, mock.patch.object(
+            attention.MultiheadAttention, "_remat_name", patched_remat_name
+        ):
+
+            @jax.jit
+            def jit_fn():
+                base_outputs, _ = F(
+                    base_layer,
+                    state=base_state,
+                    is_training=False,
+                    prng_key=forward_key,
+                    inputs=base_inputs,
+                )
+                return base_outputs
+
+            jit_fn()
+
     def _test_extend_step(
         self,
         attention_cfg: attention.MultiheadAttention.Config,