Merge pull request #2926 from AI-Hypercomputer:chengnuojin-debug-sharding

PiperOrigin-RevId: 855366909

Author: Google-ML-Automation

src/MaxText/data_loader.py

@@ -20,7 +20,7 @@
 from jax.experimental import checkify
 
 from MaxText import exceptions
-from MaxText.sharding import get_input_data_sharding, maybe_shard_with_name
+from MaxText.sharding import get_input_data_sharding
 from MaxText.utils.goodput_utils import (
     GoodputEvent,
     maybe_record_goodput,
@@ -70,10 +70,9 @@ def load_next_batch_pre_sharding(self):
 
   def load_next_batch(self, *args, **kwargs):
     """Loads the next batch with sharding hint"""
-    return maybe_shard_with_name(
+    return jax.device_put(
         self.load_next_batch_pre_sharding(),
         self.input_data_shardings,
-        self.config.shard_mode,
     )
 
   def check_example_batch(self):
@@ -154,7 +153,7 @@ def _slice(data):
       self.buffer_start = slice_end
       output = jax.tree.map(_slice, self.batch_buffer)
     self.rampup_active = rampup_manager.update()
-    return maybe_shard_with_name(output, self.input_data_shardings, self.config.shard_mode)
+    return jax.device_put(output, self.input_data_shardings)
 
 
 def create_dataloader(config, mesh, data_iterator, goodput_recorder, rampup_manager):

src/MaxText/gradient_accumulation.py

@@ -65,7 +65,7 @@ def gradient_accumulation_loss_and_grad(
 
   def _maybe_shard_with_name(inputs, sharding_names):
     """Wrapper of maybe_shard_with_name with fixed shard_mode"""
-    return maybe_shard_with_name(inputs, sharding_names, config.shard_mode)
+    return maybe_shard_with_name(inputs, sharding_names, config.shard_mode, debug_sharding=config.debug_sharding)
 
   # For more efficient DP/ZeRO-1 + GA
   if config.shard_mode == ShardMode.EXPLICIT and config.ici_data_parallelism > 1:

src/MaxText/layers/attention_op.py

@@ -1198,17 +1198,10 @@ def wrap_splash_kernel(single_head_mask, shard_head_size=1):
         segment_axis_names_splash_kernel = self._logical_to_mesh_axes((Q_LENGTH,))
       else:
         segment_axis_names_splash_kernel = self._logical_to_mesh_axes((Q_LENGTH_NO_EXP,))
-    elif (
-        self.config.use_jax_splash
-        and self.config.expert_shard_attention_option == EP_AS_FSDP
-    ):
+    elif self.config.use_jax_splash and self.config.expert_shard_attention_option == EP_AS_FSDP:
       if self.config.use_max_logit_estimate > 0:
-        sa_config = dataclasses.replace(
-            sa_config, max_logit_const=self.config.use_max_logit_estimate
-        )
-      segment_axis_names_splash_kernel = nn.logical_to_mesh_axes((
-          Q_LENGTH_NO_EXP,
-      ))
+        sa_config = dataclasses.replace(sa_config, max_logit_const=self.config.use_max_logit_estimate)
+      segment_axis_names_splash_kernel = nn.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])
@@ -1327,7 +1320,13 @@ def _maybe_shard_with_pspec(inputs, pspec: jax.sharding.PartitionSpec | None):
       if pspec is None:
         return None
       sharding = NamedSharding(self.mesh, pspec)
-      return maybe_shard_with_name(inputs, sharding, shard_mode=self.config.shard_mode)
+      return maybe_shard_with_name(
+          inputs,
+          sharding,
+          shard_mode=self.config.shard_mode,
+          debug_sharding=self.config.debug_sharding,
+          extra_stack_level=1,
+      )
 
     query = _maybe_shard_with_pspec(query, axis_names_q)
     key = _maybe_shard_with_pspec(key, axis_names_kv)

src/MaxText/layers/attentions.py

@@ -525,6 +525,7 @@ def __init__(
         maybe_shard_with_logical,
         mesh=mesh,
         shard_mode=config.shard_mode,
+        debug_sharding=config.debug_sharding,
     )
 
   def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> None:

src/MaxText/layers/decoders.py

@@ -98,6 +98,7 @@ def __call__(
         sharding.maybe_shard_with_logical,
         mesh=mesh,
         shard_mode=cfg.shard_mode,
+        debug_sharding=cfg.debug_sharding,
     )
 
     if self.model_mode == MODEL_MODE_PREFILL:

src/MaxText/layers/deepseek.py

@@ -129,7 +129,11 @@ def mlp_op(self, x, deterministic, *args, **kwargs):
 
   def with_logical_constraint(self, x):
     return maybe_shard_with_logical(
-        x, logical_axes=self.logical_axis_names, mesh=self.mesh, shard_mode=self.config.shard_mode
+        x,
+        logical_axes=self.logical_axis_names,
+        mesh=self.mesh,
+        shard_mode=self.config.shard_mode,
+        debug_sharding=self.config.debug_sharding,
     )
 
   def dropout_op(self, x, deterministic):

src/MaxText/layers/deepseek_batchsplit.py

@@ -179,7 +179,7 @@ def mlp_logical_axis_names(self):
   def with_logical_constraint(self, x):
     return maybe_shard_with_logical(
       x, logical_axes=self.logical_axis_names,
-      mesh=self.mesh, shard_mode=self.config.shard_mode
+      mesh=self.mesh, shard_mode=self.config.shard_mode, debug_sharding=self.config.debug_sharding,
     )
 
   def pre_attention_norm_op(self, x):

src/MaxText/layers/linears.py

@@ -459,6 +459,7 @@ def __init__(
         maybe_shard_with_logical,
         mesh=mesh,
         shard_mode=config.shard_mode,
+        debug_sharding=config.debug_sharding,
     )
 
   def get_norm_layer(self, num_features: int):

src/MaxText/layers/llama2.py

@@ -134,6 +134,7 @@ def __init__(
         maybe_shard_with_logical,
         mesh=self.mesh,
         shard_mode=config.shard_mode,
+        debug_sharding=config.debug_sharding,
     )
 
   def __call__(

src/MaxText/layers/moe.py

@@ -444,7 +444,13 @@ def __init__(
       self.wo_bias = None
 
   def _maybe_shard_with_logical(self, inputs, logical_name):
-    return maybe_shard_with_logical(inputs, logical_name, mesh=self.mesh, shard_mode=self.config.shard_mode)
+    return maybe_shard_with_logical(
+        inputs,
+        logical_name,
+        mesh=self.mesh,
+        shard_mode=self.config.shard_mode,
+        debug_sharding=self.config.debug_sharding,
+    )
 
   def _logical_to_mesh_axes(self, logical_name):
     return logical_to_mesh_axes(logical_name, mesh=self.mesh, rules=self.config.logical_axis_rules)