diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index d096b8abd1..52128e4c08 100644 --- a/src/maxtext/layers/attention_op.py +++ b/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: diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 8127daa0f9..5f7556ae5a 100644 --- a/src/maxtext/layers/moe.py +++ b/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): diff --git a/src/maxtext/layers/multi_token_prediction.py b/src/maxtext/layers/multi_token_prediction.py index 6f9af709fd..3dd3ad3a6e 100644 --- a/src/maxtext/layers/multi_token_prediction.py +++ b/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 diff --git a/src/maxtext/layers/pipeline.py b/src/maxtext/layers/pipeline.py index 07083a824e..bf66fbcce8 100644 --- a/src/maxtext/layers/pipeline.py +++ b/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,13 +62,13 @@ 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 ) @@ -75,7 +76,7 @@ def _setup_pipeline_attributes(self): 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, ) diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py index 655c346c03..a09e1d28c8 100644 --- a/src/maxtext/models/deepseek.py +++ b/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 diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index 66b53e640f..bacdb631d2 100644 --- a/src/maxtext/utils/sharding.py +++ b/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() diff --git a/src/maxtext/utils/vocabulary_tiling.py b/src/maxtext/utils/vocabulary_tiling.py index 9633528d63..eed16f14d8 100644 --- a/src/maxtext/utils/vocabulary_tiling.py +++ b/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, ) diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 754c526b2f..1172670930 100644 --- a/tests/unit/attention_test.py +++ b/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): diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index 1cae70ffb9..8eb3bd70c9 100644 --- a/tests/unit/moe_test.py +++ b/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(