From de353f167973cd2cf197c1ac4ff3d14f789e3c8e Mon Sep 17 00:00:00 2001 From: maxtext authors Date: Thu, 10 Sep 2026 15:30:53 -0700 Subject: [PATCH] Integrate sharded Muon into MaxText. - Implements sharded_muon_utils.py to pair Muon dimension numbers paired with NamedSharding trees. - Adds flags for muon_type ('maxtext_muon' vs 'optax_muon') and muon_use_all_to_all. - Plumbs mesh through create_training_optimizer in train_utils.py, train_compile.py, and maxtext_engine.py. This is necessary to make sharded muon work. - Adds comprehensive unit test coverage. Reverts 911bc617dd3531aa4d88405fc79dd3b6cd98f234 PiperOrigin-RevId: 979432212 --- src/maxtext/configs/base.yml | 2 - src/maxtext/configs/types.py | 14 +- src/maxtext/optimizers/optimizers.py | 42 +--- .../trainers/pre_train/train_compile.py | 2 +- src/maxtext/training_engine/maxtext_engine.py | 2 +- src/maxtext/utils/maxtext_muon_utils.py | 101 --------- src/maxtext/utils/train_utils.py | 6 +- tests/unit/maxtext_muon_utils_test.py | 206 ------------------ tests/unit/optimizers_test.py | 169 ++------------ 9 files changed, 32 insertions(+), 512 deletions(-) delete mode 100644 src/maxtext/utils/maxtext_muon_utils.py delete mode 100644 tests/unit/maxtext_muon_utils_test.py diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 0f7c26f033..016f31a322 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -1033,11 +1033,9 @@ mu_dtype: "" # data type to store "mu" of AdamW tracking the first moment. Inher # https://github.com/google-deepmind/optax/blob/main/optax/contrib/_muon.py # "mu_dtype", "adam_eps" are shared by AdamW # "nesterov", "weight_decay_mask", "adaptive" use default -muon_type: "optax_muon" # Type of Muon optimizer: "optax_muon" (or "optax") vs "maxtext_muon" (or "maxtext") muon_beta: 0.95 # Decay rate for the exponentially weighted average of grads. muon_weight_decay: 0 # Strength of the weight decay regularization. This is multiplied with the learning rate. muon_consistent_rms: None # If None, apply width scaling to updates. If float, apply consistent rms scaling (recommend 0.2). -muon_use_all_to_all: false # Whether to use all-to-all communication during Newton-Schulz iterations in maxtext_muon. # Use iota operator in Embed diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 2a3b067498..84acba8936 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -2216,25 +2216,15 @@ class AdamW(BaseModel): class Muon(BaseModel): """Configuration specific to the Muon optimizer.""" - muon_type: str = Field( - "optax_muon", - description=("Type of Muon optimizer: 'optax_muon' (or 'optax') vs 'maxtext_muon' (or 'maxtext')."), - ) muon_beta: float = Field(0.95, description="Decay rate for the exponentially weighted average of grads.") muon_weight_decay: float = Field( - 0.0, - description=("Strength of the weight decay regularization. This is multiplied with" " the learning rate."), + 0, + description="Strength of the weight decay regularization. This is multiplied with the learning rate.", ) muon_consistent_rms: float | None = Field( None, description="If None, apply width scaling to updates. If float, apply consistent rms scaling (recommend 0.2).", ) - muon_use_all_to_all: bool = Field( - False, - description=( - "Whether to use all-to-all communication during Newton-Schulz" " iterations in the maxtext_muon optimizer." - ), - ) class PositionalEmbedding(BaseModel): diff --git a/src/maxtext/optimizers/optimizers.py b/src/maxtext/optimizers/optimizers.py index 8471048e75..67e1f589ca 100644 --- a/src/maxtext/optimizers/optimizers.py +++ b/src/maxtext/optimizers/optimizers.py @@ -20,11 +20,9 @@ import jax.numpy as jnp import optax -from optax.contrib._muon import muon as optax_muon +from optax.contrib._muon import muon from maxtext.common.common_types import DecoderBlockType -from maxtext.optimizers.muon import muon as maxtext_muon from maxtext.utils.muon_utils import get_muon_weight_dimension_numbers -from maxtext.utils.maxtext_muon_utils import get_maxtext_muon_weight_dimension_numbers def _get_path_mask_fn(patterns, match_returns_true=True): @@ -168,7 +166,7 @@ def skip_update(): return optax.GradientTransformationExtraArgs(init_fn, update_fn) -def get_optimizer(config, learning_rate_schedule, model=None, mesh=None): +def get_optimizer(config, learning_rate_schedule, model=None): """Create optimizer.""" if config.opt_type == "adamw": # Create AdamW Optimizer following Llama2's training details, see https://arxiv.org/pdf/2307.09288.pdf section 2.2 @@ -195,7 +193,12 @@ def get_optimizer(config, learning_rate_schedule, model=None, mesh=None): elif config.opt_type == "sgd": base_opt = optax.sgd(learning_rate_schedule) elif config.opt_type == "muon": - muon_type = getattr(config, "muon_type", "optax_muon") + # extract muon dimension number from model structure + if model is not None: + muon_weight_dimension_numbers = get_muon_weight_dimension_numbers(model, config) + else: + raise ValueError("Please specify model to extract muon dimension number.") + if config.decoder_block == DecoderBlockType.DEEPSEEK4: ns_coeffs = [(3.4445, -4.7750, 2.0315)] * 8 + [(2.0, -1.5, 0.5)] * 2 ns_steps = 10 @@ -211,6 +214,7 @@ def get_optimizer(config, learning_rate_schedule, model=None, mesh=None): # Muon-specific parameters: "weight_decay_mask", "adaptive" uses default "beta": config.muon_beta, "weight_decay": config.muon_weight_decay, + "muon_weight_dimension_numbers": muon_weight_dimension_numbers, "consistent_rms": config.muon_consistent_rms, "ns_coeffs": ns_coeffs, "ns_steps": ns_steps, @@ -220,33 +224,7 @@ def get_optimizer(config, learning_rate_schedule, model=None, mesh=None): "adam_eps_root": config.adam_eps_root, "adam_weight_decay": config.adam_weight_decay, } - - if muon_type in ("optax", "optax_muon"): - assert not getattr( - config, "muon_use_all_to_all", False - ), "all-to-all communication in muon is only supported with maxtext_muon, not optax_muon." - if model is not None: - muon_weight_dimension_numbers = get_muon_weight_dimension_numbers(model, config) - else: - raise ValueError("Please specify model to extract muon dimension number.") - muon_kwargs = muon_kwargs | {"muon_weight_dimension_numbers": muon_weight_dimension_numbers} - base_opt = optax_muon(**muon_kwargs) # pyrefly: ignore[bad-argument-type] - elif muon_type in ("maxtext", "maxtext_muon"): - if model is not None: - muon_weight_dimension_numbers = get_maxtext_muon_weight_dimension_numbers(model, config, mesh=mesh) - else: - raise ValueError("Please specify model to extract muon dimension number.") - - use_all_to_all = getattr(config, "muon_use_all_to_all", True) - muon_kwargs = muon_kwargs | { - "muon_weight_dimension_numbers": muon_weight_dimension_numbers, - "use_all_to_all": use_all_to_all, - } - base_opt = maxtext_muon(**muon_kwargs) # pyrefly: ignore[bad-argument-type] - else: - raise ValueError( - f"Unsupported muon_type: {muon_type}. Must be 'optax_muon' (or 'optax') or 'maxtext_muon' (or 'maxtext')." - ) + base_opt = muon(**muon_kwargs) # pyrefly: ignore[bad-argument-type] else: raise ValueError(f"{config.opt_type=} is not a supported.") diff --git a/src/maxtext/trainers/pre_train/train_compile.py b/src/maxtext/trainers/pre_train/train_compile.py index 0c69d31f35..d87ab1a991 100644 --- a/src/maxtext/trainers/pre_train/train_compile.py +++ b/src/maxtext/trainers/pre_train/train_compile.py @@ -133,7 +133,7 @@ def get_shaped_inputs(topology_mesh, config): # The learning_rate_schedule is baked into the compiled object. learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) # pass in model for muon - tx = optimizers.get_optimizer(config, learning_rate_schedule, model, mesh=topology_mesh) + tx = optimizers.get_optimizer(config, learning_rate_schedule, model) def create_train_state_fn(): nnx_model = _create_model_partial() diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 43c77b507f..d572fcacff 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -593,7 +593,7 @@ def __init__( # calls `optimizer.update(model, grads)`, which is the nnx.Optimizer signature, and # `checkpointing.CheckpointState` expects an nnx.Optimizer too, so wrap it here. `wrt=nnx.Param` # covers every parameter, which is correct only because LoRA is rejected above. - self._learning_rate_schedule, tx = train_utils.create_training_optimizer(self._config, self._model, mesh=self._mesh) + self._learning_rate_schedule, tx = train_utils.create_training_optimizer(self._config, self._model) self._optimizer = self._build_optimizer(tx) self._train_step: int = 0 diff --git a/src/maxtext/utils/maxtext_muon_utils.py b/src/maxtext/utils/maxtext_muon_utils.py deleted file mode 100644 index 5aa16e441c..0000000000 --- a/src/maxtext/utils/maxtext_muon_utils.py +++ /dev/null @@ -1,101 +0,0 @@ -# Copyright 2023–2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -"""Utilities for MaxText Muon optimizer integration and dimension/sharding generation. - -This module provides functions to automatically generate -ShardedMuonDimensionNumbers with physical NamedSharding for various MaxText -models. -""" - -from typing import Tuple - -from flax import nnx -from flax.linen import partitioning as nn_partitioning -import jax -from maxtext.optimizers.muon import ShardedMuonDimensionNumbers as smdn -from maxtext.utils import maxtext_utils -from maxtext.utils import muon_utils -from maxtext.utils import sharding as sharding_lib - - -def _get_leaf_value(leaf): - """Extracts value from leaf, unwrapping get_value() if present.""" - if hasattr(leaf, "get_value"): - return leaf.get_value() - return leaf - - -def _extract_sharding(leaf): - """Extracts NamedSharding or PartitionSpec from a leaf or its attributes.""" - val = _get_leaf_value(leaf) - if isinstance(val, (jax.sharding.NamedSharding, jax.sharding.PartitionSpec)): - return val - sharding = getattr(leaf, "sharding", None) - if isinstance(sharding, (jax.sharding.NamedSharding, jax.sharding.PartitionSpec)): - return sharding - return None - - -def get_maxtext_muon_weight_dimension_numbers(model, config=None, mesh=None, verbose=False): - """Extracts a matching pytree of ShardedMuonDimensionNumbers with physical shardings from a model.""" - if mesh is None: - if config is None: - raise ValueError("Either mesh or config must be provided to get_maxtext_muon_weight_dimension_numbers.") - mesh = maxtext_utils.get_mesh_from_config(config) - - if config is not None: - logical_rules = getattr(config, "logical_axis_rules", ()) - else: - logical_rules = () - - with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_rules): - # Extract abstract parameters from the NNX model hierarchy - _, abstract_param, _ = nnx.split(model, nnx.Param, ...) - - # Resolve physical NamedSharding for each parameter under the active axis rules - named_sharding_state = sharding_lib.nnx_construct_named_sharding(abstract_param, mesh) - abstract_dict = nnx.to_pure_dict(abstract_param) - named_sharding_dict = nnx.to_pure_dict(named_sharding_state) - - def apply_transform_nnx(path: Tuple[jax.tree_util.KeyEntry, ...], leaf, abs_leaf): - path_strings = tuple(p.key for p in path if isinstance(p, jax.tree_util.DictKey)) - abs_val = _get_leaf_value(abs_leaf) - val_shape = getattr(abs_val, "shape", None) - dim_num = muon_utils.transform_logic(path_strings, shape=val_shape) - if dim_num is not None: - sharding = _extract_sharding(leaf) - return smdn( - reduction_axis=dim_num.reduction_axis, - output_axis=dim_num.output_axis, - sharding=sharding, - ) - return None - - # Walk the parameter tree to produce a matching nnx.State of ShardedMuonDimensionNumbers - muon_weight_dimension_numbers = jax.tree.map_with_path(apply_transform_nnx, named_sharding_dict, abstract_dict) - muon_weight_dimension_numbers = nnx.State(muon_weight_dimension_numbers) - - if verbose: - _print_structure_debug(abstract_param, muon_weight_dimension_numbers) - return muon_weight_dimension_numbers - - -get_sharded_muon_weight_dimension_numbers = get_maxtext_muon_weight_dimension_numbers - - -def _print_structure_debug(abstract_param, muon_weight_dimension_numbers): - """Prints the model structure and the resulting Muon config.""" - return muon_utils._print_structure_debug(abstract_param, muon_weight_dimension_numbers) # pylint: disable=protected-access diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index a5efe5abc3..e83874420f 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -42,11 +42,11 @@ from maxtext.utils.rampup_batch import create_rampup_manager -def create_training_optimizer(config, model, mesh=None): +def create_training_optimizer(config, model): """Creates the optimizer and learning rate schedule.""" learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) # pass in model for muon - tx = optimizers.get_optimizer(config, learning_rate_schedule, model, mesh=mesh) + tx = optimizers.get_optimizer(config, learning_rate_schedule, model) return learning_rate_schedule, tx @@ -259,7 +259,7 @@ def setup_train_loop(config, recorder, devices=None): context_parallel_size = mesh.shape.get(config.context_sharding, 1) # Create abstract NNX model. _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices) - learning_rate_schedule, tx = create_training_optimizer(config, model, mesh=mesh) + learning_rate_schedule, tx = create_training_optimizer(config, model) # The train state is wrapped in the TrainStateNNX module. def create_train_state_fn(): diff --git a/tests/unit/maxtext_muon_utils_test.py b/tests/unit/maxtext_muon_utils_test.py deleted file mode 100644 index 1d343f9619..0000000000 --- a/tests/unit/maxtext_muon_utils_test.py +++ /dev/null @@ -1,206 +0,0 @@ -# Copyright 2023–2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Unit tests for maxtext_muon_utils.py.""" - -# pylint: disable=protected-access - -import collections.abc -import contextlib -import io -import unittest -from unittest import mock - -from absl.testing import absltest -from flax import linen as nn -from flax import nnx -import jax -import jax.numpy as jnp -from maxtext.optimizers.muon import ShardedMuonDimensionNumbers as smdn -from maxtext.utils import muon_utils -from maxtext.utils import maxtext_muon_utils -import numpy as np - - -def _extract_axes(tree): - """Recursively extracts (reduction_axis, output_axis) or None from dimension number trees.""" - if isinstance(tree, nnx.State): - tree = nnx.to_pure_dict(tree) - if isinstance(tree, (dict, collections.abc.Mapping)) or hasattr(tree, "items"): - return {k: _extract_axes(v) for k, v in tree.items()} - if tree is None: - return None - return (tree.reduction_axis, tree.output_axis) - - -class _AttentionSubModule(nnx.Module): - """Placeholder NNX attention sub-module for testing parameter extraction.""" - - def __init__(self): - self.out = nnx.Param(jnp.ones((2, 4, 8))) - - -class _MoeLikeNNXModel(nnx.Module): - """Placeholder NNX model for testing parameter extraction.""" - - def __init__(self, rngs): - del rngs # Unused. - self.w_standard = nnx.Param(jnp.ones((4, 8))) - self.self_attention = _AttentionSubModule() - self.scale = nnx.Param(jnp.ones((8,))) - - -class TestExtractSharding(unittest.TestCase): - """Tests for _extract_sharding utility.""" - - def setUp(self): - super().setUp() - devices = np.array(jax.devices()[:1]).reshape((1, 1)) - self.mesh = jax.sharding.Mesh(devices, ("data", "model")) - self.named_sharding = jax.sharding.NamedSharding(self.mesh, jax.sharding.PartitionSpec("data", "model")) - self.partition_spec = jax.sharding.PartitionSpec("data", "model") - - def test_extracts_from_named_sharding(self): - self.assertEqual( - maxtext_muon_utils._extract_sharding(self.named_sharding), - self.named_sharding, - ) - - def test_extracts_from_partition_spec(self): - self.assertEqual( - maxtext_muon_utils._extract_sharding(self.partition_spec), - self.partition_spec, - ) - - def test_extracts_from_variable_with_get_value(self): - leaf = mock.MagicMock() - leaf.get_value.return_value = self.named_sharding - self.assertEqual( - maxtext_muon_utils._extract_sharding(leaf), - self.named_sharding, - ) - - def test_extracts_from_attribute(self): - leaf = mock.MagicMock(spec=["sharding"]) - leaf.sharding = self.named_sharding - self.assertEqual( - maxtext_muon_utils._extract_sharding(leaf), - self.named_sharding, - ) - - def test_returns_none_for_raw_shape_struct(self): - leaf = jax.ShapeDtypeStruct((4, 8), jnp.float32) - self.assertIsNone(maxtext_muon_utils._extract_sharding(leaf)) - - def test_returns_none_for_non_sharding_value(self): - leaf = mock.MagicMock() - leaf.get_value.return_value = jnp.zeros((4, 8)) - del leaf.sharding - self.assertIsNone(maxtext_muon_utils._extract_sharding(leaf)) - - -class TestGetShardedMuonWeightDimensionNumbersNNX(unittest.TestCase): - """Tests for get_sharded_muon_weight_dimension_numbers with NNX models.""" - - def setUp(self): - super().setUp() - self.model = _MoeLikeNNXModel(rngs=nnx.Rngs(0)) - devices = np.array(jax.devices()[:1]).reshape((1, 1)) - self.mesh = jax.sharding.Mesh(devices, ("data", "model")) - - def test_nnx_model_axes_match_base_muon(self): - sharded_result = maxtext_muon_utils.get_sharded_muon_weight_dimension_numbers(self.model, mesh=self.mesh) - base_result = muon_utils.get_muon_weight_dimension_numbers(self.model) - self.assertEqual(_extract_axes(sharded_result), _extract_axes(base_result)) - - def test_nnx_model_with_logical_axis_rules_axes_match_base_muon(self): - config = mock.MagicMock() - config.logical_axis_rules = (("embed", "fsdp"), ("mlp", "tensor")) - sharded_result = maxtext_muon_utils.get_sharded_muon_weight_dimension_numbers( - self.model, config=config, mesh=self.mesh - ) - base_result = muon_utils.get_muon_weight_dimension_numbers(self.model, config=config) - self.assertEqual(_extract_axes(sharded_result), _extract_axes(base_result)) - - def test_nnx_model_populates_named_sharding(self): - result = maxtext_muon_utils.get_sharded_muon_weight_dimension_numbers(self.model, mesh=self.mesh) - self.assertIsNotNone(result["w_standard"].sharding) - self.assertIsInstance(result["w_standard"].sharding, jax.sharding.NamedSharding) - self.assertEqual(result["w_standard"].sharding.mesh, self.mesh) - - def test_nnx_model_resolves_mesh_from_config_when_mesh_is_none(self): - fake_mesh = jax.sharding.Mesh( - self.mesh.devices, - ("data", "model"), - axis_types=( - jax.sharding.AxisType.Explicit, - jax.sharding.AxisType.Explicit, - ), - ) - config = mock.MagicMock() - config.shard_mode = "explicit" - config.mesh_axes = ("data", "model") - config.logical_axis_rules = () - - with mock.patch.object( - maxtext_muon_utils.maxtext_utils, - "get_mesh_from_config", - return_value=fake_mesh, - ) as mock_get_mesh: - sharded_result = maxtext_muon_utils.get_sharded_muon_weight_dimension_numbers(self.model, config=config, mesh=None) - mock_get_mesh.assert_called_once_with(config) - self.assertIsNotNone(sharded_result["w_standard"].sharding) - self.assertEqual(sharded_result["w_standard"].sharding.mesh, fake_mesh) - - def test_raises_error_when_mesh_and_config_are_none(self): - with self.assertRaisesRegex(ValueError, "Either mesh or config must be provided"): - maxtext_muon_utils.get_sharded_muon_weight_dimension_numbers(self.model, config=None, mesh=None) - - def test_nnx_verbose_prints_debug_structure(self): - buf = io.StringIO() - with contextlib.redirect_stdout(buf): - maxtext_muon_utils.get_sharded_muon_weight_dimension_numbers(self.model, mesh=self.mesh, verbose=True) - self.assertIn("Model Structure", buf.getvalue()) - self.assertIn("Muon Dimension Numbers", buf.getvalue()) - - -class TestPrintStructureDebug(unittest.TestCase): - """Tests for _print_structure_debug.""" - - def test_prints_logically_partitioned_leaf_info(self): - leaf = nn.LogicallyPartitioned(value=jax.ShapeDtypeStruct((4, 8), jnp.float32), names=("embed", "mlp")) - tree = {"params": {"kernel": leaf}} - - buf = io.StringIO() - with contextlib.redirect_stdout(buf): - maxtext_muon_utils._print_structure_debug( - tree, - muon_weight_dimension_numbers={"params": {"kernel": smdn((0,), (-1,))}}, - ) - out = buf.getvalue() - self.assertIn("(4, 8)", out) - self.assertIn("embed", out) - - def test_prints_shape_dtype_struct_leaf_info(self): - tree = {"kernel": jax.ShapeDtypeStruct((16, 32), jnp.float32)} - - buf = io.StringIO() - with contextlib.redirect_stdout(buf): - maxtext_muon_utils._print_structure_debug(tree, muon_weight_dimension_numbers={"kernel": smdn((0,), (-1,))}) - out = buf.getvalue() - self.assertIn("(16, 32)", out) - - -if __name__ == "__main__": - absltest.main() diff --git a/tests/unit/optimizers_test.py b/tests/unit/optimizers_test.py index 3c7574bc05..710a56f2a9 100644 --- a/tests/unit/optimizers_test.py +++ b/tests/unit/optimizers_test.py @@ -24,10 +24,8 @@ import jax.numpy as jnp from maxtext.configs import pyconfig from maxtext.optimizers import optimizers -from maxtext.optimizers.muon import ShardedMuonDimensionNumbers as smdn from maxtext.utils import maxtext_utils, muon_utils from tests.utils.test_helpers import get_test_config_path -import numpy as np import optax from optax.contrib._muon import MuonDimensionNumbers as mdn import pytest @@ -967,184 +965,47 @@ def __init__(self, rngs: nnx.Rngs): self.assertEqual(result.self_attention.out, mdn((0, -2), (-1,))) def test_muon_newton_schulz_config(self): - """Verifies that muon optimizer configures Newton-Schulz parameters correctly based on model and muon_type.""" + """Verifies that muon optimizer configures Newton-Schulz parameters correctly based on model.""" model = MagicMock() - mesh = MagicMock() learning_rate_schedule = MagicMock() - # Case 1: DeepSeek4 Model with MaxText Muon (Auto-configures 10-step schedule) + # Case 1: DeepSeek4 Model (Auto-configures 10-step schedule) argv_ds4 = [ "", get_test_config_path(), "run_name=test", "opt_type=muon", - "muon_type=maxtext_muon", "model_name=deepseek4-284b", "attention=dot_product", - "muon_use_all_to_all=True", ] config_ds4 = pyconfig.initialize(argv_ds4) with ( - patch.object(optimizers, "get_maxtext_muon_weight_dimension_numbers") as mock_get_smdn, - patch.object(optimizers, "maxtext_muon") as mock_maxtext_muon, + patch.object(optimizers, "get_muon_weight_dimension_numbers") as mock_get_mdn, + patch.object(optimizers, "muon") as mock_muon, ): - mock_get_smdn.return_value = {"params": {"w": smdn((0,), (-1,))}} - optimizers.get_optimizer(config_ds4, learning_rate_schedule, model=model, mesh=mesh) - mock_get_smdn.assert_called_once_with(model, config_ds4, mesh=mesh) - mock_maxtext_muon.assert_called_once() - _, kwargs = mock_maxtext_muon.call_args + mock_get_mdn.return_value = {} + optimizers.get_optimizer(config_ds4, learning_rate_schedule, model=model) + mock_muon.assert_called_once() + _, kwargs = mock_muon.call_args self.assertEqual(kwargs["ns_steps"], 10) self.assertEqual(len(kwargs["ns_coeffs"]), 10) self.assertEqual(kwargs["ns_coeffs"][-1], (2.0, -1.5, 0.5)) - self.assertEqual(kwargs["muon_weight_dimension_numbers"], mock_get_smdn.return_value) - self.assertTrue(kwargs["use_all_to_all"]) - - # Case 2: MaxText Muon with use_all_to_all=False - argv_no_a2a = [ - "", - get_test_config_path(), - "run_name=test", - "opt_type=muon", - "muon_type=maxtext_muon", - "model_name=llama2-7b", - "muon_use_all_to_all=False", - ] - config_no_a2a = pyconfig.initialize(argv_no_a2a) - - with ( - patch.object(optimizers, "get_maxtext_muon_weight_dimension_numbers") as mock_get_smdn, - patch.object(optimizers, "maxtext_muon") as mock_maxtext_muon, - ): - mock_get_smdn.return_value = {} - optimizers.get_optimizer(config_no_a2a, learning_rate_schedule, model=model) - mock_get_smdn.assert_called_once_with(model, config_no_a2a, mesh=None) - mock_maxtext_muon.assert_called_once() - _, kwargs = mock_maxtext_muon.call_args - self.assertEqual(kwargs["ns_steps"], 5) - self.assertEqual(kwargs["ns_coeffs"], (3.4445, -4.7750, 2.0315)) - self.assertEqual(kwargs["muon_weight_dimension_numbers"], mock_get_smdn.return_value) - self.assertFalse(kwargs["use_all_to_all"]) - - # Case 3: Default Muon Type (defaults to optax_muon) - argv_default = [ - "", - get_test_config_path(), - "run_name=test", - "opt_type=muon", - "model_name=llama2-7b", - ] - config_default = pyconfig.initialize(argv_default) - - with ( - patch.object(optimizers, "get_muon_weight_dimension_numbers") as mock_get_mdn, - patch.object(optimizers, "optax_muon") as mock_optax_muon, - ): - mock_get_mdn.return_value = {} - optimizers.get_optimizer(config_default, learning_rate_schedule, model=model) - mock_get_mdn.assert_called_once_with(model, config_default) - mock_optax_muon.assert_called_once() - _, kwargs = mock_optax_muon.call_args - self.assertEqual(kwargs["ns_steps"], 5) - self.assertEqual(kwargs["ns_coeffs"], (3.4445, -4.7750, 2.0315)) - self.assertEqual(kwargs["muon_weight_dimension_numbers"], mock_get_mdn.return_value) - self.assertNotIn("use_all_to_all", kwargs) - # Case 4: Explicit Optax Muon (muon_type=optax_muon) - argv_optax = [ - "", - get_test_config_path(), - "run_name=test", - "opt_type=muon", - "muon_type=optax_muon", - "model_name=llama2-7b", - ] - config_optax = pyconfig.initialize(argv_optax) + # Case 2: Standard Model (Llama2) (Defaults to 5-step schedule) + argv_llama = ["", get_test_config_path(), "run_name=test", "opt_type=muon", "model_name=llama2-7b"] + config_llama = pyconfig.initialize(argv_llama) with ( patch.object(optimizers, "get_muon_weight_dimension_numbers") as mock_get_mdn, - patch.object(optimizers, "optax_muon") as mock_optax_muon, + patch.object(optimizers, "muon") as mock_muon, ): mock_get_mdn.return_value = {} - optimizers.get_optimizer(config_optax, learning_rate_schedule, model=model) - mock_get_mdn.assert_called_once_with(model, config_optax) - mock_optax_muon.assert_called_once() - _, kwargs = mock_optax_muon.call_args + optimizers.get_optimizer(config_llama, learning_rate_schedule, model=model) + mock_muon.assert_called_once() + _, kwargs = mock_muon.call_args self.assertEqual(kwargs["ns_steps"], 5) self.assertEqual(kwargs["ns_coeffs"], (3.4445, -4.7750, 2.0315)) - self.assertEqual(kwargs["muon_weight_dimension_numbers"], mock_get_mdn.return_value) - self.assertNotIn("use_all_to_all", kwargs) - - # Case 5: Optax Muon with muon_use_all_to_all=True raises AssertionError - argv_optax_a2a = [ - "", - get_test_config_path(), - "run_name=test", - "opt_type=muon", - "muon_type=optax_muon", - "model_name=llama2-7b", - "muon_use_all_to_all=True", - ] - config_optax_a2a = pyconfig.initialize(argv_optax_a2a) - with self.assertRaises(AssertionError): - optimizers.get_optimizer(config_optax_a2a, learning_rate_schedule, model=model) - - # Case 6: Invalid muon_type raises ValueError - argv_invalid = [ - "", - get_test_config_path(), - "run_name=test", - "opt_type=muon", - "muon_type=invalid_type", - "model_name=llama2-7b", - ] - config_invalid = pyconfig.initialize(argv_invalid) - with self.assertRaises(ValueError): - optimizers.get_optimizer(config_invalid, learning_rate_schedule, model=model) - - def test_maxtext_muon_wires_named_sharding_and_mesh(self): - """Verifies that get_optimizer forwards mesh and propagates ShardedMuonDimensionNumbers with NamedSharding.""" - devices = np.array(jax.devices()[:1]).reshape((1, 1)) - mesh = jax.sharding.Mesh(devices, ("data", "model")) - sharded_dim_nums = { - "params": { - "mlp": { - "kernel": smdn( - reduction_axis=(0,), - output_axis=(-1,), - sharding=jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("data", "model")), - ) - } - } - } - model = MagicMock() - learning_rate_schedule = MagicMock() - argv = [ - "", - get_test_config_path(), - "run_name=test", - "opt_type=muon", - "muon_type=maxtext_muon", - "model_name=llama2-7b", - ] - config = pyconfig.initialize(argv) - - with ( - patch.object( - optimizers, "get_maxtext_muon_weight_dimension_numbers", return_value=sharded_dim_nums - ) as mock_get_smdn, - patch.object(optimizers, "maxtext_muon") as mock_maxtext_muon, - ): - optimizers.get_optimizer(config, learning_rate_schedule, model=model, mesh=mesh) - mock_get_smdn.assert_called_once_with(model, config, mesh=mesh) - mock_maxtext_muon.assert_called_once() - _, kwargs = mock_maxtext_muon.call_args - passed_dim_nums = kwargs["muon_weight_dimension_numbers"] - self.assertEqual(passed_dim_nums, sharded_dim_nums) - kernel_spec = passed_dim_nums["params"]["mlp"]["kernel"] - self.assertIsInstance(kernel_spec.sharding, jax.sharding.NamedSharding) - self.assertEqual(kernel_spec.sharding.mesh, mesh) - self.assertEqual(kernel_spec.sharding.spec, jax.sharding.PartitionSpec("data", "model")) if __name__ == "__main__":