Integrate sharded Muon into MaxText.

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

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

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.")
 

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

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
 

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