Support shard_map for data parallel support in Softserve

GitOrigin-RevId: fbf0e9509aec5db2d006a24c8cf4cf9548e4cf36

Author: rohitc33

axlearn/common/attention.py

@@ -1023,6 +1023,8 @@ class Config(BaseQKVLinear.Config):
         num_kv_heads: Required[int] = REQUIRED
         # The layer used to project.
         layer: MultiheadInputLinear.Config = MultiheadInputLinear.default_config()
+        # Optional partition spec for query, key, value output activations.
+        output_partition_spec: Optional[PartitionSpec] = None
 
     def __init__(self, cfg: Config, *, parent: Module):
         super().__init__(cfg, parent=parent)
@@ -1073,17 +1075,18 @@ def forward(
         q_proj, k_proj, v_proj = jnp.split(
             proj, [cfg.num_heads, cfg.num_heads + cfg.num_kv_heads], axis=-2
         )
-        # 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,
-        )
+        if (spec := cfg.output_partition_spec) is 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)