Merge pull request #4568 from AI-Hypercomputer:chengnuojin-train-eval-mesh

PiperOrigin-RevId: 953055860

Author: Google-ML-Automation

src/maxtext/layers/attention_op.py

@@ -75,7 +75,7 @@
 from maxtext.layers.initializers import variable_to_logically_partitioned
 from maxtext.layers.quantizations import AqtQuantization as Quant
 from maxtext.utils import max_utils
-from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_pspec, get_logical_axis_rules_from_config
+from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_pspec, get_logical_axis_rules
 import numpy as np
 # pylint: disable=line-too-long, g-doc-args, g-doc-return-or-yield, bad-continuation, g-inconsistent-quotes
 # pytype: disable=attribute-error
@@ -633,7 +633,7 @@ def maybe_create_nnx(einsum, *args):
       self.AqtEinsum_3 = jnp.einsum
 
   def _logical_to_mesh_axes(self, logical_name):
-    logical_rules = get_logical_axis_rules_from_config(self.config)
+    logical_rules = get_logical_axis_rules()
     return logical_to_mesh_axes(logical_name, mesh=self.mesh, rules=logical_rules)
 
   def check_attention_inputs(self, query: Array, key: Array | KVTensor, value: Array | KVTensor) -> None:

src/maxtext/layers/moe.py

@@ -43,7 +43,7 @@
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
 from maxtext.utils.sharding import create_sharding, maybe_shard_with_logical, maybe_shard_with_pspec
-from maxtext.utils.sharding import logical_to_mesh_axes, remove_expert_from_partition_spec, get_logical_axis_rules_from_config
+from maxtext.utils.sharding import logical_to_mesh_axes, remove_expert_from_partition_spec, get_logical_axis_rules
 import numpy as np
 import qwix
 from qwix.contrib.sparsity import sparsity_module
@@ -644,7 +644,7 @@ def _maybe_shard_with_logical(self, inputs, logical_name):
     )
 
   def _logical_to_mesh_axes(self, logical_name):
-    logical_rules = get_logical_axis_rules_from_config(self.config)
+    logical_rules = get_logical_axis_rules()
     return logical_to_mesh_axes(logical_name, mesh=self.mesh, rules=logical_rules)
 
   def _maybe_shard_with_pspec(self, inputs, pspec: jax.sharding.PartitionSpec | None):

src/maxtext/layers/multi_token_prediction.py

@@ -177,7 +177,7 @@ def __call__(
         ("activation_batch", "activation_length", "activation_embed"),
         self.mesh,
         self.config.shard_mode,
-        self.config.logical_axis_rules,
+        sharding.get_logical_axis_rules(),
     )
 
     embedding_norm = self.embedding_norm(target_token_embedding)
@@ -200,7 +200,7 @@ def extract_fn(x):
               x.sharding_names,
               self.mesh,
               shard_mode=self.config.shard_mode,
-              rules=self.config.logical_axis_rules,
+              rules=sharding.get_logical_axis_rules(),
           )
         return x
 

src/maxtext/layers/pipeline.py

@@ -38,6 +38,7 @@
     create_sharding,
     logical_to_mesh_axes,
     logical_to_mesh,
+    get_logical_axis_rules,
 )
 from maxtext.utils import pipeline_utils
 
@@ -61,21 +62,21 @@ def _setup_pipeline_attributes(self):
     self.spmd_axis_name = "stage" if self.config.shard_mode == ShardMode.AUTO else None
 
     self.stages_in_logical = ("activation_stage", self.batch_axis_name, self.seq_len_axis_name, "activation_embed")
-    self.stages_in_spec = logical_to_mesh_axes(self.stages_in_logical, self.mesh, rules=self.config.logical_axis_rules)
+    self.stages_in_spec = logical_to_mesh_axes(self.stages_in_logical, self.mesh, rules=get_logical_axis_rules())
     self.stages_in_sharding = (
         NamedSharding(self.mesh, self.stages_in_spec) if self.config.shard_mode == ShardMode.EXPLICIT else None
     )
 
     self.state_io_logical = ("activation_stage", None, self.batch_axis_name, self.seq_len_axis_name, "activation_embed")
