deprecate get_logical_axis_rule_from_config

Author: NuojCheng

src/maxtext/configs/base.yml

@@ -1069,7 +1069,7 @@ xprof_e2e_enable_fw_power_level_event: false
 xprof_e2e_enable_fw_thermal_event: false
 profile_power_events: false # Set to true to enable TPU-specific power/thermal profiling events. Defaults to false to avoid breaking GPU xplane tracing.
 
-log_config: true # Prints the config (after defaults have been set by pyconfig logic)
+log_config: false # Prints the config (after defaults have been set by pyconfig logic)
 debug_sharding: false # Prints model weights sharding info
 
 # Checkpoint Structured logging

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
@@ -628,7 +628,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/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
 
@@ -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):

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()

tests/unit/attention_test.py

@@ -1184,22 +1184,23 @@ def test_tpu_flash_attention_ring_context_parallel(self):
     )
     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(
@@ -1255,22 +1256,23 @@ def test_tpu_flash_attention_ring_context_parallel_grad(self):
     )
     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):