update shardmap in maxtext

Author: NuojCheng

pedagogical_examples/shmap_collective_matmul.py

@@ -29,7 +29,6 @@
 from jax.experimental.pjit import pjit
 from jax.sharding import PartitionSpec as P
 from jax.sharding import Mesh
-from jax.experimental.shard_map import shard_map
 
 MESH_DATA_AXIS = "dp"
 MESH_FSDP_AXIS = "fsdp"
@@ -103,14 +102,14 @@ def matmul(activations, weights):  # pylint: disable=redefined-outer-name
 
 
 @partial(
-    shard_map,
+    jax.shard_map,
     mesh=global_mesh,
     in_specs=(
         P(MESH_FSDP_AXIS, MESH_TENSOR_AXIS, None),
         P(MESH_FSDP_AXIS, MESH_TENSOR_AXIS, None),
     ),
     out_specs=P(MESH_FSDP_AXIS, None, MESH_TENSOR_AXIS, None),
-    check_rep=False,
+    check_vma=False,
 )
 def collective_matmul(activations, weights):  # pylint: disable=redefined-outer-name
   """Collective matrix multiply"""

src/MaxText/inference/paged_attention.py

@@ -21,7 +21,6 @@
 
 import jax
 import jax.numpy as jnp
-from jax.experimental.shard_map import shard_map
 from jax.experimental.pallas.ops.tpu.paged_attention import paged_attention_kernel
 from jax.sharding import PartitionSpec as P
 from jax.sharding import Mesh
@@ -327,7 +326,7 @@ def paged_attention_v1_decode(
     q_pspec = nn.logical_to_mesh_axes((None, None, "paged_kv_heads", None))
 
     @functools.partial(
-        shard_map,
+        jax.shard_map,
         mesh=self.mesh,
         in_specs=(
             q_pspec,
@@ -338,7 +337,7 @@ def paged_attention_v1_decode(
             None,
         ),
         out_specs=q_pspec,
-        check_rep=False,
+        check_vma=False,
     )
     def wrap_paged_attention(q, k_pages, v_pages, lengths, page_indices, pages_per_compute_block):
       q = jnp.squeeze(q, axis=1)

src/MaxText/layers/attention_op.py

@@ -28,7 +28,6 @@
 from jax.experimental.pallas.ops.gpu import decode_attention as gpu_pallas_decode_attention
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_kernel
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask
-from jax.experimental.shard_map import shard_map
 from jax.sharding import Mesh
 import jax
 import jax.numpy as jnp
@@ -946,11 +945,11 @@ def gpu_ragged_attention(self, q: Array, k: Array | KVTensor, v: Array | KVTenso
     bn = nn.logical_to_mesh_axes((CACHE_BATCH, CACHE_HEADS))
 
     @functools.partial(
-        shard_map,
+        jax.shard_map,
         mesh=self.mesh,
         in_specs=(bnd, bsnd, bsnd, b, None),
         out_specs=(bnd, bn, bn),
-        check_rep=False,
+        check_vma=False,
     )
     def wrap_ragged_attention(
         q: Array, k: Array, v: Array, lengths: Array, block_size: int
@@ -1000,7 +999,7 @@ def tpu_ragged_attention(
     bsnd = nn.logical_to_mesh_axes(self.cache_logical_axis_names)
 
     @functools.partial(
-        shard_map,
+        jax.shard_map,
         mesh=self.mesh,
         in_specs=(
             bsnd,
@@ -1010,7 +1009,7 @@ def tpu_ragged_attention(
             None,
         ),
         out_specs=bsnd,
-        check_rep=False,
+        check_vma=False,
     )
     def wrap_ragged_attention(query, key, value, lengths, block_size):
       if query.shape[-2] == key.shape[-2]:
@@ -1162,7 +1161,7 @@ def wrap_splash_kernel(multi_head_mask, shard_head_size=1):
     #  'segment_axis_names_kv' maps to ['activation_kv_length', []] meaning that K and V are not sharded
     # splash_kernel is sharded over (HEAD, LENGTH)
     @functools.partial(
-        shard_map,
+        jax.shard_map,
         mesh=self.mesh,
         in_specs=(
             axis_names_q,
@@ -1176,7 +1175,7 @@ def wrap_splash_kernel(multi_head_mask, shard_head_size=1):
             sink_axis_names,  # sharding align with query heads
         ),
         out_specs=axis_names_q,
-        check_rep=False,
+        check_vma=False,
     )
     def wrap_flash_attention(
         query,

src/MaxText/layers/moe.py

@@ -25,7 +25,6 @@
 from flax import nnx
 import jax
 from jax import ad_checkpoint as adc
-from jax.experimental import shard_map
 from jax.experimental import xla_metadata
 import jax.numpy as jnp
 import numpy as np
@@ -908,7 +907,7 @@ def gmm(inputs, kernel, tiling, group_sizes, expert_assignments):
       wo_pspec = aqt.partition_spec(wo_pspec, (1,), wo_kernel.dtype, use_bias=False)
 
     @functools.partial(
-        shard_map.shard_map,
+        jax.shard_map,
         mesh=self.mesh,
         in_specs=(
             input_partition_pspec,
@@ -923,7 +922,7 @@ def gmm(inputs, kernel, tiling, group_sizes, expert_assignments):
             None,
         ),
         out_specs=(nn.logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", "activation_embed"))),
-        check_rep=False,
+        check_vma=False,
     )
     def wrapper(x, logits, pre_bias_logits, w0, w1, wo, w0_bias, w1_bias, wo_bias, rngs):
       batch_size, sequence_length, _ = x.shape