Skip to content

Deprecate get_logical_axis_rule_from_config #4568

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
NuojCheng wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Deprecate get_logical_axis_rule_from_config #4568

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/maxtext/layers/attention_op.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
4 changes: 2 additions & 2 deletions src/maxtext/layers/moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down
4 changes: 2 additions & 2 deletions src/maxtext/layers/multi_token_prediction.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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

Expand Down
11 changes: 6 additions & 5 deletions src/maxtext/layers/pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
create_sharding,
logical_to_mesh_axes,
logical_to_mesh,
get_logical_axis_rules,
)
from maxtext.utils import pipeline_utils

Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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,
)
Expand Down
10 changes: 5 additions & 5 deletions src/maxtext/models/deepseek.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

@Shuwen-Fang Shuwen-Fang Jul 23, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do other models need this update?

@NuojCheng NuojCheng Jul 23, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only models with explicit sharding onboarded have this change. I updated some other files to make sure get_logical_axis_rules are used.


import transformers

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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

Expand Down
11 changes: 0 additions & 11 deletions src/maxtext/utils/sharding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
5 changes: 3 additions & 2 deletions src/maxtext/utils/vocabulary_tiling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)

Expand Down
66 changes: 34 additions & 32 deletions tests/unit/attention_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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):
Expand Down
6 changes: 4 additions & 2 deletions tests/unit/moe_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
Loading