Merge pull request #2562 from AI-Hypercomputer:chengnuojin-move-sharding

PiperOrigin-RevId: 826212582

Author: Google-ML-Automation

src/MaxText/data_loader.py

@@ -20,7 +20,7 @@
 from jax.experimental import checkify
 
 from MaxText import exceptions
-from MaxText import maxtext_utils
+from MaxText import sharding
 from MaxText.utils.goodput_utils import (
     GoodputEvent,
     maybe_record_goodput,
@@ -37,7 +37,7 @@ def __init__(self, config, mesh, data_iterator, goodput_recorder):
     self.goodput_recorder = goodput_recorder
     self.data_iterator = data_iterator
     self.last_batch = None
-    self.input_data_shardings = maxtext_utils.get_input_data_sharding(config, mesh)
+    self.input_data_shardings = sharding.get_input_data_sharding(config, mesh)
 
   def load_next_batch(self):
     """Loads the next batch. Can keep reusing the same batch for performance reasons."""
@@ -48,7 +48,7 @@ def load_next_batch(self):
         else:
           example_batch = next(self.data_iterator)
         # Reshard data from loaded sharding to performant activation sharding
-        self.last_batch = maxtext_utils.maybe_shard_with_name(
+        self.last_batch = sharding.maybe_shard_with_name(
             example_batch,
             self.input_data_shardings,
             self.config.shard_mode,

src/MaxText/experimental/rl/grpo_trainer.py

@@ -72,6 +72,7 @@
 from MaxText import max_logging
 from MaxText import max_utils
 from MaxText import maxtext_utils
+from MaxText import sharding
 from MaxText import train_utils
 from MaxText import profiler
 from MaxText import pyconfig
@@ -566,7 +567,7 @@ def setup_train_loop(
   )[2]
   if not config.using_pipeline_parallelism:
     # The vocab tensor(s) of shape [vocab, embed] (and transpose) are not sharded by stage
-    maxtext_utils.assert_params_sufficiently_sharded(state.params, mesh, config.sharding_tolerance)
+    sharding.assert_params_sufficiently_sharded(state.params, mesh, config.sharding_tolerance)
 
   return (
       init_rng,
@@ -688,7 +689,7 @@ def train_loop(config, config_inference, recorder, state=None):
       config, model, mesh, state, state_mesh_shardings, train_step, eval_step, eval_data_iterator
   )
 
-  data_sharding = maxtext_utils.get_input_data_sharding(config, mesh)
+  data_sharding = sharding.get_input_data_sharding(config, mesh)
 
   inference_engine = offline_engine.OfflineEngine(
       config=config_inference,

src/MaxText/gradient_accumulation.py

@@ -19,7 +19,7 @@
 from jax.sharding import NamedSharding
 
 from MaxText.common_types import ShardMode
-from MaxText.maxtext_utils import maybe_shard_with_name
+from MaxText.sharding import maybe_shard_with_name
 
 
 def gradient_accumulation_loss_and_grad(

src/MaxText/layers/attention_op.py

@@ -43,7 +43,7 @@
 
 
 from MaxText import max_utils
-from MaxText.maxtext_utils import maybe_shard_with_name
+from MaxText.sharding import maybe_shard_with_name
 from MaxText.common_types import (
     DEFAULT_MASK_VALUE,
     BATCH,

src/MaxText/layers/attentions.py

@@ -53,7 +53,7 @@
     EP_AS_CONTEXT,
     AttentionType,
 )
-from MaxText.maxtext_utils import maybe_shard_with_logical
+from MaxText.sharding import maybe_shard_with_logical
 from MaxText.inference import kvcache
 from MaxText.inference import page_manager
 from MaxText.inference import paged_attention

src/MaxText/layers/decoders.py

@@ -38,6 +38,7 @@
 from MaxText.layers import pipeline
 from MaxText import maxtext_utils
 from MaxText import multimodal_utils
+from MaxText import sharding
 from MaxText.layers.attentions import attention_as_linen
 from MaxText.layers.normalizations import rms_norm
 from MaxText.layers.embeddings import attend_on_embedding, embed_as_linen, positional_embedding_as_linen
@@ -90,7 +91,7 @@ def __call__(
     cfg = self.config
     mesh = self.mesh
     _maybe_shard_with_logical = functools.partial(
-        maxtext_utils.maybe_shard_with_logical,
+        sharding.maybe_shard_with_logical,
         mesh=mesh,
         shard_mode=cfg.shard_mode,
     )
@@ -722,7 +723,7 @@ def __call__(
         moe_layer = RemattedBlockLayers[1]
         num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers
         num_moe_layers_outside_pp = num_moe_layers - self.config.pipeline_parallel_layers
-        logical_axis_rules_pp_as_dp = maxtext_utils.logical_axis_rules_pp_act_as_dp(self.config.logical_axis_rules)
+        logical_axis_rules_pp_as_dp = sharding.logical_axis_rules_pp_act_as_dp(self.config.logical_axis_rules)
         # We chose not to pipeline the dense layers, only sparse for SPMD.
         with self.mesh, nn.partitioning.axis_rules(logical_axis_rules_pp_as_dp):
           y, _ = self.scan_decoder_layers(
@@ -749,7 +750,7 @@ def __call__(
         y = self.pipeline_module(y, *broadcast_args, partition_spec=partition_spec)
         remaining_layers = self.config.num_decoder_layers - self.config.pipeline_parallel_layers
         if remaining_layers > 0:
-          logical_axis_rules_pp_as_dp = maxtext_utils.logical_axis_rules_pp_act_as_dp(self.config.logical_axis_rules)
+          logical_axis_rules_pp_as_dp = sharding.logical_axis_rules_pp_act_as_dp(self.config.logical_axis_rules)
           with self.mesh, nn.partitioning.axis_rules(logical_axis_rules_pp_as_dp):
             y, _ = self.scan_decoder_layers(
                 cfg,

src/MaxText/layers/linears.py

@@ -31,7 +31,7 @@
 
 from MaxText import max_logging
 from MaxText import max_utils
-from MaxText.maxtext_utils import maybe_shard_with_logical
+from MaxText.sharding import maybe_shard_with_logical
 from MaxText.common_types import DecoderBlockType, ShardMode, DType, Array, Config
 from MaxText.common_types import MODEL_MODE_TRAIN, MODEL_MODE_PREFILL, EP_AS_CONTEXT
 from MaxText.layers import nnx_wrappers, quantizations

src/MaxText/layers/llama2.py

@@ -27,7 +27,7 @@
 from MaxText.inference import page_manager
 from MaxText.common_types import Config
 from MaxText import max_utils
-from MaxText.maxtext_utils import maybe_shard_with_logical
+from MaxText.sharding import maybe_shard_with_logical
 from MaxText.layers.linears import Dropout, MlpBlock
 from MaxText.layers import initializers
 from MaxText.layers import nnx_wrappers

src/MaxText/layers/models.py

@@ -34,7 +34,7 @@
 from MaxText.layers.encoders import VisionEncoder
 from MaxText.layers.quantizations import AqtQuantization as Quant
 from MaxText.layers.multi_token_prediction import MultiTokenPredictionBlock
-from MaxText.maxtext_utils import all_gather_over_fsdp
+from MaxText.sharding import all_gather_over_fsdp
 
 # ------------------------------------------------------------------------------
 # The network: Transformer Definitions

src/MaxText/layers/pipeline.py

@@ -28,7 +28,7 @@
 from flax import linen as nn
 
 from MaxText.common_types import Config, MODEL_MODE_TRAIN, EP_AS_CONTEXT
-from MaxText.maxtext_utils import all_gather_over_fsdp
+from MaxText.sharding import all_gather_over_fsdp
 
 
 class Pipeline(nn.Module):