-    self.state_io_spec = logical_to_mesh_axes(self.state_io_logical, self.mesh, rules=self.config.logical_axis_rules)
+    self.state_io_spec = logical_to_mesh_axes(self.state_io_logical, self.mesh, rules=get_logical_axis_rules())
     self.state_io_sharding = (
         NamedSharding(self.mesh, self.state_io_spec) if self.config.shard_mode == ShardMode.EXPLICIT else None
     )
     self.input_sharding = (
         create_sharding(
             self.mesh,
             (None, self.batch_axis_name, self.seq_len_axis_name, "activation_embed"),
-            rules=self.config.logical_axis_rules,
+            rules=get_logical_axis_rules(),
         )
         if self.config.shard_mode == ShardMode.EXPLICIT
         else None
@@ -84,7 +85,7 @@ def _setup_pipeline_attributes(self):
         create_sharding(
             self.mesh,
             (self.batch_axis_name, self.seq_len_axis_name, "activation_embed"),
-            rules=self.config.logical_axis_rules,
+            rules=get_logical_axis_rules(),
         )
         if self.config.shard_mode == ShardMode.EXPLICIT
         else None
@@ -114,7 +115,7 @@ def _maybe_shard_with_logical(self, inputs, logical_axes):
         logical_axes,
         shard_mode=self.config.shard_mode,
         mesh=self.mesh,
-        rules=self.config.logical_axis_rules,
+        rules=get_logical_axis_rules(),
         debug_sharding=self.config.debug_sharding,
         extra_stack_level=1,
     )

src/maxtext/models/deepseek.py

@@ -42,7 +42,7 @@
 from maxtext.utils import max_utils
 from maxtext.utils.sharding import create_sharding
 from maxtext.utils.sharding import maybe_shard_with_logical
-from maxtext.utils.sharding import get_logical_axis_rules_from_config
+from maxtext.utils.sharding import get_logical_axis_rules
 
 import transformers
 
@@ -79,9 +79,9 @@ def __init__(
     batch_size, sequence_length = max_utils.get_batch_seq_len_for_mode(self.config, self.model_mode)
     self.dummy_inputs_shape = (batch_size, sequence_length, self.config.emb_dim)
 
-    self.out_sharding = create_sharding(self.mesh, self.logical_axis_names, rules=self.config.logical_axis_rules)
+    self.out_sharding = create_sharding(self.mesh, self.logical_axis_names, rules=get_logical_axis_rules())
     self.mlp_intermediate_sharding = create_sharding(
-        self.mesh, self.mlp_logical_axis_names, rules=self.config.logical_axis_rules
+        self.mesh, self.mlp_logical_axis_names, rules=get_logical_axis_rules()
     )
 
     self.pre_self_attention_layer_norm = RMSNorm(
@@ -190,7 +190,7 @@ def with_logical_constraint(self, x):
         shard_mode=self.config.shard_mode,
         debug_sharding=self.config.debug_sharding,
         extra_stack_level=1,
-        rules=get_logical_axis_rules_from_config(self.config),
+        rules=get_logical_axis_rules(),
     )
 
   def dropout_op(self, x, deterministic):
@@ -527,7 +527,7 @@ def extract_fn(x):
               x.sharding_names,
               self.mesh,
               shard_mode=self.config.shard_mode,
-              rules=self.config.logical_axis_rules,
+              rules=get_logical_axis_rules(),
           )
         return x
 

src/maxtext/utils/sharding.py

@@ -56,17 +56,6 @@ def get_logical_axis_rules():
   return flax_get_logical_axis_rules()
 
 
-def get_logical_axis_rules_from_config(config):
-  """Get the logical axis rules from the config.
-  When we use pipeline parallelism or in eval step, we may use
-  a different logical axis rule than the one in config.
-  Plan to deprecate (b/536927795) and use get_logical_axis_rules instead.
-  """
-  if config.using_pipeline_parallelism or config.eval_interval != -1:
-    return None
-  return config.logical_axis_rules
-
-
 def _get_sharding_desc(inputs, extra_stack_level):
   """Get the inputs sharding description using inspect module"""
   frame = inspect.currentframe()

src/maxtext/utils/vocabulary_tiling.py

@@ -25,6 +25,7 @@
     maybe_shard_with_name,
     all_gather_over_fsdp,
     create_sharding,
+    get_logical_axis_rules,
 )
 from maxtext.common.common_types import ShardMode
 from maxtext.utils import max_utils
@@ -128,7 +129,7 @@ def _reshape(inputs, out_shape, out_sharding):
   labels = _maybe_shard_with_name(labels, label_spec)
   segmentation = _maybe_shard_with_name(segmentation, label_spec)
   # TODO (chengnuojin) all gather only embedding table instead of all params after NNX module is enabled
-  gathered_params = all_gather_over_fsdp(params, param_spec, model.mesh, config.logical_axis_rules, config.shard_mode)
+  gathered_params = all_gather_over_fsdp(params, param_spec, model.mesh, get_logical_axis_rules(), config.shard_mode)
 
   # Customized forward and backward maps for the embedding tiling
   @jax.custom_vjp
@@ -354,7 +355,7 @@ def _reshape(inputs, out_shape, out_sharding):
       head_params,
       nnx.get_partition_spec(head_params),
       model.mesh,
-      config.logical_axis_rules,
+      get_logical_axis_rules(),
       config.shard_mode,
   )
 

tests/unit/attention_test.py

@@ -1190,22 +1190,23 @@ def test_tpu_flash_attention_ring_context_parallel(self, context_parallel_load_b
     )
     generic_state = nnx.state(attention_as_mha_generic)
 
