Copybara import of the project:

--
e43c2e1 by NuojCheng <[email protected]>:

remove size one mesh axis

COPYBARA_INTEGRATE_REVIEW=#2737 from AI-Hypercomputer:chengnuojin-remove-sizeone e43c2e1
PiperOrigin-RevId: 842434394

Author: NuojCheng

src/MaxText/inference/paged_attention.py

@@ -30,6 +30,7 @@
 
 from MaxText.inference import page_manager
 from MaxText.inference import paged_attention_kernel_v2
+from MaxText.sharding import logical_to_mesh_axes
 from MaxText.common_types import Array, DType, AxisNames, BATCH, LENGTH, HEAD, D_KV, MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE
 from MaxText.layers.initializers import variable_to_logically_partitioned
 
@@ -322,8 +323,8 @@ def paged_attention_v1_decode(
       page_state: page_manager.PageState,
   ) -> Array:
     """Apply Paged Attention v1 in decode only."""
-    kv_pages_pspec = nn.logical_to_mesh_axes(("paged_kv_heads", None, None, None))
-    q_pspec = nn.logical_to_mesh_axes((None, None, "paged_kv_heads", None))
+    kv_pages_pspec = logical_to_mesh_axes(("paged_kv_heads", None, None, None), self.mesh)
+    q_pspec = logical_to_mesh_axes((None, None, "paged_kv_heads", None), self.mesh)
 
     @functools.partial(
         jax.shard_map,

src/MaxText/layers/attention_mla.py

@@ -21,7 +21,6 @@
 from jax.sharding import Mesh, NamedSharding
 import jax.numpy as jnp
 
-from flax import linen as nn
 from flax import nnx
 
 from MaxText.common_types import (
@@ -56,6 +55,7 @@
 from MaxText.inference import page_manager
 from MaxText.inference import paged_attention
 from MaxText.inference.kvcache import KVQuant
+from MaxText.sharding import create_sharding
 from MaxText.layers import nnx_wrappers
 from MaxText.layers.attentions import Attention
 from MaxText.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned
@@ -515,8 +515,8 @@ def mla_query_projection(self, inputs_q: Array, inputs_positions: Array, model_m
     else:
       query_logical_name = self.query_axis_names
       wqa_logical_name = (KV_BATCH, LENGTH_NO_EXP, Q_LORA_UP_PROJ)
-    query_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(query_logical_name))
-    wqa_out_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(wqa_logical_name))
+    query_sharding = create_sharding(self.mesh, query_logical_name)
+    wqa_out_sharding = create_sharding(self.mesh, wqa_logical_name)
     # Set softmax scaling.
     self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
     self.softmax_scale = self.qk_head_dim**-0.5
@@ -555,7 +555,7 @@ def mla_get_key_value(self, low_rank_main, key_rope, model_mode):
       key_logical_name = self.key_axis_names
       value_logical_name = self.value_axis_names
 
-    wkva_out_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(key_logical_name))
+    wkva_out_sharding = create_sharding(self.mesh, key_logical_name)
     kv_out = self.wkv_b(low_rank_main, out_sharding=wkva_out_sharding)
 
     # Split kv_out into key_nope and value parts.
@@ -664,7 +664,7 @@ def mla_kv_projection(self, inputs: Array, inputs_positions: Array, decoder_segm
       wka_logical_name = (KV_BATCH_NO_EXP, LENGTH, KV_LORA_UP_PROJ)
     else:
       wka_logical_name = (KV_BATCH, LENGTH_NO_EXP, KV_LORA_UP_PROJ)
-    wkva_out_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(wka_logical_name))
+    wkva_out_sharding = create_sharding(self.mesh, wka_logical_name)
     low_rank = self.wkv_a(inputs, out_sharding=wkva_out_sharding)
     low_rank_main, low_rank_rope = jnp.split(low_rank, [self.kv_lora_rank], axis=-1)
     low_rank_main = self.kv_norm(low_rank_main)