-    attention_as_mha_flash_cp = Attention(
-        config=cfg_cp,
-        num_query_heads=cfg_cp.num_query_heads,
-        num_kv_heads=cfg_cp.num_kv_heads,
-        head_dim=cfg_cp.head_dim,
-        max_target_length=cfg_cp.max_target_length,
-        max_prefill_predict_length=cfg_cp.max_prefill_predict_length,
-        inputs_q_shape=lnx.shape,
-        inputs_kv_shape=lnx.shape,
-        mesh=mesh_cp,
-        attention_kernel="flash",
-        dtype=cfg_cp.dtype,
-        dropout_rate=cfg_cp.dropout_rate,
-        model_mode=MODEL_MODE_PREFILL,
-        rngs=self.nnx_rng,
-    )
+    with nn_partitioning.axis_rules(cfg_cp.logical_axis_rules):
+      attention_as_mha_flash_cp = Attention(
+          config=cfg_cp,
+          num_query_heads=cfg_cp.num_query_heads,
+          num_kv_heads=cfg_cp.num_kv_heads,
+          head_dim=cfg_cp.head_dim,
+          max_target_length=cfg_cp.max_target_length,
+          max_prefill_predict_length=cfg_cp.max_prefill_predict_length,
+          inputs_q_shape=lnx.shape,
+          inputs_kv_shape=lnx.shape,
+          mesh=mesh_cp,
+          attention_kernel="flash",
+          dtype=cfg_cp.dtype,
+          dropout_rate=cfg_cp.dropout_rate,
+          model_mode=MODEL_MODE_PREFILL,
+          rngs=self.nnx_rng,
+      )
     nnx.update(attention_as_mha_flash_cp, generic_state)
 
     mha_generic_flash_cp_output = attention_test_util.forward_with_context_expert_parallelism(
@@ -1266,22 +1267,23 @@ def test_tpu_flash_attention_ring_context_parallel_grad(self, context_parallel_l
     )
     generic_state = nnx.state(attention_as_mha_generic)
 
-    attention_as_mha_flash_cp = Attention(
-        config=cfg_cp,
-        num_query_heads=cfg_cp.num_query_heads,
-        num_kv_heads=cfg_cp.num_kv_heads,
-        head_dim=cfg_cp.head_dim,
-        max_target_length=cfg_cp.max_target_length,
-        max_prefill_predict_length=cfg_cp.max_prefill_predict_length,
-        inputs_q_shape=lnx.shape,
-        inputs_kv_shape=lnx.shape,
-        mesh=mesh_cp,
-        attention_kernel="flash",
-        dtype=cfg_cp.dtype,
-        dropout_rate=cfg_cp.dropout_rate,
-        model_mode=MODEL_MODE_PREFILL,
-        rngs=self.nnx_rng,
-    )
+    with nn_partitioning.axis_rules(cfg_cp.logical_axis_rules):
+      attention_as_mha_flash_cp = Attention(
+          config=cfg_cp,
+          num_query_heads=cfg_cp.num_query_heads,
+          num_kv_heads=cfg_cp.num_kv_heads,
+          head_dim=cfg_cp.head_dim,
+          max_target_length=cfg_cp.max_target_length,
+          max_prefill_predict_length=cfg_cp.max_prefill_predict_length,
+          inputs_q_shape=lnx.shape,
+          inputs_kv_shape=lnx.shape,
+          mesh=mesh_cp,
+          attention_kernel="flash",
+          dtype=cfg_cp.dtype,
+          dropout_rate=cfg_cp.dropout_rate,
+          model_mode=MODEL_MODE_PREFILL,
+          rngs=self.nnx_rng,
+      )
     nnx.update(attention_as_mha_flash_cp, generic_state)
 
     def generic_loss(lnx):

tests/unit/moe_test.py

@@ -1506,7 +1506,8 @@ def test_fused_vs_sparse_softmax(self):
     copy_weights(self.dense_model, sparse_model)
 
     inputs = self._inputs()
-    sparse_out, _, _ = sparse_model(inputs)
+    with nn_partitioning.axis_rules(sparse_cfg.logical_axis_rules):
+      sparse_out, _, _ = sparse_model(inputs)
     fused_out, lb_loss, bias_updates = self.fused_model(inputs)
 
     np.testing.assert_allclose(
@@ -1647,7 +1648,8 @@ def test_prefused_vs_sparse_softmax(self):
     copy_weights_prefused(self.dense_model, prefused_model)
 
     inputs = self._inputs()
-    sparse_out, _, _ = sparse_model(inputs)
+    with nn_partitioning.axis_rules(sparse_cfg.logical_axis_rules):
+      sparse_out, _, _ = sparse_model(inputs)
     prefused_out, lb_loss, bias_updates = prefused_model(inputs)
 
     np.testing.assert_allclose(