@@ -759,7 +759,7 @@ def __call__(
     else:
       out = self._maybe_shard_with_logical(out, self.out_axis_names)
 
-    out_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(out_logical_name))
+    out_sharding = create_sharding(self.mesh, out_logical_name)
     out = self.out_projection(out, out_sharding=out_sharding)
     out = checkpoint_name(out, "out_proj")
     return out, kv_cache

src/MaxText/layers/attention_op.py

@@ -43,7 +43,7 @@
 
 
 from MaxText import max_utils
-from MaxText.sharding import maybe_shard_with_name
+from MaxText.sharding import maybe_shard_with_name, logical_to_mesh_axes
 from MaxText.common_types import (
     DEFAULT_MASK_VALUE,
     BATCH,
@@ -530,6 +530,9 @@ def maybe_create_nnx(einsum, *args):
       self.AqtEinsum_2 = jnp.einsum
       self.AqtEinsum_3 = jnp.einsum
 
+  def _logical_to_mesh_axes(self, logical_name):
+    return logical_to_mesh_axes(logical_name, mesh=self.mesh, rules=self.config.logical_axis_rules)
+
   def check_attention_inputs(self, query: Array, key: Array | KVTensor, value: Array | KVTensor) -> None:
     """Check attention inputs."""
 
@@ -950,10 +953,10 @@ def gpu_ragged_attention(self, q: Array, k: Array | KVTensor, v: Array | KVTenso
     q_for_gqa = q.squeeze(axis=1)
 
     # Define logical axis names - clearer and avoids repeated calls.
-    b = nn.logical_to_mesh_axes(self.ragged_lengths_names)
-    bsnd = nn.logical_to_mesh_axes(self.cache_logical_axis_names)
-    bnd = nn.logical_to_mesh_axes((CACHE_BATCH, CACHE_HEADS, CACHE_KV))
-    bn = nn.logical_to_mesh_axes((CACHE_BATCH, CACHE_HEADS))
+    b = self._logical_to_mesh_axes(self.ragged_lengths_names)
+    bsnd = self._logical_to_mesh_axes(self.cache_logical_axis_names)
+    bnd = self._logical_to_mesh_axes((CACHE_BATCH, CACHE_HEADS, CACHE_KV))
+    bn = self._logical_to_mesh_axes((CACHE_BATCH, CACHE_HEADS))
 
     @functools.partial(
         jax.shard_map,
@@ -1006,8 +1009,8 @@ def tpu_ragged_attention(
     """Ragged Attention."""
     if isinstance(query, KVTensor):
       raise TypeError("Ragged attention does not currently support quantized tensors.")
-    b = nn.logical_to_mesh_axes(self.ragged_lengths_names)
-    bsnd = nn.logical_to_mesh_axes(self.cache_logical_axis_names)
+    b = self._logical_to_mesh_axes(self.ragged_lengths_names)
+    bsnd = self._logical_to_mesh_axes(self.cache_logical_axis_names)
 
     @functools.partial(
         jax.shard_map,
@@ -1050,23 +1053,23 @@ def tpu_flash_attention(
     value = jnp.transpose(value, axes=(0, 2, 1, 3))
     segment_axis_names_q = None
     segment_axis_names_kv = None
-    sink_axis_names = nn.logical_to_mesh_axes((HEAD,))
+    sink_axis_names = self._logical_to_mesh_axes((HEAD,))
     if decoder_segment_ids is not None:
       if self.config.expert_shard_attention_option == EP_AS_CONTEXT:
-        segment_axis_names_q = nn.logical_to_mesh_axes((BATCH_NO_EXP, Q_LENGTH))
-        segment_axis_names_kv = nn.logical_to_mesh_axes((BATCH_NO_EXP, KV_LENGTH))
+        segment_axis_names_q = self._logical_to_mesh_axes((BATCH_NO_EXP, Q_LENGTH))
+        segment_axis_names_kv = self._logical_to_mesh_axes((BATCH_NO_EXP, KV_LENGTH))
       else:
-        segment_axis_names_q = nn.logical_to_mesh_axes((BATCH, Q_LENGTH_NO_EXP))
-        segment_axis_names_kv = nn.logical_to_mesh_axes((BATCH, KV_LENGTH))
+        segment_axis_names_q = self._logical_to_mesh_axes((BATCH, Q_LENGTH_NO_EXP))
+        segment_axis_names_kv = self._logical_to_mesh_axes((BATCH, KV_LENGTH))
 
     if self.config.expert_shard_attention_option == EP_AS_CONTEXT:
-      axis_names_splash_kernel = nn.logical_to_mesh_axes(self.flash_axis_names_splash_kernel_ep)
-      axis_names_q = nn.logical_to_mesh_axes(self.flash_axis_names_q_ep)
-      axis_names_kv = nn.logical_to_mesh_axes(self.flash_axis_names_kv_ep)
+      axis_names_splash_kernel = self._logical_to_mesh_axes(self.flash_axis_names_splash_kernel_ep)
+      axis_names_q = self._logical_to_mesh_axes(self.flash_axis_names_q_ep)
+      axis_names_kv = self._logical_to_mesh_axes(self.flash_axis_names_kv_ep)
     else:
-      axis_names_splash_kernel = nn.logical_to_mesh_axes(self.flash_axis_names_splash_kernel)
-      axis_names_q = nn.logical_to_mesh_axes(self.flash_axis_names_q)
-      axis_names_kv = nn.logical_to_mesh_axes(self.flash_axis_names_kv)
+      axis_names_splash_kernel = self._logical_to_mesh_axes(self.flash_axis_names_splash_kernel)
+      axis_names_q = self._logical_to_mesh_axes(self.flash_axis_names_q)
+      axis_names_kv = self._logical_to_mesh_axes(self.flash_axis_names_kv)
 
     global global_block_q, global_block_kv, global_block_kv_compute, global_block_q_dkv, global_block_kv_dkv
     global global_block_kv_dkv_compute, global_block_q_dq, global_block_kv_dq, global_use_fused_bwd_kernel
@@ -1197,9 +1200,9 @@ def wrap_splash_kernel(single_head_mask, shard_head_size=1):
       shard_head_size = np.prod(logical_axis_rules_head)
       splash_kernel = wrap_splash_kernel(single_head_mask, int(shard_head_size))
       if self.config.expert_shard_attention_option == EP_AS_CONTEXT:
-        segment_axis_names_splash_kernel = nn.logical_to_mesh_axes((Q_LENGTH,))
+        segment_axis_names_splash_kernel = self._logical_to_mesh_axes((Q_LENGTH,))
       else:
-        segment_axis_names_splash_kernel = nn.logical_to_mesh_axes((Q_LENGTH_NO_EXP,))
+        segment_axis_names_splash_kernel = self._logical_to_mesh_axes((Q_LENGTH_NO_EXP,))
     else:
       # Create multi-head mask
       multi_head_mask = splash_attention_mask.MultiHeadMask(masks=(mask,) * query.shape[1])

src/MaxText/layers/attentions.py

@@ -23,7 +23,7 @@
 import jax
 import jax.numpy as jnp
 
-from flax import nnx, linen as nn
+from flax import nnx
 
 from MaxText.common_types import (
     DecoderBlockType,
@@ -53,7 +53,7 @@
     EP_AS_CONTEXT,
     AttentionType,
 )
-from MaxText.sharding import maybe_shard_with_logical
+from MaxText.sharding import maybe_shard_with_logical, create_sharding
 from MaxText.inference import kvcache
 from MaxText.inference import page_manager
 from MaxText.inference import paged_attention
@@ -1003,7 +1003,7 @@ def __call__(
 
     inputs_q = self._maybe_shard_with_logical(inputs_q, input_axis_names)
     inputs_kv = self._maybe_shard_with_logical(inputs_kv, input_axis_names)
-    qkv_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(input_axis_names))
+    qkv_sharding = create_sharding(self.mesh, input_axis_names)
 
     # apply projection.
     if self.config.fused_qkv:

src/MaxText/layers/decoders.py

@@ -22,7 +22,7 @@
 import jax
 import jax.numpy as jnp
 from jax.ad_checkpoint import checkpoint_name
-from jax.sharding import Mesh, NamedSharding
+from jax.sharding import Mesh
 
 from flax import linen as nn
 from flax import nnx
@@ -32,6 +32,7 @@
 from MaxText.common_types import MODEL_MODE_TRAIN, MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE
 from MaxText import max_logging
 from MaxText import max_utils
+from MaxText.sharding import create_sharding
 from MaxText.inference import page_manager
 from MaxText.layers import linears
 from MaxText.layers import quantizations
@@ -607,16 +608,7 @@ def apply_output_head(self, shared_embedding: nn.Module | nnx.Module, y, determi
 
     cfg = self.config
     if cfg.shard_mode == ShardMode.EXPLICIT:
-      norm_out_sharding = NamedSharding(
-          self.mesh,
-          nn.logical_to_mesh_axes(
-              (
-                  "activation_batch",
-                  "activation_length_no_exp",
-                  "activation_embed",
-              )
-          ),
-      )
+      norm_out_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length_no_exp", "activation_embed"))
     else:
       norm_out_sharding = None
 
@@ -631,17 +623,10 @@ def apply_output_head(self, shared_embedding: nn.Module | nnx.Module, y, determi
     y = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(y, deterministic=deterministic)
 
     if model_mode in (MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE):
-      out_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes((None, None, "activation_vocab")))
+      out_sharding = create_sharding(self.mesh, (None, None, "activation_vocab"))
     else:
-      out_sharding = NamedSharding(
-          self.mesh,
-          nn.logical_to_mesh_axes(
-              (
-                  "activation_embed_and_logits_batch",
-                  "activation_length_no_exp",
-                  "activation_vocab",
-              )
-          ),
+      out_sharding = create_sharding(
+          self.mesh, ("activation_embed_and_logits_batch", "activation_length_no_exp", "activation_vocab")
       )
 
     # [batch, length, emb_dim] -> [batch, length, vocab_size]

src/MaxText/layers/deepseek.py

@@ -19,7 +19,7 @@
 from functools import partial
 
 from jax.ad_checkpoint import checkpoint_name
-from jax.sharding import Mesh, NamedSharding
+from jax.sharding import Mesh
 import jax.numpy as jnp
 
 from flax import linen as nn
@@ -32,7 +32,7 @@
 from MaxText.layers import moe
 from MaxText.layers import quantizations
 from MaxText.layers.quantizations import AqtQuantization as Quant
-from MaxText.sharding import maybe_shard_with_logical
+from MaxText.sharding import maybe_shard_with_logical, create_sharding
 from MaxText.inference import page_manager
 from MaxText.common_types import MODEL_MODE_PREFILL
 
@@ -75,7 +75,7 @@ def self_attention_with_norm(
       mesh=mesh,
       shard_mode=cfg.shard_mode,
   )
-  lnx_sharding = NamedSharding(mesh, nn.logical_to_mesh_axes(logical_axis_names))
+  lnx_sharding = create_sharding(mesh, logical_axis_names)
   lnx = _maybe_shard_with_logical(lnx, logical_axis_names)
 
   attention_layer = attention_mla.mla_as_linen(
@@ -189,8 +189,8 @@ def __call__(
       inputs = inputs[0]
 
     _maybe_shard_with_logical = partial(maybe_shard_with_logical, mesh=self.mesh, shard_mode=self.config.shard_mode)
-    lnx_out_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(logical_axis_names))
-    mlp_intermediate_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(mlp_logical_axis_names))
+    lnx_out_sharding = create_sharding(self.mesh, logical_axis_names)
+    mlp_intermediate_sharding = create_sharding(self.mesh, mlp_logical_axis_names)
     inputs = _maybe_shard_with_logical(inputs, logical_axis_names)
     inputs = checkpoint_name(inputs, "decoder_layer_input")
 
@@ -273,8 +273,8 @@ def __call__(
       inputs = inputs[0]
 
     _maybe_shard_with_logical = partial(maybe_shard_with_logical, mesh=self.mesh, shard_mode=self.config.shard_mode)
-    lnx_out_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(logical_axis_names))
-    lnx_intermediate_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(mlp_logical_axis_names))
+    lnx_out_sharding = create_sharding(self.mesh, logical_axis_names)
+    lnx_intermediate_sharding = create_sharding(self.mesh, mlp_logical_axis_names)
     inputs = _maybe_shard_with_logical(inputs, logical_axis_names)
     inputs = checkpoint_name(inputs, "decoder_layer_input")
 

src/MaxText/layers/embeddings.py

@@ -22,11 +22,11 @@
 import jax.numpy as jnp
 from jax.sharding import Mesh, NamedSharding
 
-from flax import linen as nn
 from flax import nnx
 
 from MaxText import max_logging
 from MaxText import max_utils
+from MaxText.sharding import logical_to_mesh_axes, create_sharding
 from MaxText.common_types import ShardMode, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, Array, Config, DType
 from MaxText.layers import nnx_wrappers
 from MaxText.layers.initializers import Initializer, default_embed_init, variable_to_logically_partitioned
@@ -169,7 +169,7 @@ def __call__(self, inputs: Array, model_mode: str = MODEL_MODE_TRAIN) -> Array:
             "activation_embed",
         )
     )
-    out_pspec = nn.logical_to_mesh_axes(output_axis_names)
+    out_pspec = logical_to_mesh_axes(output_axis_names, self.mesh)
 
     out_sharding = NamedSharding(self.mesh, out_pspec) if self.config.shard_mode == ShardMode.EXPLICIT else None
 
@@ -751,7 +751,7 @@ def __init__(
     self.attention_scaling = attention_scaling
 
     self.freqs_sharding = (
-        NamedSharding(mesh, nn.logical_to_mesh_axes(("activation_batch", "activation_length_no_exp", "q_heads")))
+        create_sharding(mesh, ("activation_batch", "activation_length_no_exp", "q_heads"))
         if shard_mode == ShardMode.EXPLICIT
         else None
     )
@@ -877,7 +877,7 @@ def __call__(self, inputs: Array, position: None | Array = None) -> Array:
     inputs_complex = first_half + 1j * second_half  # shape: [B, S, N, half_dim]
     # Apply the rotary transformation via complex multiplication.
     rotated_sharding = (
-        NamedSharding(self.mesh, nn.logical_to_mesh_axes(("activation_batch", "activation_length_no_exp", None, None)))
+        create_sharding(self.mesh, ("activation_batch", "activation_length_no_exp", None, None))
         if self.shard_mode == ShardMode.EXPLICIT
         else None
     )

src/MaxText/layers/llama2.py

@@ -19,15 +19,14 @@
 import functools
 import jax.numpy as jnp
 from jax.ad_checkpoint import checkpoint_name
-from jax.sharding import Mesh, NamedSharding
+from jax.sharding import Mesh
 
-from flax import linen as nn
 from flax import nnx
 
 from MaxText.inference import page_manager
 from MaxText.common_types import Config
 from MaxText import max_utils
-from MaxText.sharding import maybe_shard_with_logical
+from MaxText.sharding import maybe_shard_with_logical, create_sharding
 from MaxText.layers.linears import Dropout, MlpBlock
 from MaxText.layers import initializers
 from MaxText.layers import nnx_wrappers
@@ -157,7 +156,7 @@ def __call__(
       inputs = inputs[0]
     inputs = self._maybe_shard_with_logical(inputs, self.activation_axis_names)
     inputs = checkpoint_name(inputs, "decoder_layer_input")
-    lnx_sharding = NamedSharding(self.mesh, nn.logical_to_mesh_axes(self.activation_axis_names))
+    lnx_sharding = create_sharding(self.mesh, self.activation_axis_names)
     lnx = self.pre_self_attention_layer_norm(inputs, out_sharding=lnx_sharding)
     lnx = self._maybe_shard_with_logical(lnx, self.activation_axis_names)
 
@@ -185,9 +184,8 @@ def __call__(
     hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names)
 
     # MLP block.
-    mlp_intermediate_sharding = NamedSharding(
-        self.mesh,
-        nn.logical_to_mesh_axes(("activation_batch", "activation_length_no_exp", "activation_mlp")),
+    mlp_intermediate_sharding = create_sharding(
+        self.mesh, ("activation_batch", "activation_length_no_exp", "activation_mlp")
     )
     mlp_lnx = self.mlp(
         hidden_states,