From cbcaaeb9311fb09b29394bbf0b28b9577ead1310 Mon Sep 17 00:00:00 2001 From: darisoy Date: Wed, 29 Jul 2026 23:35:08 +0000 Subject: [PATCH 01/12] Update qwen3-next-80b-a3b.yml config parameters --- .../configs/models/qwen3-next-80b-a3b.yml | 38 +++++++++++-------- 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/src/maxtext/configs/models/qwen3-next-80b-a3b.yml b/src/maxtext/configs/models/qwen3-next-80b-a3b.yml index 765977f1b5..327a395afb 100644 --- a/src/maxtext/configs/models/qwen3-next-80b-a3b.yml +++ b/src/maxtext/configs/models/qwen3-next-80b-a3b.yml @@ -18,35 +18,41 @@ decoder_block: "qwen3_next" # Core Architectural Parameters -base_emb_dim: 2048 -base_num_decoder_layers: 48 -base_num_query_heads: 16 -base_num_kv_heads: 2 -head_dim: 256 -vocab_size: 151936 +base_emb_dim: 3072 +base_num_decoder_layers: 40 +base_num_query_heads: 64 +base_num_kv_heads: 8 +head_dim: 128 +vocab_size: 128008 normalization_layer_epsilon: 1.0e-6 # MoE Specific Parameters -# Set base_mlp_dim to match base_moe_mlp_dim to pass validation for fully MoE models. -base_mlp_dim: 512 -base_moe_mlp_dim: 512 -num_experts: 512 +# base_mlp_dim sizes the dense-prefix layer's MLP +# base_moe_mlp_dim sizes every other (MoE) layer's routed + shared experts. +base_mlp_dim: 8192 +base_moe_mlp_dim: 1536 +num_experts: 128 shared_experts: 1 -num_experts_per_tok: 10 +num_experts_per_tok: 8 norm_topk_prob: true +# The first layer is a dense MLP (no MoE) and always uses full attention. +first_num_dense_layers: 1 + # Qwen3-Next Specific Parameters for Linear Attention (Gated Delta Net) -inhomogeneous_layer_cycle_interval: 4 +inhomogeneous_layer_cycle_interval: 3 gdn_conv_kernel_dim: 4 gdn_key_head_dim: 128 gdn_value_head_dim: 128 -gdn_num_key_heads: 16 -gdn_num_value_heads: 32 +gdn_num_key_heads: 32 +gdn_num_value_heads: 64 gdn_chunk_size: 64 # RoPE Settings -rope_max_timescale: 10000000 -partial_rotary_factor: 0.25 +rope_max_timescale: 10000 +partial_rotary_factor: 1.0 + +mhc_expansion_rate: 4 # General Model Settings enable_dropout: false From 00302b21ae30ec90d783df971574eedec9c40b2e Mon Sep 17 00:00:00 2001 From: darisoy Date: Thu, 30 Jul 2026 01:42:56 +0000 Subject: [PATCH 02/12] Integrate mHC into Qwen3NextDecoderLayer (mhc_expansion_rate=4) --- src/maxtext/models/qwen3.py | 63 +++++++++++++++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 3cac10ff25..e16d7b5e49 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -37,6 +37,8 @@ from maxtext.layers import attentions from maxtext.layers import initializers as max_initializers from maxtext.layers import moe +from maxtext.layers import mhc +from maxtext.common.common_types import HyperConnectionType from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.embeddings import Qwen3OmniMoeVisionPosEmbedInterpolate, PositionalEmbedding @@ -1294,6 +1296,15 @@ def __init__( # Instantiate our `Qwen3NextSparseMoeBlock`. self.mlp = Qwen3NextSparseMoeBlock(config=cfg, mesh=self.mesh, quant=self.quant, rngs=rngs) + self.is_mhc_enabled = getattr(cfg, "mhc_expansion_rate", 1) > 1 + if self.is_mhc_enabled: + self.mhc_attention = mhc.ManifoldConstrainedHyperConnections( + config=cfg, dim=cfg.emb_dim, mesh=self.mesh, rngs=rngs + ) + self.mhc_mlp = mhc.ManifoldConstrainedHyperConnections( + config=cfg, dim=cfg.emb_dim, mesh=self.mesh, rngs=rngs + ) + def __call__( self, inputs: jnp.ndarray, @@ -1309,6 +1320,58 @@ def __call__( # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) if isinstance(inputs, tuple): inputs = inputs[0] + + if self.is_mhc_enabled: + new_kv_cache = None + + def attention_branch(inputs): + nonlocal new_kv_cache + if isinstance(self.attention, Qwen3NextFullAttention): + out, new_kv_cache = cast(Qwen3NextFullAttention, self.attention)( + inputs, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + ) + else: + out, new_kv_cache = cast(Qwen3NextGatedDeltaNet, self.attention)( + inputs, + model_mode=model_mode, + kv_cache=kv_cache, + decoder_segment_ids=decoder_segment_ids, + attention_metadata=attention_metadata, + ) + return out + + intermediate_inputs, _ = self.mhc_attention( + self.input_layernorm, + attention_branch, + x=inputs, + mhc_type=HyperConnectionType.MLP_DENSE, + ) + + def mlp_branch(inputs): + mlp_output, load_balance_loss = self.mlp(inputs, deterministic=deterministic) + if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: + self.moe_lb_loss = nnx.Intermediate(load_balance_loss) + return mlp_output + + layer_output, _ = self.mhc_mlp( + self.post_attention_layernorm, + mlp_branch, + x=intermediate_inputs, + mhc_type=HyperConnectionType.MLP_DENSE, + ) + + layer_output = nn.with_logical_constraint( + layer_output, + self.activation_axis_names, + ) + return layer_output, new_kv_cache + residual = inputs # First LayerNorm, applied before the attention block. From 0360d026ee5cfc794b79d74ec8a053e4fa5c0683 Mon Sep 17 00:00:00 2001 From: darisoy Date: Wed, 29 Jul 2026 23:41:49 +0000 Subject: [PATCH 03/12] Implement Newton-Schulz log-depth custom VJP triangular matrix solver for GDN (PR #4577 / cl/955390793) --- tests/unit/qwen3_next_vs_reference_test.py | 35 ++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/unit/qwen3_next_vs_reference_test.py b/tests/unit/qwen3_next_vs_reference_test.py index e9efad376f..7ffa0c0990 100644 --- a/tests/unit/qwen3_next_vs_reference_test.py +++ b/tests/unit/qwen3_next_vs_reference_test.py @@ -22,6 +22,7 @@ import jax import jax.numpy as jnp from jax.sharding import Mesh +from jax.test_util import check_grads from maxtext.configs import pyconfig from maxtext.layers import normalizations from maxtext.layers.normalizations import Qwen3NextRMSNorm, Qwen3NextRMSNormGated @@ -1037,6 +1038,40 @@ def run_jax(x): ) print("test_qwen3_next_sparse_moe_block passed!") + def test_invert_unit_lower_triangular_log_depth(self): + """Test for loss at chunk_size 256.""" + jax.config.update("jax_enable_x64", True) # Use float64 for precise testing + chunk_size = 256 + + # Generate a random matrix and make it strictly lower triangular + key = jax.random.PRNGKey(chunk_size) + S_random = jax.random.normal(key, (chunk_size, chunk_size), dtype=jnp.float64) / chunk_size + S = jnp.tril(S_random, k=-1) + + # The matrix to invert is (I + S) + identity = jnp.eye(chunk_size, dtype=jnp.float64) + matrix_to_invert = identity + S + + # Using our custom function + A = qwen3.invert_unit_lower_triangular_log_depth(S) + + # The product A @ (I + S) should be exactly the identity matrix + # Wait, due to numerical precision, we should check for max error (loss) + reconstructed_identity = A @ matrix_to_invert + + # Compute loss for forward pass + loss = jnp.max(jnp.abs(reconstructed_identity - identity)) + + # We expect the loss to be very small, around numerical precision + self.assertLess(loss, 1e-10, f"Failed for chunk_size {chunk_size} with loss {loss}") + + # Verify backward pass accuracy using jax.test_util.check_grads + # This uses finite differences to check the correctness of the custom VJP + # We check the gradients for the function. + # `check_grads` will assert if finite difference gradients + # don't match the custom VJP gradients. + check_grads(qwen3.invert_unit_lower_triangular_log_depth, (S,), order=1, modes=["rev"]) + def test_gated_delta_net_full(self): """Tests the full Qwen3NextGatedDeltaNet layer for numerical correctness.""" print("Running test_gated_delta_net_full...") From 4960fea253d1059b6f6aec96a8db6278250d0acb Mon Sep 17 00:00:00 2001 From: darisoy Date: Fri, 31 Jul 2026 19:23:28 +0000 Subject: [PATCH 04/12] Implement Hybrid GDN v3 Tokamax forward + Newton-Schulz Custom VJP backward (with shard_map sharding support) --- .../dockerfiles/maxtext_runner.Dockerfile | 3 + src/maxtext/configs/base.yml | 4 + src/maxtext/configs/types.py | 8 + src/maxtext/models/hybrid_gdn.py | 279 ++++++++++++++++++ src/maxtext/models/qwen3.py | 110 +++++++ 5 files changed, 404 insertions(+) create mode 100644 src/maxtext/models/hybrid_gdn.py diff --git a/src/dependencies/dockerfiles/maxtext_runner.Dockerfile b/src/dependencies/dockerfiles/maxtext_runner.Dockerfile index d85511c848..02b138931d 100644 --- a/src/dependencies/dockerfiles/maxtext_runner.Dockerfile +++ b/src/dependencies/dockerfiles/maxtext_runner.Dockerfile @@ -14,6 +14,9 @@ ENV MAXTEXT_REPO_ROOT=/deps # Set the working directory in the container WORKDIR /deps +# Install GDN v3 Tokamax commit +RUN pip install --no-deps --no-cache-dir --force-reinstall git+https://github.com/openxla/tokamax.git@b626dd8b54d708047788cf2ec538cba63a4e3739 + # Copy assets separately COPY ${PACKAGE_DIR}/maxtext/assets/ "${MAXTEXT_ASSETS_ROOT}" diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index db0b1a68a7..6b717f91c8 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -1271,6 +1271,10 @@ gdn_num_value_heads: 32 gdn_chunk_size: 64 # Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel. use_qk_norm_in_gdn: true +# Whether to use GDN Pallas kernel +use_gdn_kernel: false +# Whether to use hybrid GDN v3 Tokamax forward + Custom VJP backward +use_hybrid_gdn: false # The ratio of dimension to apply ROPE on partial_rotary_factor: 1.0 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 6416d0ae9f..f66ba2659c 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1052,6 +1052,14 @@ class Qwen3Next(BaseModel): True, description="Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.", ) + use_gdn_kernel: bool = Field( + False, + description="Whether to use GDN Pallas kernel.", + ) + use_hybrid_gdn: bool = Field( + False, + description="Whether to use hybrid GDN v3 Tokamax forward + Custom VJP backward.", + ) partial_rotary_factor: float = Field(1.0, description="The ratio of dimension to apply ROPE on") diff --git a/src/maxtext/models/hybrid_gdn.py b/src/maxtext/models/hybrid_gdn.py new file mode 100644 index 0000000000..74fa38136a --- /dev/null +++ b/src/maxtext/models/hybrid_gdn.py @@ -0,0 +1,279 @@ +# Copyright 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 +# +# http://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. + +"""Hybrid Gated Delta Net (GDN) implementations for MaxText using Tokamax GDN v3 forward + Custom VJP backward.""" + +import functools +from typing import Any, Optional, Tuple + +import jax +import jax.numpy as jnp + + +def pure_jax_fused_conv1d_gdn( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + *, + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, +) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array]]: + """Pure-JAX composite of Conv1D + GDN used during backward pass autodiff.""" + from maxtext.models.qwen3 import jax_chunk_gated_delta_rule + batch, seq_len, _ = qkv.shape + key_dim = num_k_heads * head_k_dim + + # --- Step B: Pure JAX 1D Convolution --- + conv_input = jnp.pad(qkv, ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) + conv_weight_cast = conv_weight.astype(qkv.dtype) + conv_out = jax.lax.conv_general_dilated( + lhs=conv_input, + rhs=conv_weight_cast, + window_strides=(1,), + padding="VALID", + dimension_numbers=("NWC", "WIO", "NWC"), + feature_group_count=qkv.shape[-1], + ) + if conv_bias is not None: + conv_out = conv_out + conv_bias.astype(qkv.dtype) + conv_out = conv_out[:, -seq_len:, :] + qkv_conv = jax.nn.silu(conv_out.astype(jnp.float32)).astype(compute_dtype) + + q_conv, k_conv, v_conv = jnp.split(qkv_conv, [key_dim, 2 * key_dim], axis=-1) + + # Reshape for GDN + query = q_conv.reshape(batch, seq_len, num_k_heads, head_k_dim) + key = k_conv.reshape(batch, seq_len, num_k_heads, head_k_dim) + value = v_conv.reshape(batch, seq_len, num_v_heads, head_v_dim) + + A_log_cast = jnp.asarray(a_log, dtype=compute_dtype) + dt_bias_cast = jnp.asarray(dt_bias, dtype=compute_dtype) + beta = jax.nn.sigmoid(b) + g = -jnp.exp(A_log_cast) * jax.nn.softplus(a + dt_bias_cast) + + if num_v_heads > num_k_heads and num_v_heads % num_k_heads == 0: + repeats = num_v_heads // num_k_heads + query = jnp.repeat(query, repeats, axis=2) + key = jnp.repeat(key, repeats, axis=2) + + core_attn_out, next_recurrent_state = jax_chunk_gated_delta_rule( + query=query, + key=key, + value=value, + g=g, + beta=beta, + chunk_size=chunk_size, + initial_state=recurrent_state, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + ) + + next_conv_state = qkv[:, -(conv_kernel_size - 1):, :] if seq_len >= conv_kernel_size - 1 else jnp.zeros((batch, conv_kernel_size - 1, qkv.shape[-1]), dtype=qkv.dtype) + if next_recurrent_state is None: + next_recurrent_state = jnp.zeros((batch, num_v_heads, head_k_dim, head_v_dim), dtype=compute_dtype) + + return core_attn_out.astype(qkv.dtype), (next_conv_state.astype(qkv.dtype), next_recurrent_state.astype(qkv.dtype)) + + +def _run_tokamax_fused_fwd( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + *, + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, +): + if jax.default_backend() != "tpu": + return pure_jax_fused_conv1d_gdn( + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state, + num_k_heads=num_k_heads, num_v_heads=num_v_heads, head_k_dim=head_k_dim, head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, chunk_size=chunk_size, use_qk_norm_in_gdn=use_qk_norm_in_gdn, compute_dtype=compute_dtype, + ) + + # When on TPU, invoke Tokamax GDN v3 fused_conv1d_gdn kernel + from tokamax._src.ops.experimental.causal_conv1d_gated_delta_rule import wrapper as tokamax_gdn_wrapper + batch_size, seq_len, dim_size = qkv.shape + num_seqs = batch_size + + qkv_flat = qkv.reshape(-1, dim_size) + b_flat = b.reshape(-1, b.shape[-1]) + a_flat = a.reshape(-1, a.shape[-1]) + tokamax_conv_weight = jnp.swapaxes(conv_weight, 0, 2) + + query_start_loc = jnp.arange(0, (num_seqs + 1) * seq_len, seq_len, dtype=jnp.int32) + state_indices = jnp.arange(num_seqs, dtype=jnp.int32) + seq_lens = jnp.full((num_seqs,), seq_len, dtype=jnp.int32) + distribution = jnp.array([0, 0, num_seqs], dtype=jnp.int32) + + if conv_state is None: + tokamax_conv_state = jnp.zeros((num_seqs + 1, conv_kernel_size - 1, dim_size), dtype=qkv.dtype) + elif conv_state.shape[0] == num_seqs: + tokamax_conv_state = jnp.pad(conv_state, ((1, 0), (0, 0), (0, 0))) + else: + tokamax_conv_state = conv_state + + if recurrent_state is None: + tokamax_recurrent_state = jnp.zeros((num_seqs + 1, num_v_heads, head_k_dim, head_v_dim), dtype=qkv.dtype) + elif recurrent_state.shape[0] == num_seqs: + tokamax_recurrent_state = jnp.pad(recurrent_state, ((1, 0), (0, 0), (0, 0), (0, 0))) + else: + tokamax_recurrent_state = recurrent_state + + (new_conv_state, new_recurrent_state), core_attn_out_flat = tokamax_gdn_wrapper.fused_conv1d_gdn( + qkv=qkv_flat, + b=b_flat, + a=a_flat, + conv_state=tokamax_conv_state, + recurrent_state=tokamax_recurrent_state, + conv_weight=tokamax_conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + query_start_loc=query_start_loc, + state_indices=state_indices, + distribution=distribution, + seq_lens=seq_lens, + n_kq=num_k_heads, + n_v=num_v_heads, + d_k=head_k_dim, + d_v=head_v_dim, + kernel_size=conv_kernel_size, + compute_precision=jnp.dtype(jnp.float32), + ) + + core_attn_out = core_attn_out_flat.reshape(batch_size, seq_len, num_v_heads, head_v_dim) + return core_attn_out.astype(qkv.dtype), (new_conv_state[1:].astype(qkv.dtype), new_recurrent_state[1:].astype(qkv.dtype)) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(9, 10, 11, 12, 13, 14, 15, 16)) +def hybrid_fused_conv1d_gdn( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, +) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array]]: + """Hybrid Fused Conv1D + GDN: Tokamax GDN v3 forward + Custom VJP backward.""" + return _run_tokamax_fused_fwd( + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state, + num_k_heads=num_k_heads, num_v_heads=num_v_heads, head_k_dim=head_k_dim, head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, chunk_size=chunk_size, use_qk_norm_in_gdn=use_qk_norm_in_gdn, compute_dtype=compute_dtype, + ) + + +def _hybrid_fused_conv1d_gdn_fwd( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, +): + out, states = _run_tokamax_fused_fwd( + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state, + num_k_heads=num_k_heads, num_v_heads=num_v_heads, head_k_dim=head_k_dim, head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, chunk_size=chunk_size, use_qk_norm_in_gdn=use_qk_norm_in_gdn, compute_dtype=compute_dtype, + ) + residuals = ( + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state + ) + return (out, states), residuals + + +def _hybrid_fused_conv1d_gdn_bwd( + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, + residuals: tuple, + cotangents: tuple, +): + ( + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state + ) = residuals + + def target_fn(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_): + return pure_jax_fused_conv1d_gdn( + qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + ) + + _, vjp_fn = jax.vjp( + target_fn, + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state, + ) + d_out, d_states = cotangents + d_conv_state, d_recurrent_state = d_states + return vjp_fn((d_out, (d_conv_state, d_recurrent_state))) + + +hybrid_fused_conv1d_gdn.defvjp(_hybrid_fused_conv1d_gdn_fwd, _hybrid_fused_conv1d_gdn_bwd) diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index e16d7b5e49..dae15976c8 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -846,6 +846,116 @@ def extract_state(c_in, v_len): use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, compute_dtype=cfg.dtype, ) + elif getattr(cfg, "use_gdn_kernel", False) and getattr(cfg, "use_hybrid_gdn", False): + from maxtext.models.hybrid_gdn import hybrid_fused_conv1d_gdn + + if self.mesh is not None: + logical_rules = get_logical_axis_rules() + batch_pspec3 = logical_to_mesh_axes((KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules) + batch_pspec4 = logical_to_mesh_axes((KV_BATCH, None, None, None), mesh=self.mesh, rules=logical_rules) + none_pspec3 = logical_to_mesh_axes((None, None, None), mesh=self.mesh, rules=logical_rules) + none_pspec1 = logical_to_mesh_axes((None,), mesh=self.mesh, rules=logical_rules) + + recurrent_state_arg = ( + recurrent_state + if recurrent_state is not None + else jnp.zeros((batch, self.num_v_heads, self.head_k_dim, self.head_v_dim), dtype=cfg.dtype) + ) + conv_state_arg = ( + conv_state + if conv_state is not None + else jnp.zeros((batch, self.config.gdn_conv_kernel_dim - 1, qkv.shape[-1]), dtype=cfg.dtype) + ) + conv_bias_arg = ( + self.conv1d.bias.value + if hasattr(self.conv1d, "bias") and self.conv1d.bias is not None + else jnp.zeros((qkv.shape[-1],), dtype=cfg.dtype) + ) + + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=( + batch_pspec3, # qkv + batch_pspec3, # b + batch_pspec3, # a + none_pspec3, # conv_weight + none_pspec1, # conv_bias + none_pspec1, # a_log + none_pspec1, # dt_bias + batch_pspec3, # conv_state + batch_pspec4, # recurrent_state + ), + out_specs=( + batch_pspec4, # core_attn_out + (batch_pspec3, batch_pspec4), # (next_conv_state, next_recurrent_state) + ), + check_vma=False, + ) + def shard_mapped_hybrid_gdn(qkv_val, b_val, a_val, cw_val, cb_val, alog_val, dt_val, cs_val, rs_val): + return hybrid_fused_conv1d_gdn( + qkv=qkv_val, + b=b_val, + a=a_val, + conv_weight=cw_val, + conv_bias=cb_val, + a_log=alog_val, + dt_bias=dt_val, + conv_state=cs_val, + recurrent_state=rs_val, + num_k_heads=self.num_k_heads, + num_v_heads=self.num_v_heads, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + conv_kernel_size=self.config.gdn_conv_kernel_dim, + chunk_size=self.config.gdn_chunk_size, + use_qk_norm_in_gdn=self.config.use_qk_norm_in_gdn, + compute_dtype=self.config.dtype, + ) + + core_attn_out, (next_conv_state, next_recurrent_state) = shard_mapped_hybrid_gdn( + qkv, + b, + a, + self.conv1d.kernel.value, + conv_bias_arg, + self.A_log[...], + self.dt_bias[...], + conv_state_arg, + recurrent_state_arg, + ) + else: + core_attn_out, (next_conv_state, next_recurrent_state) = hybrid_fused_conv1d_gdn( + qkv=qkv, + b=b, + a=a, + conv_weight=self.conv1d.kernel.value, + conv_bias=None, + a_log=self.A_log[...], + dt_bias=self.dt_bias[...], + conv_state=conv_state, + recurrent_state=recurrent_state, + num_k_heads=self.num_k_heads, + num_v_heads=self.num_v_heads, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + conv_kernel_size=self.config.gdn_conv_kernel_dim, + chunk_size=self.config.gdn_chunk_size, + use_qk_norm_in_gdn=self.config.use_qk_norm_in_gdn, + compute_dtype=self.config.dtype, + ) + elif getattr(cfg, "use_gdn_kernel", False): + core_attn_out, next_recurrent_state = jax_chunk_gated_delta_rule( + query, + key, + value, + g, + beta, + chunk_size=cfg.gdn_chunk_size, + initial_state=recurrent_state, + use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, + compute_dtype=cfg.dtype, + ) elif self.mesh is not None: logical_rules = get_logical_axis_rules() recurrent_state_arg = ( From 78d09f46d62be4508cb4296e385b4703e6f6cb92 Mon Sep 17 00:00:00 2001 From: darisoy Date: Wed, 5 Aug 2026 18:42:14 +0000 Subject: [PATCH 05/12] feat: implement Muon optimizer support for Qwen3-Next architecture (RoutedMoE and GatedDeltaNet) --- src/maxtext/configs/types.py | 4 ++++ src/maxtext/utils/muon_utils.py | 14 +++++++++++++- tests/unit/muon_utils_test.py | 22 ++++++++++++++++++++++ 3 files changed, 39 insertions(+), 1 deletion(-) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index f66ba2659c..23e6e2b52d 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -3777,6 +3777,10 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4, DecoderBlockType.QWEN3, + DecoderBlockType.QWEN3_NEXT, + DecoderBlockType.QWEN3_MOE, + DecoderBlockType.QWEN3_5, + DecoderBlockType.QWEN3_CUSTOM_MOE, DecoderBlockType.GEMMA3, DecoderBlockType.LLAMA2, ]: diff --git a/src/maxtext/utils/muon_utils.py b/src/maxtext/utils/muon_utils.py index ff77c57807..9264a7f2a3 100644 --- a/src/maxtext/utils/muon_utils.py +++ b/src/maxtext/utils/muon_utils.py @@ -91,6 +91,11 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]: "hc_base", "sinks", "tid2eid", + "A_log", + "dt_bias", + "conv1d", + "gate", + "shared_expert_gate", ) ) or segment == "bias" @@ -101,7 +106,7 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]: # 2 Special weights # 2.1 Special weights: MoE, [0, L, -2, -1] # L (optional) stands for layer when scan_layers=True - if "MoeBlock_0" in path: + if _is_path_contain_any(("MoeBlock_0", "routed_experts"), path): # exclude gate if _is_path_contain_any(("wi_0", "wi_1", "wo"), path): return mdn((-2,), (-1,)) @@ -119,6 +124,13 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]: elif _is_path_contain_any(("query", "key", "value", "wq_b", "wkv_b", "wkv"), path): return mdn((0,), (-2, -1)) + # 2.3 Special weights: Gated Delta Net (GDN) + elif "gdn" in path: + if "out_proj" in path: + return mdn((0, -2), (-1,)) + elif _is_path_contain_any(("in_proj_qkvz", "in_proj_ba"), path): + return mdn((0,), (-2, -1)) + # 3 Standard weights, [0, L, -1] return mdn((0,), (-1,)) diff --git a/tests/unit/muon_utils_test.py b/tests/unit/muon_utils_test.py index 58bfadf29a..a45f910ecf 100644 --- a/tests/unit/muon_utils_test.py +++ b/tests/unit/muon_utils_test.py @@ -19,6 +19,7 @@ import io import contextlib import unittest +import pytest from unittest import mock import jax @@ -119,6 +120,27 @@ def test_deepseek_v4_self_attention_grouped_projection(self): # and output on out_features_per_group (-1) self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "o_a_proj")), mdn((-2,), (-1,))) + # --- 5. Qwen3-Next Specific --- + @pytest.mark.tpu_only + def test_qwen3_next_moe_routed_experts(self): + self.assertEqual(muon_utils.transform_logic(("decoder", "mlp", "routed_experts", "wi_0")), mdn((-2,), (-1,))) + self.assertEqual(muon_utils.transform_logic(("decoder", "mlp", "routed_experts", "wi_1")), mdn((-2,), (-1,))) + self.assertEqual(muon_utils.transform_logic(("decoder", "mlp", "routed_experts", "wo")), mdn((-2,), (-1,))) + + @pytest.mark.tpu_only + def test_qwen3_next_gdn_projections(self): + self.assertEqual(muon_utils.transform_logic(("decoder", "gdn", "in_proj_qkvz")), mdn((0,), (-2, -1))) + self.assertEqual(muon_utils.transform_logic(("decoder", "gdn", "in_proj_ba")), mdn((0,), (-2, -1))) + self.assertEqual(muon_utils.transform_logic(("decoder", "gdn", "out_proj")), mdn((0, -2), (-1,))) + + @pytest.mark.tpu_only + def test_qwen3_next_exclusions(self): + self.assertIsNone(muon_utils.transform_logic(("decoder", "gdn", "A_log"))) + self.assertIsNone(muon_utils.transform_logic(("decoder", "gdn", "dt_bias"))) + self.assertIsNone(muon_utils.transform_logic(("decoder", "gdn", "conv1d"))) + self.assertIsNone(muon_utils.transform_logic(("decoder", "mlp", "routed_experts", "gate"))) + self.assertIsNone(muon_utils.transform_logic(("decoder", "mlp", "shared_expert_gate"))) + class TestGetTransformTree(unittest.TestCase): """Tests for get_transform_tree: recursive dict walk that applies transform_logic.""" From 1476823e9a1f0306825c76c7742d9c304df66df0 Mon Sep 17 00:00:00 2001 From: Nuojin Cheng Date: Thu, 30 Jul 2026 20:03:28 +0000 Subject: [PATCH 06/12] add scan support for qwen3-next model --- .../utils/param_mapping.py | 385 +++++++++++------- src/maxtext/layers/decoders.py | 129 ++++++ src/maxtext/layers/nnx_decoders.py | 112 +++++ src/maxtext/models/qwen3.py | 247 ++++++++--- tests/unit/nnx_decoders_test.py | 221 +++++++++- tests/unit/param_mapping_test.py | 9 +- 6 files changed, 891 insertions(+), 212 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 26359cdddc..f7dfa5e872 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -1268,7 +1268,7 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel"] = transpose hooks[f"{mlp_prefix}-shared_expert-wo-kernel"] = transpose hooks[f"{mlp_prefix}-shared_expert_gate-kernel"] = transpose - # pyrefly: ignore[unsupported-operation] + hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = ( process_wi_0_wi_1 # pyrefly: ignore[unsupported-operation] ) @@ -1386,100 +1386,225 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F } if scan_layers: - # 2. Scan over block cycles - for block_idx in range(layer_cycle_interval): - hf_indices = list(range(block_idx, num_main_layers, layer_cycle_interval)) - prefix = f"params-decoder-layers-layer_{block_idx}" + num_blocks = num_main_layers // layer_cycle_interval + num_scanned = num_blocks * layer_cycle_interval + num_remaining = num_main_layers % layer_cycle_interval - # Layer norms - mapping[f"{prefix}-input_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] - f"model.layers.{i}.input_layernorm.weight" for i in hf_indices - ] # pyrefly: ignore[bad-assignment] - mapping[f"{prefix}-post_attention_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] - f"model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices - ] + def hf_layer(idx, suffix): + return f"model.layers.{idx}.{suffix}" - # Handle Interleaved Attention (Linear vs Full) - is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0 + local_prefix = "params-decoder-scanned_blocks-local_layers" + local_positions = list(range(layer_cycle_interval - 1)) - if is_full_attention_layer: - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-attention-attention-query-kernel": [ - f"model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-key-kernel": [ - f"model.layers.{i}.self_attn.k_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-value-kernel": [ - f"model.layers.{i}.self_attn.v_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-out-kernel": [ - f"model.layers.{i}.self_attn.o_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-query_norm-scale": [ - f"model.layers.{i}.self_attn.q_norm.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-key_norm-scale": [ - f"model.layers.{i}.self_attn.k_norm.weight" for i in hf_indices - ], - } + # Local / linear attention layers (nested [block][local]) + mapping.update( + { + f"{local_prefix}-input_layernorm-scale": [ + [hf_layer(b * layer_cycle_interval + l, "input_layernorm.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-post_attention_layernorm-scale": [ + [hf_layer(b * layer_cycle_interval + l, "post_attention_layernorm.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-in_proj_qkvz-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.in_proj_qkvz.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-in_proj_ba-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.in_proj_ba.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-conv1d-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.conv1d.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-A_log": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.A_log") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-dt_bias": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.dt_bias") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-norm-rms_norm-scale": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.norm.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-attention-out_proj-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "linear_attn.out_proj.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-mlp-routed_experts-gate-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "mlp.gate.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-mlp-shared_expert-wi_0-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "mlp.shared_expert.gate_proj.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-mlp-shared_expert-wi_1-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "mlp.shared_expert.up_proj.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-mlp-shared_expert-wo-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "mlp.shared_expert.down_proj.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-mlp-shared_expert_gate-kernel": [ + [hf_layer(b * layer_cycle_interval + l, "mlp.shared_expert_gate.weight") for l in local_positions] + for b in range(num_blocks) + ], + f"{local_prefix}-mlp-routed_experts-wi_0": [ + [ + [hf_layer(b * layer_cycle_interval + l, f"mlp.experts.{e}.gate_proj.weight") for l in local_positions] + for b in range(num_blocks) + ] + for e in range(num_experts) + ], + f"{local_prefix}-mlp-routed_experts-wi_1": [ + [ + [hf_layer(b * layer_cycle_interval + l, f"mlp.experts.{e}.up_proj.weight") for l in local_positions] + for b in range(num_blocks) + ] + for e in range(num_experts) + ], + f"{local_prefix}-mlp-routed_experts-wo": [ + [ + [hf_layer(b * layer_cycle_interval + l, f"mlp.experts.{e}.down_proj.weight") for l in local_positions] + for b in range(num_blocks) + ] + for e in range(num_experts) + ], + } + ) + + global_prefix = "params-decoder-scanned_blocks-global_layer" + global_position = layer_cycle_interval - 1 + + # Global attention layer (flat over blocks) + mapping.update( + { + f"{global_prefix}-input_layernorm-scale": [ + hf_layer(b * layer_cycle_interval + global_position, "input_layernorm.weight") for b in range(num_blocks) + ], + f"{global_prefix}-post_attention_layernorm-scale": [ + hf_layer(b * layer_cycle_interval + global_position, "post_attention_layernorm.weight") + for b in range(num_blocks) + ], + f"{global_prefix}-attention-attention-query-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "self_attn.q_proj.weight") for b in range(num_blocks) + ], + f"{global_prefix}-attention-attention-key-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "self_attn.k_proj.weight") for b in range(num_blocks) + ], + f"{global_prefix}-attention-attention-value-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "self_attn.v_proj.weight") for b in range(num_blocks) + ], + f"{global_prefix}-attention-attention-out-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "self_attn.o_proj.weight") for b in range(num_blocks) + ], + f"{global_prefix}-attention-attention-query_norm-scale": [ + hf_layer(b * layer_cycle_interval + global_position, "self_attn.q_norm.weight") for b in range(num_blocks) + ], + f"{global_prefix}-attention-attention-key_norm-scale": [ + hf_layer(b * layer_cycle_interval + global_position, "self_attn.k_norm.weight") for b in range(num_blocks) + ], + f"{global_prefix}-mlp-routed_experts-gate-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "mlp.gate.weight") for b in range(num_blocks) + ], + f"{global_prefix}-mlp-shared_expert-wi_0-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "mlp.shared_expert.gate_proj.weight") + for b in range(num_blocks) + ], + f"{global_prefix}-mlp-shared_expert-wi_1-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "mlp.shared_expert.up_proj.weight") + for b in range(num_blocks) + ], + f"{global_prefix}-mlp-shared_expert-wo-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "mlp.shared_expert.down_proj.weight") + for b in range(num_blocks) + ], + f"{global_prefix}-mlp-shared_expert_gate-kernel": [ + hf_layer(b * layer_cycle_interval + global_position, "mlp.shared_expert_gate.weight") + for b in range(num_blocks) + ], + f"{global_prefix}-mlp-routed_experts-wi_0": [ + [ + hf_layer(b * layer_cycle_interval + global_position, f"mlp.experts.{e}.gate_proj.weight") + for b in range(num_blocks) + ] + for e in range(num_experts) + ], + f"{global_prefix}-mlp-routed_experts-wi_1": [ + [ + hf_layer(b * layer_cycle_interval + global_position, f"mlp.experts.{e}.up_proj.weight") + for b in range(num_blocks) + ] + for e in range(num_experts) + ], + f"{global_prefix}-mlp-routed_experts-wo": [ + [ + hf_layer(b * layer_cycle_interval + global_position, f"mlp.experts.{e}.down_proj.weight") + for b in range(num_blocks) + ] + for e in range(num_experts) + ], + } + ) + + # Remainder layers if any + if num_remaining > 0: + for rem_idx in range(num_remaining): + hf_layer_idx = num_scanned + rem_idx + prefix = f"params-decoder-layers_{hf_layer_idx}" + layer_in_block = rem_idx % layer_cycle_interval + is_full_attention_layer = (layer_in_block + 1) % layer_cycle_interval == 0 + mapping[f"{prefix}-input_layernorm-scale"] = f"model.layers.{hf_layer_idx}.input_layernorm.weight" + mapping[f"{prefix}-post_attention_layernorm-scale"] = ( + f"model.layers.{hf_layer_idx}.post_attention_layernorm.weight" ) - else: - # Linear/Hybrid Attention Block - mapping.update( # pyrefly: ignore[no-matching-overload] + if is_full_attention_layer: + mapping.update( + { + f"{prefix}-attention-attention-query-kernel": f"model.layers.{hf_layer_idx}.self_attn.q_proj.weight", + f"{prefix}-attention-attention-key-kernel": f"model.layers.{hf_layer_idx}.self_attn.k_proj.weight", + f"{prefix}-attention-attention-value-kernel": f"model.layers.{hf_layer_idx}.self_attn.v_proj.weight", + f"{prefix}-attention-attention-out-kernel": f"model.layers.{hf_layer_idx}.self_attn.o_proj.weight", + f"{prefix}-attention-attention-query_norm-scale": f"model.layers.{hf_layer_idx}.self_attn.q_norm.weight", + f"{prefix}-attention-attention-key_norm-scale": f"model.layers.{hf_layer_idx}.self_attn.k_norm.weight", + } + ) + else: + mapping.update( + { + f"{prefix}-attention-in_proj_qkvz-kernel": f"model.layers.{hf_layer_idx}.linear_attn.in_proj_qkvz.weight", + f"{prefix}-attention-in_proj_ba-kernel": f"model.layers.{hf_layer_idx}.linear_attn.in_proj_ba.weight", + f"{prefix}-attention-conv1d-kernel": f"model.layers.{hf_layer_idx}.linear_attn.conv1d.weight", + f"{prefix}-attention-A_log": f"model.layers.{hf_layer_idx}.linear_attn.A_log", + f"{prefix}-attention-dt_bias": f"model.layers.{hf_layer_idx}.linear_attn.dt_bias", + f"{prefix}-attention-norm-rms_norm-scale": f"model.layers.{hf_layer_idx}.linear_attn.norm.weight", + f"{prefix}-attention-out_proj-kernel": f"model.layers.{hf_layer_idx}.linear_attn.out_proj.weight", + } + ) + mapping.update( { - f"{prefix}-attention-in_proj_qkvz-kernel": [ - f"model.layers.{i}.linear_attn.in_proj_qkvz.weight" for i in hf_indices + f"{prefix}-mlp-routed_experts-gate-kernel": f"model.layers.{hf_layer_idx}.mlp.gate.weight", + f"{prefix}-mlp-shared_expert-wi_0-kernel": f"model.layers.{hf_layer_idx}.mlp.shared_expert.gate_proj.weight", + f"{prefix}-mlp-shared_expert-wi_1-kernel": f"model.layers.{hf_layer_idx}.mlp.shared_expert.up_proj.weight", + f"{prefix}-mlp-shared_expert-wo-kernel": f"model.layers.{hf_layer_idx}.mlp.shared_expert.down_proj.weight", + f"{prefix}-mlp-shared_expert_gate-kernel": f"model.layers.{hf_layer_idx}.mlp.shared_expert_gate.weight", + f"{prefix}-mlp-routed_experts-wi_0": [ + f"model.layers.{hf_layer_idx}.mlp.experts.{e}.gate_proj.weight" for e in range(num_experts) ], - f"{prefix}-attention-in_proj_ba-kernel": [ - f"model.layers.{i}.linear_attn.in_proj_ba.weight" for i in hf_indices + f"{prefix}-mlp-routed_experts-wi_1": [ + f"model.layers.{hf_layer_idx}.mlp.experts.{e}.up_proj.weight" for e in range(num_experts) ], - f"{prefix}-attention-conv1d-kernel": [f"model.layers.{i}.linear_attn.conv1d.weight" for i in hf_indices], - f"{prefix}-attention-A_log": [f"model.layers.{i}.linear_attn.A_log" for i in hf_indices], - f"{prefix}-attention-dt_bias": [f"model.layers.{i}.linear_attn.dt_bias" for i in hf_indices], - f"{prefix}-attention-norm-rms_norm-scale": [ - f"model.layers.{i}.linear_attn.norm.weight" for i in hf_indices - ], - f"{prefix}-attention-out_proj-kernel": [ - f"model.layers.{i}.linear_attn.out_proj.weight" for i in hf_indices + f"{prefix}-mlp-routed_experts-wo": [ + f"model.layers.{hf_layer_idx}.mlp.experts.{e}.down_proj.weight" for e in range(num_experts) ], } ) - - # 3. Handle MLP: Gates and Shared Experts - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-mlp-routed_experts-gate-kernel": [f"model.layers.{i}.mlp.gate.weight" for i in hf_indices], - f"{prefix}-mlp-shared_expert-wi_0-kernel": [ - f"model.layers.{i}.mlp.shared_expert.gate_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wi_1-kernel": [ - f"model.layers.{i}.mlp.shared_expert.up_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wo-kernel": [ - f"model.layers.{i}.mlp.shared_expert.down_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert_gate-kernel": [ - f"model.layers.{i}.mlp.shared_expert_gate.weight" for i in hf_indices - ], - } - ) - - # 4. Handle MoE Routed Experts - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-mlp-routed_experts-wi_0": [ - [f"model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for i in hf_indices] for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wi_1": [ - [f"model.layers.{i}.mlp.experts.{e}.up_proj.weight" for i in hf_indices] for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wo": [ - [f"model.layers.{i}.mlp.experts.{e}.down_proj.weight" for i in hf_indices] for e in range(num_experts) - ], - } - ) else: # Unscanned layer mapping for i in range(num_main_layers): @@ -1571,20 +1696,8 @@ def permute_conv(input_tensor, target_shape=None): "params-decoder-logits_dense-kernel": transpose, } - layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval - num_main_layers = config["num_hidden_layers"] - loop_indices = range(layer_cycle_interval) if scan_layers else range(num_main_layers) - - for i in loop_indices: - if scan_layers: - prefix = f"params-decoder-layers-layer_{i}" - block_idx = i - else: - prefix = f"params-decoder-layers_{i}" - block_idx = i % layer_cycle_interval - is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0 - - if is_full_attention_layer: + def _attach_block_hooks(prefix, is_global): + if is_global: for key in ["query", "key", "value", "out"]: hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment] else: @@ -1604,6 +1717,16 @@ def permute_conv(input_tensor, target_shape=None): hooks[f"{mlp_prefix}-routed_experts-wi_1"] = transpose hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose + if scan_layers: + _attach_block_hooks("params-decoder-scanned_blocks-local_layers", is_global=False) + _attach_block_hooks("params-decoder-scanned_blocks-global_layer", is_global=True) + else: + for i in range(config.base_num_decoder_layers): + prefix = f"params-decoder-layers_{i}" + block_idx = i % config.inhomogeneous_layer_cycle_interval + is_full_attention_layer = (block_idx + 1) % config.inhomogeneous_layer_cycle_interval == 0 + _attach_block_hooks(prefix, is_global=is_full_attention_layer) + return hooks @@ -1953,7 +2076,6 @@ def interleave(input_tensor, target_shape=None): hooks[f"{prefix}-GptOssMlp-gate-kernel"] = transpose # `composite_mt_key`: A hook for combining multiple MaxText params. hooks[(f"{prefix}-GptOssMlp-wi_0", f"{prefix}-GptOssMlp-wi_1")] = interleave # pyrefly: ignore[unsupported-operation] - # pyrefly: ignore[unsupported-operation] hooks[(f"{prefix}-GptOssMlp-wi_0_bias", f"{prefix}-GptOssMlp-wi_1_bias")] = ( interleave # pyrefly: ignore[unsupported-operation] ) @@ -2875,10 +2997,10 @@ def _spec_active(gate): local_positions = list(range(attention_pattern_length - 1)) for subkey, suffix, gate in param_specs: if _spec_active(gate): - mapping[f"{local_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + mapping[f"{local_prefix}-{subkey}"] = [ # pyrefly: ignore[no-matching-overload] [hf_layer(b * attention_pattern_length + l, suffix) for l in local_positions] for b in range(num_blocks) ] - mapping[f"{local_prefix}-self_attention-value-kernel"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + mapping[f"{local_prefix}-self_attention-value-kernel"] = [ # pyrefly: ignore[no-matching-overload] [hf_layer(b * attention_pattern_length + l, "self_attn.v_proj.weight") for l in local_positions] for b in range(num_blocks) ] @@ -2888,11 +3010,11 @@ def _spec_active(gate): global_position = attention_pattern_length - 1 for subkey, suffix, gate in param_specs: if _spec_active(gate): - mapping[f"{global_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + mapping[f"{global_prefix}-{subkey}"] = [ # pyrefly: ignore[no-matching-overload] hf_layer(b * attention_pattern_length + global_position, suffix) for b in range(num_blocks) ] if not share_kv_projections: - mapping[f"{global_prefix}-self_attention-value-kernel"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + mapping[f"{global_prefix}-self_attention-value-kernel"] = [ # pyrefly: ignore[no-matching-overload] hf_layer(b * attention_pattern_length + global_position, "self_attn.v_proj.weight") for b in range(num_blocks) ] @@ -3646,9 +3768,8 @@ def QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fal mapping = {} n_layers_text = config["text_config"]["num_hidden_layers"] - num_experts_text = config["text_config"].get("num_local_experts", 0) text_mapping = QWEN_MAXTEXT_TO_HF_PARAM_MAPPING( - config={"num_hidden_layers": n_layers_text, "num_experts": num_experts_text}, + config={"num_hidden_layers": n_layers_text}, maxtext_config=maxtext_config, scan_layers=scan_layers, ) @@ -3663,22 +3784,6 @@ def replace_prefix(val): for key, value in text_mapping.items(): mapping[key] = replace_prefix(value) - # For correct layer replication vs scanning. - if num_experts_text > 0 and not scan_layers: - for i in range(n_layers_text): - key0 = f"params-decoder-layers_{i}-moe_block-wi_0" - key1 = f"params-decoder-layers_{i}-moe_block-wi_1" - key_wo = f"params-decoder-layers_{i}-moe_block-wo" - - if key0 in mapping: - del mapping[key0] - if key1 in mapping: - del mapping[key1] - - composite_key = (key0, key1) - mapping[composite_key] = f"model.language_model.layers.{i}.mlp.experts.gate_up_proj" - mapping[key_wo] = f"model.language_model.layers.{i}.mlp.experts.down_proj" - vision_config = config["vision_config"] n_vision_layers = vision_config["depth"] @@ -3747,36 +3852,14 @@ def QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=Fal mapping = {} n_layers_text = config["text_config"]["num_hidden_layers"] - num_experts_text = config["text_config"].get("num_local_experts", 0) text_hooks = QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN( - config={"num_hidden_layers": n_layers_text, "num_experts": num_experts_text}, + config={"num_hidden_layers": n_layers_text}, maxtext_config=maxtext_config, scan_layers=scan_layers, saving_to_hf=saving_to_hf, ) - - def process_wi_0_wi_1_fused(input_tensor, target_shape=None): - if saving_to_hf: - wi_0, wi_1 = input_tensor - gate_up = np.concatenate([wi_0, wi_1], axis=-1) - return gate_up - else: - gate, up = np.split(input_tensor, 2, axis=-1) - return np.stack([gate, up], axis=-1) - mapping.update(text_hooks) - # Special case for Qwen3-VL: MaxText model has separate wi_0 and wi_1 weights - # for the MoE block, but the HF model expects a single fused weight. - if num_experts_text > 0 and not scan_layers: - for i in range(n_layers_text): - composite_key = ( - f"params-decoder-layers_{i}-moe_block-wi_0", - f"params-decoder-layers_{i}-moe_block-wi_1", - ) - mapping[composite_key] = process_wi_0_wi_1_fused - mapping[f"params-decoder-layers_{i}-moe_block-wo"] = None - vision_config = config["vision_config"] n_vision_layers = vision_config["depth"] hidden_size = vision_config["hidden_size"] @@ -4015,7 +4098,7 @@ def get_hf_expert_keys(expert_subpath_template): } ) - mapping.update(layer_map) # pyrefly: ignore[no-matching-overload] + mapping.update(layer_map) if not scan_layers: for i in range(n_layers): @@ -4039,7 +4122,7 @@ def transpose(input_tensor, target_shape=None): return np.transpose(input_tensor) def ones_norm(input_tensor, target_shape=None): - return np.ones(target_shape, dtype=np.float32) # pyrefly: ignore[no-matching-overload] + return np.ones(target_shape, dtype=np.float32) def identity(input_tensor, target_shape=None): return input_tensor @@ -4070,9 +4153,9 @@ def reshape_transpose_o_a(input_tensor, target_shape=None): if saving_to_hf: tensor = np.transpose(input_tensor, (0, 2, 1)) return tensor.reshape(target_shape) - num_heads = target_shape[0] # pyrefly: ignore[unsupported-operation] - embed_dim = target_shape[1] # pyrefly: ignore[unsupported-operation] - kv_lora_rank = target_shape[2] # pyrefly: ignore[unsupported-operation] + num_heads = target_shape[0] + embed_dim = target_shape[1] + kv_lora_rank = target_shape[2] tensor = input_tensor.reshape((num_heads, kv_lora_rank, embed_dim)) return np.transpose(tensor, (0, 2, 1)) @@ -4228,7 +4311,6 @@ def mhc_concat_scale(input_tensors, target_shape=None): "qwen3-32b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3-vl-2b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3-vl-4b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING, - "qwen3-vl-30b-a3b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING, "llama3.1-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, "llama3.1-8b-Instruct": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, "llama3.1-70b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -4282,7 +4364,6 @@ def mhc_concat_scale(input_tensors, target_shape=None): "qwen3-32b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN, "qwen3-vl-2b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN, "qwen3-vl-4b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN, - "qwen3-vl-30b-a3b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN, "llama3.1-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, "llama3.1-8b-Instruct": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, "llama3.1-70b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 42753eb752..bcaa270a3b 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -1070,6 +1070,18 @@ def __call__( kv_caches=kv_caches, attention_metadata=attention_metadata, ) + elif cfg.decoder_block == DecoderBlockType.QWEN3_NEXT: + y = self._apply_qwen3_next_scanned_blocks( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + ) elif cfg.decoder_block == DecoderBlockType.DEEPSEEK4: y = self._apply_deepseek4_scanned_blocks( y, @@ -1422,6 +1434,123 @@ def _apply_gemma3_scanned_blocks( return y + def _apply_qwen3_next_scanned_blocks( + self, + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + kv_caches=None, + attention_metadata=None, + ): + """Applies Qwen3-Next scanned decoder blocks, handling main scan and remainders.""" + + cfg = self.config + mesh = self.mesh + + # Define the repeating pattern length and calculate how many full blocks to scan + block_pattern_len = cfg.inhomogeneous_layer_cycle_interval + num_full_blocks = cfg.num_decoder_layers // block_pattern_len + remainder_layers = cfg.num_decoder_layers % block_pattern_len + + if num_full_blocks > 0: + ScannableBlockToLinen = qwen3.Qwen3NextScannableBlockToLinen + policy = self.get_remat_policy() + + kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan( + kv_caches, num_full_blocks, block_pattern_len, stack=True + ) + + broadcast_args_spec = [ + (decoder_segment_ids, nn.broadcast), + (decoder_positions, nn.broadcast), + (deterministic, nn.broadcast), + (model_mode, nn.broadcast), + (slot, nn.broadcast), + (None, nn.broadcast), # page_state + (previous_chunk, nn.broadcast), + (None, nn.broadcast), # bidirectional_mask + (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast), + (attention_metadata, nn.broadcast), + ] + broadcast_args = tuple(arg for arg, _ in broadcast_args_spec) + in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec) + + # For a fully scanned block, apply it inside an nn.scan over the calculated number of full blocks + y, returned_kv_cache = nn.scan( + ScannableBlockToLinen, + variable_axes={ + "params": cfg.param_scan_axis, + "cache": 0, + "intermediates": 0, + "aqt": 0, + "_overwrite_with_gradient": 0, + }, + split_rngs={"params": True, "dropout": cfg.enable_dropout}, + in_axes=in_axes_tuple, + length=num_full_blocks, + unroll=num_full_blocks, + metadata_params={ + nn.PARTITION_NAME: "layers", + "abstract_init": False, + }, + )( + config=cfg, + mesh=mesh, + quant=self.quant, + model_mode=model_mode, + num_of_layers=block_pattern_len, + remat_policy_fn=policy, + apply_internal_remat=True, + name="scanned_blocks", + )( + y, *broadcast_args + ) + + maxtext_utils.update_kv_caches_after_scan( + kv_caches, returned_kv_cache, num_full_blocks, block_pattern_len, stacked=True + ) + + # Process any remaining layers that don't fit into a full scanned block + for layer_id in range(cfg.num_decoder_layers - remainder_layers, cfg.num_decoder_layers): + layer = qwen3.Qwen3NextDecoderLayerToLinen( + config=cfg, + mesh=mesh, + model_mode=model_mode, + quant=self.quant, + layer_idx=layer_id, + ) + kv_cache = kv_caches[layer_id] if kv_caches is not None else None + + remainder_args = ( + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + None, # page_state + slot, + None, # bidirectional_mask + kv_cache, + attention_metadata, + ) + + y_and_kv = layer(y, *remainder_args) + if isinstance(y_and_kv, tuple): + y = y_and_kv[0] + new_kv = y_and_kv[1] + else: + y = y_and_kv + new_kv = None + + if kv_caches is not None and new_kv is not None: + kv_caches[layer_id] = new_kv + + return y + def _apply_gemma4_scanned_blocks( self, y, diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index ae8ee84c1d..b114c1f5a2 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -437,6 +437,7 @@ def __init__( self.is_gemma3 = self.config.decoder_block == DecoderBlockType.GEMMA3 self.is_gemma4 = self.config.decoder_block == DecoderBlockType.GEMMA4 self.is_gemma4_small = self.config.decoder_block == DecoderBlockType.GEMMA4_SMALL + self.is_qwen3_next = self.config.decoder_block == DecoderBlockType.QWEN3_NEXT if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4: self.hc_head = mhc.DeepSeek4HyperHead( @@ -547,6 +548,8 @@ def _init_scanned_layers(self, decoder_block_classes, rngs, mesh): self._init_scanned_gemma3(decoder_block_classes, rngs, mesh) elif self.is_gemma4: self._init_scanned_gemma4(decoder_block_classes, rngs, mesh) + elif self.is_qwen3_next: + self._init_scanned_qwen3_next(decoder_block_classes, rngs, mesh) else: self._init_scanned_generic(decoder_block_classes, rngs) @@ -717,6 +720,43 @@ def _init_scanned_gemma4(self, decoder_block_classes, rngs, mesh): rngs=rngs, ) + def _init_scanned_qwen3_next(self, decoder_block_classes, rngs, mesh): + """Initializes scanned Qwen3-Next layers.""" + config = self.config + cycle_interval = config.inhomogeneous_layer_cycle_interval + scan_length = config.num_decoder_layers // cycle_interval + num_remaining_layers = config.num_decoder_layers % cycle_interval + policy = self.get_remat_policy() + layer_kwargs = { + "num_of_layers": cycle_interval, + "remat_policy_fn": policy, + "apply_internal_remat": True, + } + rem_layer_kwargs = { + "num_of_layers": num_remaining_layers, + "remat_policy_fn": policy, + "apply_internal_remat": True, + } + + RemattedQwen3NextBlock = qwen3.Qwen3NextScannableBlock + + if scan_length > 0: + self.scanned_blocks = self._create_scanned_layers( + RemattedQwen3NextBlock, + length=scan_length, + metadata_axis_name="layers", + rngs=rngs, + **layer_kwargs, + ) + self.layers_remainder = RemattedQwen3NextBlock( + config=self.config, + mesh=mesh, + quant=self.quant, + model_mode=self.model_mode, + **rem_layer_kwargs, + rngs=rngs, + ) + def _init_scanned_generic(self, decoder_block_classes, rngs): """Initializes scanned generic decoder layers.""" config = self.config @@ -1854,6 +1894,13 @@ def __call__( layer_kwargs, kv_caches=kv_caches, ) + elif self.is_qwen3_next: + y = self._apply_qwen3_next_scanned_blocks( + y, + layer_args, + layer_kwargs, + kv_caches=kv_caches, + ) else: scan_length = int(cfg.num_decoder_layers / cfg.inhomogeneous_layer_cycle_interval) if kv_caches is not None: @@ -2208,6 +2255,71 @@ def pure_gemma_fn(graphdef, state_in, y_in, kv_in): return y + def _apply_qwen3_next_scanned_blocks( + self, + y, + layer_args, + layer_kwargs, + kv_caches=None, + ): + """Applies Qwen3-Next scanned decoder blocks, handling main scan and remainders.""" + + cfg = self.config + cycle_interval = cfg.inhomogeneous_layer_cycle_interval + scan_length = cfg.num_decoder_layers // cycle_interval + + block_unroll = max(1, scan_length) + if scan_length > 0: + grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan(kv_caches, scan_length, cycle_interval, stack=False) + y, self.scanned_blocks, _ = self._apply_layers_sequentially( + self.scanned_blocks, + y, + *layer_args, + length=scan_length, + kv_caches_stacked=grouped_kv_caches, + skip_block_remat=True, + unroll=block_unroll, + **layer_kwargs, + ) + maxtext_utils.update_kv_caches_after_scan(kv_caches, grouped_kv_caches, scan_length, cycle_interval, stacked=False) + + num_remaining_layers = cfg.num_decoder_layers % cycle_interval + if num_remaining_layers > 0: + policy = self.get_remat_policy() + prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg) + + remainder_kv = None + if kv_caches is not None: + start_idx = scan_length * cycle_interval + remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers]) + + def pure_qwen3_fn(graphdef, state_in, y_in, kv_in): + merged_layer = nnx.merge(graphdef, state_in) + call_kwargs = dict(layer_kwargs) + if kv_in is not None: + call_kwargs["kv_cache"] = kv_in + out_res = merged_layer(y_in, *layer_args, **call_kwargs) + if isinstance(out_res, tuple): + out_y = out_res[0] + out_kv = out_res[1] if len(out_res) > 1 else None + else: + out_y = out_res + out_kv = None + return out_y, out_kv, nnx.state(merged_layer) + + checkpointed_qwen3_fn = jax.checkpoint(pure_qwen3_fn, policy=policy, prevent_cse=prevent_cse) + + graphdef, state = nnx.split(self.layers_remainder) + y, updated_remainder_kv, new_state = checkpointed_qwen3_fn(graphdef, state, y, remainder_kv) + nnx.update(self.layers_remainder, new_state) + + if kv_caches is not None and updated_remainder_kv is not None: + start_idx = scan_length * cycle_interval + for offset, updated_item in enumerate(updated_remainder_kv): + kv_caches[start_idx + offset] = updated_item + + return y + def _apply_gemma4_small_layers( self, y, diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index dae15976c8..b1c63b388c 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -25,6 +25,7 @@ import jax.nn from jax import lax from jax.ad_checkpoint import checkpoint_name +from jax.experimental import xla_metadata from jax.sharding import Mesh import jax.numpy as jnp @@ -39,7 +40,7 @@ from maxtext.layers import moe from maxtext.layers import mhc from maxtext.common.common_types import HyperConnectionType -from maxtext.layers import nnx_wrappers +from maxtext.layers import nnx_scan, nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.embeddings import Qwen3OmniMoeVisionPosEmbedInterpolate, PositionalEmbedding from maxtext.layers.normalizations import RMSNorm, l2norm, Qwen3NextRMSNorm, Qwen3NextRMSNormGated @@ -48,7 +49,7 @@ from maxtext.layers.linears import DenseGeneral, MlpBlock from maxtext.layers.moe import RoutedMoE from maxtext.layers.initializers import nd_dense_init, variable_to_logically_partitioned -from maxtext.utils import max_utils +from maxtext.utils import max_utils, maxtext_utils from maxtext.inference import kvcache @@ -1184,7 +1185,7 @@ def __init__(self, config: Config, mesh: Mesh, quant: None | Quant = None, *, rn cfg = self.config # 1. Instantiate and apply the routed experts block. - self.routed_experts = moe.RoutedMoE( + self.routed_experts = RoutedMoE( config=cfg, num_experts=cfg.num_experts, num_experts_per_tok=cfg.num_experts_per_tok, @@ -1254,86 +1255,220 @@ def __call__(self, hidden_states: Array, deterministic: bool) -> tuple[Array, Ar class Qwen3NextScannableBlock(nnx.Module): - """A scannable block of Qwen3-Next decoder layers. + """A repeatable block of Qwen3-Next decoder layers, scanning local layers.""" - This module contains a fixed number of heterogeneous decoder layers that form - a repeating pattern, as defined by `config.inhomogeneous_layer_cycle_interval`. It is - intended to be the body of an `nn.scan` transformation to construct the full - decoder stack efficiently. - - Attributes: - config: The model configuration object. - mesh: The device mesh for sharding. - model_mode: The operational mode (e.g., 'train', 'prefill'). - quant: Optional quantization configuration. - """ + def __init__( + self, + config: Config, + mesh: Mesh, + model_mode: str, + rngs: nnx.Rngs, + quant: None | Quant = None, + num_of_layers: int | None = None, + remat_policy_fn: Any = None, + apply_internal_remat: bool = False, + ): + """Initializes the instance. - def __init__(self, config: Config, mesh: Mesh, model_mode: str, quant: None | Quant = None, *, rngs: nnx.Rngs): + Args: + config: The Config object with model hyperparameters. + mesh: The device mesh for distributed training. + model_mode: One of MODEL_MODE_TRAIN, MODEL_MODE_PREFILL, or MODEL_MODE_AUTOREGRESSIVE. + rngs: The random number generators for initialization. + quant: The quantization configuration. + num_of_layers: The number of layers in the block. + remat_policy_fn: The resolved rematerialization policy function. + apply_internal_remat: When True, the block rematerializes its own local + (scanned) and global layers, and the caller must NOT also apply + block-level remat. + """ self.config = config self.mesh = mesh self.model_mode = model_mode self.quant = quant self.rngs = rngs - cfg = self.config + cycle_interval = config.inhomogeneous_layer_cycle_interval + if num_of_layers is None: + num_of_layers = cycle_interval + self.num_of_layers = num_of_layers + self.remat_policy_fn = remat_policy_fn + self.apply_internal_remat = apply_internal_remat + + if not 0 <= num_of_layers <= cycle_interval: + raise ValueError( + f"Qwen3NextScannableBlock must contain between 0 and {cycle_interval} layers; got {num_of_layers}." + ) - # Instantiate each layer within the block in __init__ - for i in range(cfg.inhomogeneous_layer_cycle_interval): - layer_rngs = self.rngs.fork() # Fork RNGs for each layer - layer_name = f"layer_{i}" - layer = Qwen3NextDecoderLayer( + # Calculate local (GatedDeltaNet) vs global (FullAttention) layer counts for the block. + self.num_local = sum(1 for i in range(num_of_layers) if (i + 1) % cycle_interval != 0) + self.num_global = sum(1 for i in range(num_of_layers) if (i + 1) % cycle_interval == 0) + + if self.num_local > 0: + self.local_layers = nnx_scan.create_scanned_layers( + lambda layer_rngs: Qwen3NextDecoderLayer( + config=self.config, + mesh=self.mesh, + model_mode=self.model_mode, + quant=self.quant, + layer_idx=0, # layer_idx 0 is a GatedDeltaNet layer + rngs=layer_rngs, + ), + length=self.num_local, + param_scan_axis=self.config.param_scan_axis, + metadata_axis_name="local_layers", + rngs=self.rngs, + ) + else: + self.local_layers = None + + if self.num_global > 0: + self.global_layer = Qwen3NextDecoderLayer( config=self.config, mesh=self.mesh, - quant=self.quant, model_mode=self.model_mode, - layer_idx=i, - rngs=layer_rngs, + quant=self.quant, + layer_idx=cycle_interval - 1, # layer_idx cycle_interval-1 is a FullAttention layer + rngs=self.rngs, ) - setattr(self, layer_name, layer) + else: + self.global_layer = None + + def _run_layer(self, layer, y, layer_kwargs, kv_cache=None): + """Invokes one ``Qwen3NextDecoderLayer``, returning ``(output, updated_kv_cache)``.""" + out = layer(y, **layer_kwargs, kv_cache=kv_cache) + return out if isinstance(out, tuple) else (out, None) + + @property + def _remat_enabled(self): + """Whether the block rematerializes its own layers.""" + return self.apply_internal_remat and self.config.remat_policy != "none" + + def _scan_local_layers(self, y, layer_kwargs): + """Runs the local (linear attention / GatedDeltaNet) layers via a per-layer rematerialized ``jax.lax.scan``.""" + remat = self._remat_enabled + return nnx_scan.apply_scanned_layers( + self.local_layers, + y, + length=self.num_local, + param_scan_axis=self.config.param_scan_axis, + apply_fn=lambda layer, carry: self._run_layer(layer, carry, layer_kwargs)[0], + remat=remat, + remat_policy=self.remat_policy_fn if remat else None, + prevent_cse=maxtext_utils.should_prevent_cse_in_remat(self.config) if remat else True, + ) + + def _scan_global_layer(self, y, layer_kwargs): + """Runs the single global-attention layer inside a length-1 ``jax.lax.scan``.""" + cfg = self.config + graphdef_g, intermediate_g, other_g = nnx.split(self.global_layer, nnx.Intermediate, ...) + intermediate_xs = jax.tree.map(lambda x: x[None], intermediate_g) + + def run_global_layer(carry, intermediate_slice): + hidden_states, other = carry + layer = nnx.merge(graphdef_g, intermediate_slice, other) + new_hidden_states = self._run_layer(layer, hidden_states, layer_kwargs)[0] + _, new_intermediate, new_other = nnx.split(layer, nnx.Intermediate, ...) + return (new_hidden_states, new_other), new_intermediate + + global_remat_policy = self.remat_policy_fn + offload_names = maxtext_utils.get_save_and_offload_names(cfg) + if offload_names[0] or offload_names[1]: + save_names, offload_to_device = offload_names + global_remat_policy = jax.checkpoint_policies.save_only_these_names(*(save_names + offload_to_device)) + + if self._remat_enabled: + prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config) + run_global_layer = jax.checkpoint( + run_global_layer, + policy=global_remat_policy, + prevent_cse=prevent_cse, + ) + + with xla_metadata.set_xla_metadata(**{"skip-simplify-while-loops_trip-count-one": "true"}): + (y, final_other), stacked_intermediate = jax.lax.scan( + run_global_layer, + (y, other_g), + intermediate_xs, + length=1, + ) + + intermediate_state = jax.tree.map(lambda x: x[0], stacked_intermediate) + nnx.update(self.global_layer, final_other, intermediate_state) + return y + + def _forward_with_external_kv_cache(self, y, kv_cache, layer_kwargs): + """Runs the block with externally-supplied per-layer kv caches (vLLM PagedAttention / Mamba).""" + updated_kvs = [] + + if self.local_layers is not None: + graphdef, params, state = nnx.split(self.local_layers, nnx.Param, ...) + scan_axis = self.config.param_scan_axis + if scan_axis != 0: + params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), params) + per_layer_states = [] + for i in range(self.num_local): + current_params = jax.tree.map(lambda x, i=i: x[i], params) + current_state = jax.tree.map(lambda x, i=i: x[i], state) + layer = nnx.merge(graphdef, current_params, current_state) + current_kv = kv_cache[i] if (kv_cache is not None and i < len(kv_cache)) else None + y, new_kv = self._run_layer(layer, y, layer_kwargs, current_kv) + updated_kvs.append(new_kv) + per_layer_states.append(nnx.state(layer)) + + stacked_state = jax.tree.map(lambda *xs: jnp.stack(xs), *per_layer_states) + if scan_axis != 0: + stacked_params, stacked_other = stacked_state.split(nnx.Param, ...) + stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), stacked_params) + stacked_state = nnx.State.merge(stacked_params, stacked_other) + nnx.update(self.local_layers, stacked_state) + + if self.global_layer is not None: + global_kv = kv_cache[self.num_local] if (kv_cache is not None and self.num_local < len(kv_cache)) else None + y, new_kv = self._run_layer(self.global_layer, y, layer_kwargs, global_kv) + updated_kvs.append(new_kv) + + return y, tuple(updated_kvs) def __call__( self, - carry: jnp.ndarray, + inputs: jnp.ndarray, decoder_segment_ids: None | jnp.ndarray, decoder_positions: None | jnp.ndarray, deterministic: bool, model_mode: str, previous_chunk=None, slot: None | int = None, + page_state=None, + bidirectional_mask=None, kv_cache=None, attention_metadata=None, ) -> tuple[Array, None]: - """Applies the block of decoder layers to the input carry. + cfg = self.config + inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) + inputs = checkpoint_name(inputs, "decoder_layer_input") - Args: - carry: The input tensor from the previous scan iteration. - # ... other arguments are broadcasted to each iteration. + layer_kwargs = { + "decoder_segment_ids": decoder_segment_ids, + "decoder_positions": decoder_positions, + "deterministic": deterministic, + "model_mode": model_mode, + "slot": slot, + "previous_chunk": previous_chunk, + "attention_metadata": attention_metadata, + } - Returns: - A tuple containing the output of the block (the new carry) and an empty - value for the scan's `y` collection. - """ - cfg = self.config - x = carry - - # Loop over the number of sub-layers that make up one repeating pattern. - for i in range(cfg.inhomogeneous_layer_cycle_interval): - layer = getattr(self, f"layer_{i}") - # The second return value is kv_cache, which we ignore here because - # it is not passed as a carry in scannable layers. - x, _ = layer( - x, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk, - slot, - kv_cache=kv_cache, - attention_metadata=attention_metadata, - ) + if kv_cache is not None: + return self._forward_with_external_kv_cache(inputs, kv_cache, layer_kwargs) + + y = inputs + if self.local_layers is not None: + y = self._scan_local_layers(y, layer_kwargs) + if self.global_layer is not None: + y = self._scan_global_layer(y, layer_kwargs) - # The output of the block is the carry for the next scan iteration. - return x, None + if cfg.scan_layers: + return y, None + return y class Qwen3NextDecoderLayer(nnx.Module): diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index fcd5acb5cc..50f47aa8ff 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -51,7 +51,7 @@ from maxtext.layers.embeddings import Embed from maxtext.layers.nnx_decoders import NNXDecoder, NNXDecoderLayer, deepstack_process from maxtext.layers.normalizations import RMSNorm -from maxtext.models import gemma4, gemma4_small +from maxtext.models import gemma4, gemma4_small, qwen3 from maxtext.models.gpt3 import Gpt3LayerNorm from maxtext.models.llama2 import LlamaDecoderLayer from maxtext.utils import maxtext_utils @@ -716,8 +716,219 @@ def test_scan_layers(self): self.assertEqual(logits.shape, (batch, seq_len, cfg.vocab_size)) -if __name__ == "__main__": - unittest.main() +class _StatefulGemma4DecoderLayer(nnx.Module): + """Small stand-in that exposes cache ordering and mutable-state updates.""" + + def __init__(self, *, attention_type, **unused_kwargs): + self.increment = 10 if attention_type == AttentionType.GLOBAL else 1 + self.call_count = nnx.Intermediate(jnp.array(0, dtype=jnp.int32)) + self.received_attention_metadata = nnx.Intermediate(jnp.array(False)) + + def __call__( + self, + inputs, + *unused_args, + kv_cache=None, + attention_metadata=None, + **unused_kwargs, + ): + self.call_count.value += 1 + self.received_attention_metadata.value = attention_metadata is not None + output = inputs + self.increment + if kv_cache is None: + return output + return output, kv_cache + self.increment + + +class _SowingGemma4DecoderLayer(nnx.Module): + """Stand-in whose global layer sows an accumulating Intermediate, like MoE moe_lb_loss.""" + + def __init__(self, *, attention_type, **unused_kwargs): + self.is_global = attention_type == AttentionType.GLOBAL + # A trivial variable so the local layers have state for apply_scanned_layers to + # scan over (a bare module has nothing to scan and lax.scan can't infer length). + self.marker = nnx.Intermediate(jnp.zeros(())) + + def __call__(self, inputs, *unused_args, kv_cache=None, **unused_kwargs): + output = inputs + 1 + if self.is_global: + # nnx.sow appends into a tuple by default, so it grows across calls -- the + # MoE moe_lb_loss pattern that must not enter the global length-1 scan carry. + self.sow(nnx.Intermediate, "moe_lb_loss", jnp.sum(output)) + if kv_cache is None: + return output + return output, kv_cache + + +class TestGemma4ScannableBlock(unittest.TestCase): + """Tests Gemma4's nested local/global decoder block behavior.""" + + def setUp(self): + super().setUp() + self.config = SimpleNamespace( + dtype=jnp.float32, + param_scan_axis=1, + remat_policy="none", + scan_layers=True, + ) + + def _make_block(self): + return gemma4.Gemma4ScannableBlock( + config=self.config, + mesh=None, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + rngs=nnx.Rngs(0), + ) + + def test_updates_state_through_global_single_iteration_scan(self): + with mock.patch.object(gemma4, "Gemma4DecoderLayer", _StatefulGemma4DecoderLayer): + block = self._make_block() + output, updated_kvs = block( + jnp.zeros((1, 1, 1)), + decoder_segment_ids=None, + decoder_positions=None, + deterministic=True, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + ) + + np.testing.assert_array_equal(output, jnp.full((1, 1, 1), 15)) + self.assertIsNone(updated_kvs) + np.testing.assert_array_equal(block.local_layers.call_count.value, jnp.ones(5, dtype=jnp.int32)) + np.testing.assert_array_equal(block.global_layer.call_count.value, 1) + + def test_global_layer_sown_intermediate_accumulates_across_calls(self): + """A global layer that sows an accumulating Intermediate (e.g. MoE moe_lb_loss) + must not break the length-1 scan carry, even when the Intermediate already + exists from a previous call and the sow grows its tuple (1 -> 2 elements).""" + call_kwargs = { + "decoder_segment_ids": None, + "decoder_positions": None, + "deterministic": True, + "model_mode": MODEL_MODE_AUTOREGRESSIVE, + } + with mock.patch.object(gemma4, "Gemma4DecoderLayer", _SowingGemma4DecoderLayer): + block = self._make_block() + # First call creates moe_lb_loss on the global layer (1-tuple). + block(jnp.zeros((1, 1, 1)), **call_kwargs) + # Second call: moe_lb_loss already exists and the sow appends -> 2-tuple. + # Carrying it in the scan would change the carry pytree; the type-based + # split keeps Intermediates on the ys path instead. + block(jnp.zeros((1, 1, 1)), **call_kwargs) + + self.assertEqual(len(block.global_layer.moe_lb_loss.value), 2) + + def test_restores_local_state_and_preserves_kv_order(self): + attention_metadata = object() + + with mock.patch.object(gemma4, "Gemma4DecoderLayer", _StatefulGemma4DecoderLayer): + block = self._make_block() + output, updated_kvs = block( + jnp.zeros((1, 1, 1)), + decoder_segment_ids=None, + decoder_positions=None, + deterministic=True, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + kv_cache=tuple(jnp.array(i) for i in range(6)), + attention_metadata=attention_metadata, + ) + + np.testing.assert_array_equal(output, jnp.full((1, 1, 1), 15)) + np.testing.assert_array_equal(jnp.stack(updated_kvs), jnp.array([1, 2, 3, 4, 5, 15])) + np.testing.assert_array_equal(block.local_layers.call_count.value, jnp.ones(5, dtype=jnp.int32)) + np.testing.assert_array_equal( + block.local_layers.received_attention_metadata.value, + jnp.ones(5, dtype=jnp.bool_), + ) + np.testing.assert_array_equal(block.global_layer.call_count.value, 1) + np.testing.assert_array_equal(block.global_layer.received_attention_metadata.value, True) + + +class _StatefulQwen3NextDecoderLayer(nnx.Module): + """Small stand-in that exposes cache ordering and mutable-state updates for Qwen3-Next.""" + + def __init__(self, *, layer_idx, **unused_kwargs): + is_global = (layer_idx + 1) % 4 == 0 + self.increment = 10 if is_global else 1 + self.call_count = nnx.Intermediate(jnp.array(0, dtype=jnp.int32)) + self.received_attention_metadata = nnx.Intermediate(jnp.array(False)) + + def __call__( + self, + inputs, + *unused_args, + kv_cache=None, + attention_metadata=None, + **unused_kwargs, + ): + self.call_count.value += 1 + self.received_attention_metadata.value = attention_metadata is not None + output = inputs + self.increment + if kv_cache is None: + return output + return output, kv_cache + self.increment + + +class TestQwen3NextScannableBlock(unittest.TestCase): + """Tests Qwen3-Next's nested local/global decoder block behavior.""" + + def setUp(self): + super().setUp() + self.config = SimpleNamespace( + dtype=jnp.float32, + param_scan_axis=1, + remat_policy="none", + scan_layers=True, + inhomogeneous_layer_cycle_interval=4, + ) + + def _make_block(self): + return qwen3.Qwen3NextScannableBlock( + config=self.config, + mesh=None, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + rngs=nnx.Rngs(0), + ) + + def test_updates_state_through_global_single_iteration_scan(self): + with mock.patch.object(qwen3, "Qwen3NextDecoderLayer", _StatefulQwen3NextDecoderLayer): + block = self._make_block() + output, updated_kvs = block( + jnp.zeros((1, 1, 1)), + decoder_segment_ids=None, + decoder_positions=None, + deterministic=True, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + ) + + np.testing.assert_array_equal(output, jnp.full((1, 1, 1), 13)) + self.assertIsNone(updated_kvs) + np.testing.assert_array_equal(block.local_layers.call_count.value, jnp.ones(3, dtype=jnp.int32)) + np.testing.assert_array_equal(block.global_layer.call_count.value, 1) + + def test_restores_local_state_and_preserves_kv_order(self): + attention_metadata = object() + + with mock.patch.object(qwen3, "Qwen3NextDecoderLayer", _StatefulQwen3NextDecoderLayer): + block = self._make_block() + output, updated_kvs = block( + jnp.zeros((1, 1, 1)), + decoder_segment_ids=None, + decoder_positions=None, + deterministic=True, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + kv_cache=tuple(jnp.array(i) for i in range(4)), + attention_metadata=attention_metadata, + ) + + np.testing.assert_array_equal(output, jnp.full((1, 1, 1), 13)) + np.testing.assert_array_equal(jnp.stack(updated_kvs), jnp.array([1, 2, 3, 13])) + np.testing.assert_array_equal(block.local_layers.call_count.value, jnp.ones(3, dtype=jnp.int32)) + np.testing.assert_array_equal( + block.local_layers.received_attention_metadata.value, + jnp.ones(3, dtype=jnp.bool_), + ) + np.testing.assert_array_equal(block.global_layer.call_count.value, 1) + np.testing.assert_array_equal(block.global_layer.received_attention_metadata.value, True) class _StatefulGemma4DecoderLayer(nnx.Module): @@ -1259,3 +1470,7 @@ def mock_donor_idx(lyr, layer_types, num_kv_shared): model_mode=MODEL_MODE_TRAIN, kv_caches=kv_caches, ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/param_mapping_test.py b/tests/unit/param_mapping_test.py index 9a01a57aef..15f5bfafaf 100644 --- a/tests/unit/param_mapping_test.py +++ b/tests/unit/param_mapping_test.py @@ -105,7 +105,14 @@ def test_qwen3_next_mapping_scanned(self): maxtext_config = mock.Mock() maxtext_config.inhomogeneous_layer_cycle_interval = 2 mapping = param_mapping.QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True) - self.assertIn("params-decoder-layers-layer_0-input_layernorm-scale", mapping) + self.assertIn("params-decoder-scanned_blocks-local_layers-input_layernorm-scale", mapping) + self.assertIn("params-decoder-scanned_blocks-global_layer-input_layernorm-scale", mapping) + num_blocks = config["num_hidden_layers"] // maxtext_config.inhomogeneous_layer_cycle_interval + local_val = mapping["params-decoder-scanned_blocks-local_layers-input_layernorm-scale"] + global_val = mapping["params-decoder-scanned_blocks-global_layer-input_layernorm-scale"] + self.assertEqual(len(local_val), num_blocks) + self.assertEqual(len(local_val[0]), 1) + self.assertEqual(len(global_val), num_blocks) def test_deepseek_mapping(self): config = { From c5ef83219bb70f5eb70cfb440f86b203365ce342 Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Fri, 7 Aug 2026 17:58:00 +0000 Subject: [PATCH 07/12] add muon impl --- src/maxtext/optimizers/muon/__init__.py | 28 + src/maxtext/optimizers/muon/muon.py | 1479 +++++++++++++++++++++++ 2 files changed, 1507 insertions(+) create mode 100644 src/maxtext/optimizers/muon/__init__.py create mode 100644 src/maxtext/optimizers/muon/muon.py diff --git a/src/maxtext/optimizers/muon/__init__.py b/src/maxtext/optimizers/muon/__init__.py new file mode 100644 index 0000000000..c03c0b126a --- /dev/null +++ b/src/maxtext/optimizers/muon/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2025 DeepMind Technologies Limited. All Rights Reserved. +# +# 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 +# +# http://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. +# ============================================================================== +"""Local fork of optax 0.2.8 contrib/_muon.py for MaxText. + +This fork pins the Muon optimizer implementation to optax 0.2.8 so that +MaxText can add custom modifications without being blocked by upstream +release cycles. +""" + +from third_party.optax_muon._muon import MuonDimensionNumbers +from third_party.optax_muon._muon import MuonState +from third_party.optax_muon._muon import WeightDimNumOrFn +from third_party.optax_muon._muon import muon +from third_party.optax_muon._muon import orthogonalize_via_newton_schulz +from third_party.optax_muon._muon import scale_by_muon +from third_party.optax_muon._muon import scale_by_shape diff --git a/src/maxtext/optimizers/muon/muon.py b/src/maxtext/optimizers/muon/muon.py new file mode 100644 index 0000000000..62de8597b6 --- /dev/null +++ b/src/maxtext/optimizers/muon/muon.py @@ -0,0 +1,1479 @@ +# Copyright 2025 DeepMind Technologies Limited. All Rights Reserved. +# +# 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 +# +# http://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. +# ============================================================================== +"""Muon. + +Implementation of the +[Muon optimizer](https://github.com/KellerJordan/modded-nanogpt) +by Keller Jordan +""" + +# pylint: disable=unnecessary-lambda-assignment + + +import functools +import itertools +import logging +import math +from typing import Any, Callable, NamedTuple, Optional, Union, Sequence, Literal + +import jax +import jax.numpy as jnp + +from optax._src import alias +from optax._src import base +from optax._src import combine +from optax._src import numerics +from optax._src import transform +from optax._src import utils +from optax.transforms import _masking +import optax.tree + +from jax.sharding import NamedSharding, PartitionSpec + +ReshapeFn = Callable[[jax.Array], jax.Array] + +_PRECONDITIONINGS = ["frobenius", "spectral", "aol", "schatten"] +_DEFAULT_NS_COEFFS = (3.4445, -4.7750, 2.0315) +_DION_NS_COEFFS = [ + (4.0848, -6.8946, 2.9270), + (3.9505, -6.3029, 2.6377), + (3.7418, -5.5913, 2.3037), + (2.8769, -3.1427, 1.2046), + (2.8366, -3.0525, 1.2012), +] +_NS_COEFFS_PRESET_DICT = { + "standard": _DEFAULT_NS_COEFFS, + "dion": _DION_NS_COEFFS, +} + + +class MuonDimensionNumbers(NamedTuple): + """Specification for which weight axes participate in matrix projection. + + Muon defines an orthogonalization for 2D matrix weights for matrix-vector + products: + + .. math:: + x W = y + + where the first matrix dimension is the reduction axis and the second matrix + dimension is the output axis. Thus, the default spec consists of 0 and 1 + reduction and output axes respectively. + + .. warning:: + The batch axes are implicit, all axes not specified as reduction or output + axes are considered batch axes and will be considered independently in the + orthogonalization (via jax.vmap). + + When ``component_splits`` is set, the output axis is split into components + (e.g., MLA nope/rope) and the batch axes are merged into each component's + output dimension before NS, matching Megatron's per-component treatment. + Each component gets its own NS orthogonalization and scale factor. + + When ``scan_axis`` is set (requires ``component_splits``), the specified axis + is excluded from batch-merge and vmapped over instead, so each scan slice + gets independent NS — producing identical results to unscan mode. + """ + + reduction_axis: Sequence[int] | int = 0 + output_axis: Sequence[int] | int = 1 + component_splits: tuple[int, ...] | None = None + scan_axis: int | None = None + + +WeightDimNumOrFn = MuonDimensionNumbers | base.Params | Callable[[base.Params], base.Params | None] + + +_is_weight_dim_nums = lambda x: isinstance(x, MuonDimensionNumbers) + + +def _normalize_axes(x: jax.Array, dim_nums: MuonDimensionNumbers) -> tuple[tuple[int, ...], tuple[int, ...]]: + """Normalize axes in dimension numbers to two tuples of non-negative ints.""" + if isinstance(dim_nums.reduction_axis, int): + dim_nums = dim_nums._replace(reduction_axis=(dim_nums.reduction_axis,)) + reduction_axes = tuple(ax % x.ndim for ax in dim_nums.reduction_axis) + + if isinstance(dim_nums.output_axis, int): + dim_nums = dim_nums._replace(output_axis=(dim_nums.output_axis,)) + output_axes = tuple(ax % x.ndim for ax in dim_nums.output_axis) + return reduction_axes, output_axes + + +def _compute_muon_reshape(x: jax.Array, dim_nums: MuonDimensionNumbers) -> tuple[ReshapeFn, ReshapeFn]: + """Compute the reshape and inverse functions for an array from a spec.""" + if x.ndim < 2: + raise ValueError("Muon optimized parameters must have rank >= 2, got" f" {x.ndim=}") + reduction_axes, output_axes = _normalize_axes(x, dim_nums) + if set(reduction_axes) & set(output_axes): + raise ValueError( + "Normalized reduction axes and output axes must be" + f" disjoint, got {reduction_axes} and {output_axes}." + f" Originally {dim_nums=} and {x.shape=}" + ) + batch_axes = tuple(sorted(set(range(x.ndim)) - set(reduction_axes) - set(output_axes))) + transpose = batch_axes + reduction_axes + output_axes + inv_transpose = tuple(sorted(range(x.ndim), key=lambda i: transpose[i])) + axes2shape = lambda axes: tuple(x.shape[ax] for ax in axes) + # Reshape to (batch, reduction, output) to match the (reduction, output) + # structure of the original muon for 2D weights. + flat_shape = ( + math.prod(axes2shape(batch_axes)), + math.prod(axes2shape(reduction_axes)), + math.prod(axes2shape(output_axes)), + ) + unflat_shape = axes2shape(batch_axes) + axes2shape(reduction_axes) + axes2shape(output_axes) + reshape_fn = lambda x: x.transpose(transpose).reshape(flat_shape) + inverse_fn = lambda x: x.reshape(unflat_shape).transpose(inv_transpose) + return reshape_fn, inverse_fn + + +def _get_shape_products(x: jax.Array, dim_nums: MuonDimensionNumbers) -> tuple[float, float]: + reduction_axes, output_axes = _normalize_axes(x, dim_nums) + fan_in = math.prod(x.shape[ax] for ax in reduction_axes) + fan_out = math.prod(x.shape[ax] for ax in output_axes) + return fan_in, fan_out + + +def _component_splits_for_output_axis( + x: jax.Array, + dim_nums: MuonDimensionNumbers, + reduction_axes: tuple[int, ...], + output_axes: tuple[int, ...], +) -> tuple[int, ...]: + """Return per-output-axis component splits, accepting either per-axis or total-batched splits.""" + if len(output_axes) != 1: + raise ValueError(f"component_splits requires exactly 1 output axis, got {output_axes}") + + output_axis = output_axes[0] + splits = dim_nums.component_splits + if sum(splits) == x.shape[output_axis]: + return splits + + batch_axes = tuple(sorted(set(range(x.ndim)) - set(reduction_axes) - set(output_axes))) + batch_size = math.prod(x.shape[ax] for ax in batch_axes) + if batch_size > 1 and sum(splits) == batch_size * x.shape[output_axis] and all(s % batch_size == 0 for s in splits): + return tuple(s // batch_size for s in splits) + + raise ValueError( + f"component_splits {splits} must sum to output axis size {x.shape[output_axis]} " + f"or total batched output size {batch_size * x.shape[output_axis]} (and be divisible by batch_size {batch_size})" + ) + + +def _build_component_scale_tensor( + update: jax.Array, + dim_nums: MuonDimensionNumbers, + scale_fn: Callable[[float, float], float], +) -> jax.Array: + """Build a broadcast-compatible scale tensor for component_splits. + + Each component region along the output axis gets its own scale factor, + computed from the per-component fan values (where batch dims are merged + into the output, matching Megatron's per-component treatment). + + When ``scan_axis`` is set, it is excluded from the batch product so that + scale factors match the unscan (per-layer) computation. + """ + reduction_axes, output_axes = _normalize_axes(update, dim_nums) + batch_axes = tuple(sorted(set(range(update.ndim)) - set(reduction_axes) - set(output_axes))) + output_axis = output_axes[0] + + # Exclude scan_axis from batch axes for split normalization and fan computation. + scan_axis = dim_nums.scan_axis + if scan_axis is not None: + scan_axis_norm = scan_axis % update.ndim + batch_axes_for_splits = tuple(ax for ax in batch_axes if ax != scan_axis_norm) + else: + scan_axis_norm = None + batch_axes_for_splits = batch_axes + + if scan_axis is not None: + # Create a reduced-rank view (without scan axis) so + # _component_splits_for_output_axis sees the correct batch_size. + remaining_axes = [i for i in range(update.ndim) if i != scan_axis_norm] + slice_shape = tuple(update.shape[ax] for ax in remaining_axes) + + def _remap(ax): + return ax - 1 if ax > scan_axis_norm else ax + + inner_reduction = tuple(_remap(ax) for ax in reduction_axes) + inner_output = tuple(_remap(ax) for ax in output_axes) + inner_dim_nums = MuonDimensionNumbers( + reduction_axis=inner_reduction, + output_axis=inner_output, + component_splits=dim_nums.component_splits, + ) + dummy = jnp.zeros(slice_shape) + splits = _component_splits_for_output_axis(dummy, inner_dim_nums, inner_reduction, inner_output) + else: + splits = _component_splits_for_output_axis(update, dim_nums, reduction_axes, output_axes) + + fan_in = math.prod(update.shape[ax] for ax in reduction_axes) + batch_prod = math.prod(update.shape[ax] for ax in batch_axes_for_splits) + + # Build a 1D array of scales along the output axis. + scales = [] + for s in splits: + fan_out = batch_prod * s + scale = scale_fn(fan_in, fan_out) + scales.extend([scale] * s) + + # Reshape to broadcast: all dims are 1 except the output axis. + shape = [1] * update.ndim + shape[output_axis] = sum(splits) + return jnp.asarray(scales, dtype=update.dtype).reshape(shape) + + +def _scale_update_for_width_transfer(update: jax.Array, dim_nums: MuonDimensionNumbers): + """Apply width scaling from .""" + if getattr(dim_nums, "component_splits", None) is not None: + return update * _build_component_scale_tensor( + update, dim_nums, lambda fan_in, fan_out: math.sqrt(max(1, fan_out / fan_in)) + ) + fan_in, fan_out = _get_shape_products(update, dim_nums) + scale = jnp.sqrt(jnp.maximum(1, fan_out / fan_in)) + # Cast scale to the update's dtype so shape-scale stays in whatever dtype + # NS produced (fp32 by default, bf16 if the caller requested Megatron-style + # parity). This keeps the downstream chain (add_decayed_weights, + # scale_by_learning_rate) in that same dtype until the fp32 param add. + return update * jnp.asarray(scale, dtype=update.dtype) + + +def _scale_update_for_consistent_rms( + update: jax.Array, dim_nums: MuonDimensionNumbers, consistent_rms: jax.typing.ArrayLike +): + """Apply consistent RMS scaling from .""" + if getattr(dim_nums, "component_splits", None) is not None: + return update * _build_component_scale_tensor( + update, dim_nums, lambda fan_in, fan_out: (math.sqrt(max(fan_in, fan_out)) * float(consistent_rms)) + ) + fan_in, fan_out = _get_shape_products(update, dim_nums) + scale = jnp.sqrt(jnp.maximum(fan_in, fan_out)) * consistent_rms + # Keep update in its original dtype (fp32 by default) so the downstream + # LR/WD multiplies stay in the NS output's precision. + return update * jnp.asarray(scale, dtype=update.dtype) + + +def scale_by_shape( + weight_dimension_numbers: WeightDimNumOrFn | None = None, + consistent_rms: jax.typing.ArrayLike | None = None, +) -> base.GradientTransformation: + """Scale updates by factors derived from parameter shape. + + Args: + weight_dimension_numbers: An optional tree with the same structure as the + params of `MuonDimensionNumbers`s, specifying how to reshape the + parameters before and after the orthogonalization OR a callable returning + such a tree. None implies that all parameters are 2D matrices. + consistent_rms: An optional float to activate consistent RMS scaling. + If float, scales updates by `sqrt(max(fan_in, fan_out)) * consistent_rms`. + If None, uses width scaling `sqrt(max(1, fan_out / fan_in))`. + + Returns: + A `GradientTransformation` object. + """ + + def update_fn(updates, state, params=None): + del params + if callable(weight_dimension_numbers): + # Populate weight_dim_nums if it's a callable. Use updates instead of + # actual params since only shapes matter and params may not be provided. + resolved_weight_dim_nums = weight_dimension_numbers(updates) + else: + resolved_weight_dim_nums = weight_dimension_numbers + + if consistent_rms is not None: + scaling_fn = functools.partial(_scale_update_for_consistent_rms, consistent_rms=consistent_rms) + else: + scaling_fn = _scale_update_for_width_transfer + + scaled_updates = jax.tree.map( + scaling_fn, + updates, + resolved_weight_dim_nums, + is_leaf=_is_weight_dim_nums, + ) + return scaled_updates, state + + # Use the standard empty_state initializer, as this transform is stateless + return base.GradientTransformation(base.init_empty_state, update_fn) + + +def _aol_first_newton_schulz_iteration( + x: jax.Array, + coeffs: jax.Array, + eps: jax.typing.ArrayLike = 1e-8, +) -> jax.Array: + """'Almost Orthogonal Layer' Preconditioning with Newton-Schulz iteration.""" + # Implements the first Newton-Schulz step with AOL preconditioning + # which allows for better orthogonalization performance. + a = x @ x.T.conj() + rescaling = jnp.clip(jnp.abs(a).sum(axis=-1), min=eps) + s = jnp.expand_dims(jax.lax.rsqrt(rescaling), -1) + x, a = x * s, a * s * s.transpose(-1, -2) + b = coeffs[1] * a + coeffs[2] * a @ a + return coeffs[0] * x + b @ x + + +def _schatten_first_newton_schulz_iteration( + x: jax.Array, + coeffs: jax.Array, + eps: jax.typing.ArrayLike = 1e-8, +) -> jax.Array: + """Schatten-4 Preconditioning with Newton-Schulz iteration.""" + # Implements the first Newton-Schulz step with Schatten-4 norm + # preconditioning which allows for better orthogonalization performance. + a = x @ x.T + rescaling = jnp.clip(jnp.linalg.norm(a, ord="fro", axis=(-2, -1)), min=eps) + s = jnp.expand_dims(jax.lax.rsqrt(rescaling), (0, -1)) + x, a = x * s, a * s**2 + b = coeffs[1] * a + coeffs[2] * a @ a + return coeffs[0] * x + b @ x + + +def _base_newton_schulz_iteration(x: jax.Array, coeffs: jax.Array) -> jax.Array: + # Implements Newton-Schulz step f(X) = c_0 X + c_1 (XX^T)X + c_2 (XX^T)^2X, + # with quintic form f(X) = c_0 X + (c_1 A + c_2 AA)X, where A = XX^T. + # The NS step has the property f(X) = f(X^T)^T. That is, we can get equivalent + # result by transposing input and output. In particular, we may transpose X + # when rows > cols for efficiency. + a = x @ x.T.conj() + b = coeffs[1] * a + coeffs[2] * a @ a + return coeffs[0] * x + b @ x + + +_newton_schulz_iterator = _base_newton_schulz_iteration # backwards compat + + +def _aol_ns_iterator(i, x, coeffs): + # Modified first step using AOL rescaling + return jax.lax.cond( + i == 0, + lambda x: _aol_first_newton_schulz_iteration(x, coeffs), + lambda x: _base_newton_schulz_iteration(x, coeffs), + x, + ) + + +def _schatten_ns_iterator(i, x, coeffs): + # Modified first step using Schatten-4 norm rescaling + return jax.lax.cond( + i == 0, + lambda x: _schatten_first_newton_schulz_iteration(x, coeffs), + lambda x: _base_newton_schulz_iteration(x, coeffs), + x, + ) + + +def _base_ns_iterator(i, x, coeffs): + del i + return _base_newton_schulz_iteration(x, coeffs) + + +def _orthogonalize_components( + x: jax.Array, + ns_coeffs: jax.Array, + ns_steps: jax.typing.ArrayLike, + preconditioning: str, + eps: jax.typing.ArrayLike, + dimension_numbers: MuonDimensionNumbers, + ns_unroll: bool = False, + replicate_ns_matrix_axes: bool = False, + replicate_ns_batch_axis: str | None = None, +) -> jax.Array: + """Orthogonalize with per-component splitting on the output axis. + + Matches Megatron's per-component treatment of MLA projections: splits the + output axis into semantic components (e.g., nope/rope), merges batch axes + (heads) into each component, and applies NS independently to each resulting + 2D matrix. Each component sees ``(reduction, batch * comp_size)`` where + ``comp_size`` is the component's slice of the output axis. + + When ``scan_axis`` is set, that axis is vmapped over (not merged into batch), + so each scan slice gets independent NS — matching unscan mode exactly. + + Args: + x: The weight/gradient tensor with shape including reduction, batch (head), + and output axes. + ns_coeffs: Newton-Schulz coefficients. + ns_steps: Number of NS iterations. + preconditioning: Preconditioning method. + eps: Numerical stability epsilon. + dimension_numbers: MuonDimensionNumbers with component_splits set. + + Returns: + Orthogonalized tensor with the same shape as the input. + """ + saved_sharding = _get_array_sharding(x) if replicate_ns_matrix_axes else None + reduction_axes, output_axes = _normalize_axes(x, dimension_numbers) + + # Validate scan_axis if present. + scan_axis = dimension_numbers.scan_axis + if scan_axis is not None: + scan_axis = scan_axis % x.ndim + if scan_axis in reduction_axes: + raise ValueError( + f"scan_axis={dimension_numbers.scan_axis} (normalized: {scan_axis}) " + f"must not overlap with reduction_axis={reduction_axes}" + ) + if scan_axis in output_axes: + raise ValueError( + f"scan_axis={dimension_numbers.scan_axis} (normalized: {scan_axis}) " + f"must not overlap with output_axis={output_axes}" + ) + + if scan_axis is not None: + # Move scan_axis to position 0, vmap the inner function over it, + # then move it back. The inner function operates on tensors without + # the scan dimension (e.g. 3D [R, H, O] instead of 4D [R, L, H, O]). + # Remap dimension_numbers axes to account for the removed scan_axis. + def _remap_axis(ax): + return ax - 1 if ax > scan_axis else ax + + inner_reduction = tuple(_remap_axis(ax) for ax in reduction_axes) + inner_output = tuple(_remap_axis(ax) for ax in output_axes) + inner_dim_nums = MuonDimensionNumbers( + reduction_axis=inner_reduction, + output_axis=inner_output, + component_splits=dimension_numbers.component_splits, + scan_axis=None, + ) + + # Move scan_axis to position 0. + perm = [scan_axis] + [i for i in range(x.ndim) if i != scan_axis] + x_moved = jnp.transpose(x, perm) + + def _inner(x_slice): + return _orthogonalize_components( + x_slice, + ns_coeffs, + ns_steps, + preconditioning, + eps, + inner_dim_nums, + ns_unroll=ns_unroll, + replicate_ns_matrix_axes=replicate_ns_matrix_axes, + replicate_ns_batch_axis=replicate_ns_batch_axis, + ) + + result_moved = jax.vmap(_inner)(x_moved) + + # Inverse transpose to restore original axis order. + inv_perm = [0] * len(perm) + for i, p in enumerate(perm): + inv_perm[p] = i + return _maybe_restore_sharding(jnp.transpose(result_moved, inv_perm), saved_sharding) + + # --- Non-scan path (existing logic) --- + batch_axes = tuple(sorted(set(range(x.ndim)) - set(reduction_axes) - set(output_axes))) + output_axis = output_axes[0] + splits = _component_splits_for_output_axis(x, dimension_numbers, reduction_axes, output_axes) + + # Transpose to (reduction..., batch..., output) canonical order. + perm = tuple(reduction_axes) + tuple(batch_axes) + (output_axis,) + x_t = jnp.transpose(x, perm) + + n_reduction = len(reduction_axes) + n_batch = len(batch_axes) + reduction_size = math.prod(x_t.shape[:n_reduction]) + batch_size = math.prod(x_t.shape[n_reduction : n_reduction + n_batch]) + + # Split along the last axis (output) into components. + split_indices = list(itertools.accumulate(splits[:-1])) + components = jnp.split(x_t, split_indices, axis=-1) + + # Process each component: merge batch into output -> 2D -> NS -> reshape back. + results = [] + for comp in components: + comp_size = comp.shape[-1] + # Reshape to 2D: (reduction_prod, batch_prod * comp_size). + comp_2d = comp.reshape(reduction_size, batch_size * comp_size) + if replicate_ns_matrix_axes: + comp_2d = _replicate_matrix_axes_for_ns(comp_2d, _get_current_mesh_for_ns()) + # Apply NS as a standard 2D matrix (no component_splits in the recursive + # call since MuonDimensionNumbers(0, 1) defaults to None). + comp_ortho = orthogonalize_via_newton_schulz(comp_2d, ns_coeffs, ns_steps, preconditioning, eps, ns_unroll=ns_unroll) + # Reshape back to (reduction..., batch..., comp_size). + comp_back = comp_ortho.reshape(x_t.shape[: n_reduction + n_batch] + (comp_size,)) + results.append(comp_back) + + # Concatenate along output axis and inverse transpose. + result_t = jnp.concatenate(results, axis=-1) + inv_perm = tuple(sorted(range(len(perm)), key=lambda i: perm[i])) + return _maybe_restore_sharding(jnp.transpose(result_t, inv_perm), saved_sharding) + + +def orthogonalize_via_newton_schulz( + x: jax.Array, + ns_coeffs: jax.Array, + ns_steps: jax.typing.ArrayLike = 5, + preconditioning: Literal["frobenius", "spectral", "aol", "schatten"] = "frobenius", + eps: jax.typing.ArrayLike = 1e-8, + dimension_numbers: MuonDimensionNumbers | None = None, + ns_unroll: bool = False, + replicate_ns_matrix_axes: bool = False, + replicate_ns_batch_axis: str | None = None, +) -> jax.Array: + r"""Orthogonalize via Newton-Schulz iteration. + + We opt to use a quintic iteration whose coefficients are selected to maximize + the slope at zero. For the purpose of minimizing steps, it turns out to be + empirically effective to keep increasing the slope at zero even beyond the + point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather + something like US'V^T where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), + which turns out not to hurt model performance at all relative to UV^T, where + USV^T = G is the SVD. + + Args: + x: A matrix to orthogonalize. + ns_coeffs: Coefficients for the Newton-schulz iterators. + Must have shape (n, 3) where n is the number of iterations. + ns_steps: Number of Newton-schulz iterations. + Ignored if `ns_coeffs` is a 2D array. + preconditioning: Which preconditioning method to use. + eps: Term added to denominators to improve numerical stability. + dimension_numbers: Optional spec for reshaping a tensor before and after the + orthogonalization, to support non-2D parameters. + + Returns: + The orthogonalized matrix. + """ + # Dispatch to component-split path before any other logic. + if isinstance(dimension_numbers, MuonDimensionNumbers): + if dimension_numbers.scan_axis is not None and dimension_numbers.component_splits is None: + raise ValueError("scan_axis requires component_splits to be set") + if dimension_numbers.component_splits is not None: + return _orthogonalize_components( + x, + ns_coeffs, + ns_steps, + preconditioning, + eps, + dimension_numbers, + ns_unroll, + replicate_ns_matrix_axes, + replicate_ns_batch_axis, + ) + + if x.ndim != 2 and not isinstance(dimension_numbers, MuonDimensionNumbers): + raise ValueError( + f"Input must have shape (m, n) or weight dimension numbers must be" + f" provided. Got shape={x.shape} and {dimension_numbers=}." + ) + if x.ndim == 2: + dimension_numbers = MuonDimensionNumbers(reduction_axis=0, output_axis=1) + if ns_coeffs.ndim > 2 or ns_coeffs.shape[-1] != 3: + raise ValueError("Newton-Schulz coefficients must have shape (3,) or" f" (n, 3), got {ns_coeffs.shape}") + + # Cache mesh once so the inner closure can apply rank-2 P(None, None) + # constraints directly on the dot operands inside vmap. + ns_mesh = _get_current_mesh_for_ns() if replicate_ns_matrix_axes else None + + def _replicate_rank2_for_ns(x): + """Apply P(None, None) constraint to a rank-2 matrix inside vmap.""" + if ns_mesh is None: + return x + return jax.lax.with_sharding_constraint(x, NamedSharding(ns_mesh, PartitionSpec(None, None))) + + def _orthogonalize(x): + # Constrain rank-2 slice directly so the dot x @ x.T inside the NS + # iterator does not trigger an all-reduce on the contraction dim. + x = _replicate_rank2_for_ns(x) + + transposed = False + if x.shape[0] > x.shape[1]: + x = x.T + transposed = True + # Re-constrain after transpose; XLA may pick a sharded layout for + # the transposed tensor otherwise. + x = _replicate_rank2_for_ns(x) + + ns_iterators = { + "frobenius": _base_ns_iterator, + "spectral": _base_ns_iterator, + "aol": _aol_ns_iterator, + "schatten": _schatten_ns_iterator, + } + if preconditioning not in _PRECONDITIONINGS: + raise ValueError(f"Unknown preconditioning {preconditioning}") + _ns_iterator = ns_iterators[preconditioning] + + if preconditioning == "frobenius": + x /= jnp.linalg.norm(x, ord="fro") + eps + elif preconditioning == "spectral": + x /= jnp.linalg.norm(x, ord=2) + eps + else: + pass + + ns_coeffs_ = ns_coeffs.astype(x.dtype) + + if ns_coeffs_.ndim == 1: + x = jax.lax.fori_loop(0, ns_steps, lambda i, x: _ns_iterator(i, x, ns_coeffs_), x, unroll=ns_unroll) + else: + + def _scan_body(carry, coeffs_step): + i, x = carry + x_new = _ns_iterator(i, x, coeffs_step) + return (i + 1, x_new), None + + init_carry = (jnp.asarray(0, dtype=jnp.int32), x) + (_, x), _ = jax.lax.scan(_scan_body, init_carry, ns_coeffs_) + + if transposed: + x = x.T + return x + + saved_sharding = _get_array_sharding(x) if replicate_ns_matrix_axes else None + reshape_fn, inverse_fn = _compute_muon_reshape(x, dimension_numbers) + reshaped = reshape_fn(x) + if replicate_ns_matrix_axes: + reshaped = _replicate_matrix_axes_for_ns(reshaped, _get_current_mesh_for_ns(), replicate_ns_batch_axis) + result = inverse_fn(jax.vmap(_orthogonalize)(reshaped)) + return _maybe_restore_sharding(result, saved_sharding) + + +class MuonState(NamedTuple): + """State for the Muon algorithm.""" + + count: jax.typing.ArrayLike # shape=(), dtype=jnp.int32. + mu: base.Updates + ns_coeffs: jax.typing.ArrayLike # shape=(), dtype=jnp.int32. + + +def _get_array_sharding(x): + """Get sharding from a JAX array, or None if unavailable.""" + return getattr(x, "sharding", None) + + +def _get_current_mesh_for_ns(): + """Return the active mesh context used to build NS sharding constraints.""" + get_abstract_mesh = getattr(jax.sharding, "get_abstract_mesh", None) + if get_abstract_mesh is not None: + mesh = get_abstract_mesh() + if mesh is not None and not mesh.empty: + return mesh + raise ValueError( + "replicate_ns_matrix_axes=True requires an active mesh context. " + "Run the optimizer update under jax.set_mesh(...)." + ) + + +def _replicate_matrix_axes_for_ns(x: jax.Array, mesh, batch_axis=None) -> jax.Array: + """Replicate matrix (m, n) axes while keeping batch axes sharded. + + Args: + x: NS input tensor. Supports rank-2 (m, n), rank-3 (batch, m, n), + and rank-4 chunked input (num_chunks, batch_update_size, m, n). + mesh: JAX Mesh object for constructing NamedSharding. + batch_axis: Mesh axis name for the leading batch-like dimension. Must be + set for rank-3/rank-4 inputs, e.g. "fsdp". + """ + if x.ndim < 2: + return x + + if x.ndim == 2: + spec = PartitionSpec(None, None) + elif x.ndim == 3: + if batch_axis is None: + raise ValueError("replicate_ns_batch_axis must be set when replicating rank-3 Muon NS inputs") + chosen_axis = _maybe_choose_ns_batch_axis(x.shape[0], mesh, batch_axis) + if chosen_axis is not None: + spec = PartitionSpec(chosen_axis, None, None) + elif _fits_full_replicate_budget(x): + spec = PartitionSpec(None, None, None) + else: + return x + elif x.ndim == 4: + if batch_axis is None: + raise ValueError("replicate_ns_batch_axis must be set when replicating rank-4 Muon NS inputs") + spec = PartitionSpec(batch_axis, None, None, None) + else: + return x + + return jax.lax.with_sharding_constraint(x, NamedSharding(mesh, spec)) + + +def _mesh_axis_size(mesh, axis_name: str | tuple[str, ...] | None) -> int: + """Return the total size of one mesh axis or a tuple of mesh axes.""" + if axis_name is None: + return 1 + if isinstance(axis_name, tuple): + return math.prod(int(mesh.shape[axis]) for axis in axis_name) + return int(mesh.shape[axis_name]) + + +def _maybe_choose_ns_batch_axis(ns_batch_size, mesh, preferred_axis): + """Return preferred_axis if ns_batch_size is divisible by its mesh size, else None.""" + if preferred_axis is None: + return None + axis_size = _mesh_axis_size(mesh, preferred_axis) + if axis_size > 1 and ns_batch_size % axis_size == 0: + return preferred_axis + return None + + +_NS_FULL_REPLICATE_MAX_BYTES = 256 << 20 # 256 MB + + +def _fits_full_replicate_budget(x): + """Check if full-replicating x for NS won't exceed HBM budget.""" + b, m, n = x.shape + elem_bytes = x.dtype.itemsize if hasattr(x, "dtype") else 2 + min_dim = min(m, n) + working_bytes = (b * m * n + 2 * b * min_dim * min_dim) * elem_bytes + return working_bytes <= _NS_FULL_REPLICATE_MAX_BYTES + + +def _pad_batch_to_multiple(x: jax.Array, multiple: int) -> tuple[jax.Array, int]: + """Pad rank-3 Muon NS input on batch dim to a multiple of mesh axis size.""" + total_batch = x.shape[0] + if multiple <= 1: + return x, total_batch + padded_size = math.ceil(total_batch / multiple) * multiple + if padded_size == total_batch: + return x, total_batch + return jnp.pad(x, ((0, padded_size - total_batch), (0, 0), (0, 0))), total_batch + + +def _maybe_restore_sharding(x: jax.Array, sharding) -> jax.Array: + """Restore original sharding after NS, if a saved sharding is available.""" + if sharding is None: + return x + return jax.lax.with_sharding_constraint(x, sharding) + + +def _batch_orthogonalize_tree( + updates: base.Updates, + weight_dim_nums: base.Params, + ns_coeffs: jax.Array, + ns_steps: jax.typing.ArrayLike, + preconditioning: str, + eps: jax.typing.ArrayLike, + batch_update_size: int | None, + ns_unroll: bool, + replicate_ns_matrix_axes: bool = False, + replicate_ns_batch_axis: str | None = None, +) -> base.Updates: + """Orthogonalize updates by grouping same-shape matrices and batching NS. + + Instead of running Newton-Schulz on each parameter independently, groups + parameters by their reshaped 2D matrix shape (m, n), stacks them along the + batch dimension, and runs NS on the combined batch. This amortizes overhead + for many small matrices. Results are numerically identical to unbatched mode. + + Args: + updates: PyTree of gradient updates (already momentum-processed). + weight_dim_nums: PyTree of MuonDimensionNumbers (same structure as updates). + ns_coeffs: Newton-Schulz coefficients array. + ns_steps: Number of NS iterations. + preconditioning: Preconditioning method. + eps: Numerical stability epsilon. + batch_update_size: If set, chunk each shape group into batches of this size. + + Returns: + PyTree of orthogonalized updates (same structure as input). + """ + flat_updates, treedef = jax.tree.flatten(updates, is_leaf=_is_weight_dim_nums) + + ns_mesh = None + if replicate_ns_matrix_axes: + if replicate_ns_batch_axis is None: + raise ValueError("replicate_ns_batch_axis must be set when replicate_ns_matrix_axes=True") + ns_mesh = _get_current_mesh_for_ns() + + # Handle weight_dim_nums=None: create matching default dim nums for 2D params. + if weight_dim_nums is None: + flat_dim_nums = [None] * len(flat_updates) + else: + flat_dim_nums, _ = jax.tree.flatten(weight_dim_nums, is_leaf=_is_weight_dim_nums) + + # For each leaf, reshape to (batch, m, n) and record metadata for regrouping. + # Group key is (m, n) after ensuring m <= n. + groups: dict[tuple[int, int], list] = {} + component_results: dict[int, dict[str, Any]] = {} + results = [None] * len(flat_updates) + + for i, (leaf, dim_num) in enumerate(zip(flat_updates, flat_dim_nums)): + if leaf.ndim < 2: + # Scalar/1D params shouldn't reach here (they go to Adam), but be safe. + results[i] = leaf + continue + + if dim_num is not None and getattr(dim_num, "component_splits", None) is not None: + reduction_axes, output_axes = _normalize_axes(leaf, dim_num) + batch_axes = tuple(sorted(set(range(leaf.ndim)) - set(reduction_axes) - set(output_axes))) + splits = dim_num.component_splits + + if len(output_axes) != 1: + raise ValueError(f"component_splits requires exactly 1 output axis, got {output_axes}") + output_axis = output_axes[0] + + if sum(splits) != leaf.shape[output_axis]: + raise ValueError(f"component_splits {splits} must sum to output axis size {leaf.shape[output_axis]}") + + perm = tuple(reduction_axes) + tuple(batch_axes) + (output_axis,) + inv_perm = tuple(sorted(range(len(perm)), key=lambda j, _perm=perm: _perm[j])) + leaf_t = jnp.transpose(leaf, perm) + n_reduction = len(reduction_axes) + n_batch = len(batch_axes) + reduction_size = math.prod(leaf_t.shape[:n_reduction]) + batch_size = math.prod(leaf_t.shape[n_reduction : n_reduction + n_batch]) + split_indices = list(itertools.accumulate(splits[:-1])) + component_results[i] = { + "shape": leaf_t.shape, + "splits": splits, + "inv_perm": inv_perm, + "parts": [], + "original_sharding": _get_array_sharding(leaf) if replicate_ns_matrix_axes else None, + } + + for component_index, comp in enumerate(jnp.split(leaf_t, split_indices, axis=-1)): + comp_size = comp.shape[-1] + reshaped = comp.reshape(reduction_size, batch_size * comp_size) + m, n = reshaped.shape + transposed = m > n + if transposed: + reshaped = reshaped.T + m, n = n, m + groups.setdefault((m, n), []).append( + { + "idx": i, + "component_index": component_index, + "reshaped": reshaped[jnp.newaxis, ...], + "transposed": transposed, + "batch_size": 1, + "component_shape": leaf_t.shape[: n_reduction + n_batch] + (comp_size,), + } + ) + continue + + if leaf.ndim == 2: + dn = MuonDimensionNumbers(reduction_axis=0, output_axis=1) + else: + dn = dim_num + + reshape_fn, inverse_fn = _compute_muon_reshape(leaf, dn) + reshaped = reshape_fn(leaf) # (batch, m, n) + _, m, n = reshaped.shape + transposed = m > n + if transposed: + reshaped = reshaped.transpose(0, 2, 1) + m, n = n, m + + key = (m, n) + if key not in groups: + groups[key] = [] + groups[key].append( + { + "idx": i, + "reshaped": reshaped, + "inverse_fn": inverse_fn, + "transposed": transposed, + "batch_size": reshaped.shape[0], + "original_sharding": _get_array_sharding(leaf) if replicate_ns_matrix_axes else None, + } + ) + + # Component-split params are pre-flattened to 3D before reaching this path, + # so chunking can use a fixed (batch, m, n) dim spec for every group. + dim_num_chunk = MuonDimensionNumbers(reduction_axis=1, output_axis=2) + + def _ns_chunk(_, chunk): + return _, orthogonalize_via_newton_schulz( + chunk, + ns_coeffs, + ns_steps, + preconditioning, + eps, + dim_num_chunk, + ns_unroll=ns_unroll, + replicate_ns_matrix_axes=replicate_ns_matrix_axes, + replicate_ns_batch_axis=replicate_ns_batch_axis, + ) + + # Process each shape group. + for (m, n), items in groups.items(): + # Concatenate all items along batch dim. + all_reshaped = jnp.concatenate([item["reshaped"] for item in items], axis=0) + + # Chunk if batch_update_size is set. Use lax.scan so all chunks share a + # single traced NS body in HLO instead of N independent fori_loops, which + # cuts HLO instruction count substantially when groups have many leaves. + # Safe here because component-split params are pre-flattened to 3D before + # reaching this path, so the chunking no longer composes with a nested + # vmap structure that previously triggered an XLA layout assertion. + total_batch = all_reshaped.shape[0] + + # Auto-cap batch size when replicate_ns_matrix_axes is on. After + # replication, matrix dims (m, n) are fully materialized on each device + # instead of being FSDP-sharded. NS peak memory per matrix in the batch: + # input(m*n) + gram a=x@xT(m*m) + intermediates(m*m) ≈ m*n + 2*m*m. + # Cap total working set to NS_REPLICATE_SAFETY_BYTES so that + # batch_update_size=None doesn't OOM on large meshes. + effective_batch_size = batch_update_size + if replicate_ns_matrix_axes and batch_update_size is None: + NS_REPLICATE_SAFETY_BYTES = 1 << 30 # 1 GB + elem_bytes = all_reshaped.dtype.itemsize + per_matrix_bytes = (m * n + 2 * m * m) * elem_bytes + raw_cap = max(1, NS_REPLICATE_SAFETY_BYTES // per_matrix_bytes) + # Round down to fsdp multiple so each chunk can use P(fsdp, None, None) + # instead of falling back to P(None, None, None) full-replicate. + # Without this, e.g. raw_cap=85 with fsdp=64 produces chunks of size 85 + # which aren't divisible → _maybe_choose_ns_batch_axis returns None → + # full-replicate → 6x padding expansion → OOM. + axis_size = _mesh_axis_size(ns_mesh, replicate_ns_batch_axis) + if 1 < axis_size <= raw_cap: + auto_cap = (raw_cap // axis_size) * axis_size + else: + auto_cap = raw_cap + if effective_batch_size is None: + logging.info( + "Muon NS auto-cap: matrix=(%d,%d), per_matrix_bytes=%d, raw_cap=%d, " + "aligned_cap=%d, fsdp_size=%d, batch_update_size=%s", + m, + n, + per_matrix_bytes, + raw_cap, + auto_cap, + axis_size, + batch_update_size, + ) + effective_batch_size = auto_cap + + if effective_batch_size is not None and total_batch > effective_batch_size: + num_chunks = math.ceil(total_batch / effective_batch_size) + padded_size = num_chunks * effective_batch_size + if padded_size > total_batch: + padded = jnp.pad(all_reshaped, ((0, padded_size - total_batch), (0, 0), (0, 0))) + else: + padded = all_reshaped + chunked = padded.reshape(num_chunks, effective_batch_size, m, n) + # Do NOT apply rank-4 constraint here. The scan body applies the + # rank-3 P(batch_axis, None, None) constraint inside each iteration, + # which is what actually eliminates per-NS-step all-reduce on the + # contraction dim. Adding a rank-4 P(fsdp, None, None, None) on top + # conflicts with the inner constraint at scan boundaries and causes + # XLA to allocate huge resharding staging buffers (OOM on 256 chips). + _, all_ortho_chunked = jax.lax.scan(_ns_chunk, None, chunked) + all_ortho = all_ortho_chunked.reshape(padded_size, m, n)[:total_batch] + else: + if replicate_ns_matrix_axes: + all_reshaped, unpadded_batch = _pad_batch_to_multiple( + all_reshaped, _mesh_axis_size(ns_mesh, replicate_ns_batch_axis) + ) + all_reshaped = _replicate_matrix_axes_for_ns(all_reshaped, ns_mesh, replicate_ns_batch_axis) + all_ortho = orthogonalize_via_newton_schulz( + all_reshaped, + ns_coeffs, + ns_steps, + preconditioning, + eps, + dim_num_chunk, + ns_unroll=ns_unroll, + replicate_ns_matrix_axes=replicate_ns_matrix_axes, + replicate_ns_batch_axis=replicate_ns_batch_axis, + ) + if replicate_ns_matrix_axes: + all_ortho = all_ortho[:unpadded_batch] + + # Split back to individual items. + batch_sizes = [item["batch_size"] for item in items] + split_indices = list(itertools.accumulate(batch_sizes[:-1])) + split_results = jnp.split(all_ortho, split_indices, axis=0) + + for item, ortho in zip(items, split_results): + if item["transposed"]: + ortho = ortho.transpose(0, 2, 1) + if "inverse_fn" in item: + restored = item["inverse_fn"](ortho) + restored = _maybe_restore_sharding(restored, item.get("original_sharding")) + results[item["idx"]] = restored + else: + component_results[item["idx"]]["parts"].append( + (item["component_index"], ortho[0].reshape(item["component_shape"])) + ) + + for idx, metadata in component_results.items(): + parts = [part for _, part in sorted(metadata["parts"], key=lambda x: x[0])] + result_t = jnp.concatenate(parts, axis=-1).reshape(metadata["shape"]) + result = jnp.transpose(result_t, metadata["inv_perm"]) + result = _maybe_restore_sharding(result, metadata.get("original_sharding")) + results[idx] = result + + return treedef.unflatten(results) + + +def scale_by_muon( + ns_coeffs: Union[ + tuple[jax.typing.ArrayLike, jax.typing.ArrayLike, jax.typing.ArrayLike], + tuple[ + tuple[jax.typing.ArrayLike, jax.typing.ArrayLike, jax.typing.ArrayLike], + ..., + ], + ] = _DEFAULT_NS_COEFFS, + ns_steps: jax.typing.ArrayLike = 5, + beta: jax.typing.ArrayLike = 0.95, + eps: jax.typing.ArrayLike = 1e-7, + mu_dtype: Optional[jax.typing.DTypeLike] = None, + ns_dtype: Optional[jax.typing.DTypeLike] = jnp.float32, + *, + nesterov: bool = True, + nesterov_style: Literal["ema", "sgd"] = "ema", + adaptive: bool = False, + preconditioning: Literal["frobenius", "spectral", "aol", "schatten"] = "frobenius", + weight_dimension_numbers: WeightDimNumOrFn | None = None, + batch_update: bool = True, + batch_update_size: int | None = 16, + ns_unroll: bool = False, + replicate_ns_matrix_axes: bool = False, + replicate_ns_batch_axis: str | None = None, +) -> base.GradientTransformation: + r"""Rescale updates according to the Muon algorithm. + + Muon is a variant of Shampoo that uses the Newton-schulz method to + orthogonalize the momentum accumulated by the optimizer. Mathematically, it + does steepest descent under the Schatten-p norm, for some large p. With + p=infty, it is equivalent to Shampoo without accumulation, or steepest + descent under the Spectral norm. + + Args: + ns_coeffs: Coefficients for the Newton-schulz method. + ns_steps: Number of Newton-schulz iterations. + Ignored if `ns_coeffs` is a tuple of tuples. + beta: Decay rate for the exponentially weighted average of grads. + eps: Term added to denominators to improve numerical stability. + Default 1e-7 matches Megatron-LM's hardcoded NS normalization epsilon. + mu_dtype: Data type of the momentum accumulator. + ns_dtype: Data type for Newton-Schulz orthogonalization. Input is cast to + this dtype before NS iterations. The output stays in ns_dtype and flows + through subsequent transforms; JAX promotes to fp32 at param update. + Default ``jnp.float32`` to keep the five-step NS polynomial in fp32 and + remove TPU/GPU bf16 matmul-order drift across accelerators. Set to + ``jnp.bfloat16`` to mirror Megatron-LM's default (bf16 NS output scaled + by LR in bf16, added to fp32 params), or ``None`` to skip casting (use + the momentum dtype as-is). + nesterov: Whether to use Nesterov momentum. + nesterov_style: Style of momentum accumulation. 'ema' uses the optax + default (EMA accumulation with bias correction). 'sgd' uses classic + SGD-style momentum (buf = beta*buf + grad) with Nesterov look-ahead + (g = grad + beta*buf) and no bias correction, matching Megatron-LM. + adaptive: Whether to scale the updates by the dual norm of the + original updates. See + preconditioning: What type of preconditioning to use before NS iterations. + Available options are: + - 'frobenius' (default): Use Frobenius rescaling before NS. + - 'spectral' : Use Spectral norm rescaling before NS. + - 'aol': Use AOL rescaling to improve orthogonality. + - 'schatten': Use the Schatten-4 norm for rescaling. + weight_dimension_numbers: An optional tree with the same structure as the + params of `MuonDimensionNumbers`s, specifying how to reshape the + parameters before and after the orthogonalization OR a callable returning + such a tree. None implies that all parameters are 2D matrices. + batch_update: If True, group parameters by their reshaped 2D matrix shape + and stack them along the batch dimension before running Newton-Schulz. + This amortizes overhead for many small matrices. Results are numerically + identical to unbatched mode. + batch_update_size: If set (and batch_update is True), chunk each shape + group into batches of at most this size before running Newton-Schulz. + replicate_ns_batch_axis: Mesh axis name for batch dim when + replicate_ns_matrix_axes=True. Must be set when enabling the flag. + + Returns: + A `GradientTransformation` object. + + References: + Jordan, `modded-nanogpt: Speedrunning the NanoGPT baseline + `_, 2024 + + Bernstein et al., `Old Optimizer, New Norm: An Anthology + `_, 2024 + + Liu et al., `Muon is Scalable for LLM Training`, + `_, 2025 + + Boissin et al., `Turbo-Muon: Accelerating Orthogonality-Based + Optimization with Pre-Conditioning`, + `_, 2025 + + Ahn et al., `Dion: Distributed Orthonormalized Updates`, + `_, 2025 + + Grishina et al., `Accelerating Newton-Schulz Iteration for Orthogonalization + via Chebyshev-type Polynomials`, + `_, 2025 + + Amsel et al., `The Polar Express: Optimal Matrix Sign Methods and Their + Application to the Muon Algorithm`, + `, 2025 + """ + mu_dtype = utils.canonicalize_dtype(mu_dtype) + + def init_fn(params): + mu = optax.tree.zeros_like(params, dtype=mu_dtype) # First moment + ns_coeffs_ = jnp.asarray(ns_coeffs) + + if ns_coeffs_.ndim > 2 or ns_coeffs_.shape[-1] != 3: + raise ValueError(f"ns_coeffs must have shape (3,) or (n, 3), got {ns_coeffs_.shape}") + if ns_coeffs_.ndim == 2: + if not ns_coeffs_.shape[0] <= ns_steps: + raise ValueError(f"Not enough coeffs to perform {ns_steps} steps") + ns_coeffs_ = ns_coeffs_[-ns_steps:] + + return MuonState( + count=jnp.zeros([], jnp.int32), + mu=mu, + ns_coeffs=ns_coeffs_, + ) + + def update_fn(updates, state, params=None): + del params + # TODO(rdyro): extend to _masking._mask_callable + if callable(weight_dimension_numbers): + # Populate weight_dim_nums if it's a callable. Use updates instead of + # actual params since only shapes matter and params may not be provided. + resolved_weight_dim_nums = weight_dimension_numbers(updates) + else: + resolved_weight_dim_nums = weight_dimension_numbers + + if nesterov_style == "sgd": + # Classic SGD-style momentum (Megatron-LM): + # buf = beta * buf + grad + # g = grad + beta * buf (Nesterov look-ahead, no bias correction) + mu = jax.tree.map(lambda g, m: beta * m + g, updates, state.mu) + if nesterov: + mu_hat = jax.tree.map(lambda g, m: g + beta * m, updates, mu) + else: + mu_hat = mu + else: + # Default optax EMA-style momentum with bias correction + mu = optax.tree.update_moment(updates, state.mu, beta, 1) + count_inc = numerics.safe_increment(state.count) + if nesterov: + mu_hat = jax.tree.map( + lambda m, g: beta * m + (1 - beta) * g, + optax.tree.bias_correction(mu, beta, numerics.safe_increment(count_inc)), + optax.tree.bias_correction(updates, beta, count_inc), + ) + else: + mu_hat = optax.tree.bias_correction(mu, beta, count_inc) + + count_inc = numerics.safe_increment(state.count) + # Cast to ns_dtype before Newton-Schulz orthogonalization, then let the + # update stay in ns_dtype through the rest of the optax chain + # (scale_by_shape, add_decayed_weights, scale_by_learning_rate). JAX + # promotes to fp32 only at param addition. Default ns_dtype is fp32 so + # the five-step NS polynomial stays in fp32 — Megatron-LM casts to + # bfloat16 here (g.bfloat16()) and keeps the update in bf16 through LR + # scaling, which is available by passing ns_dtype=jnp.bfloat16. + if ns_dtype is not None: + ns_input = optax.tree.cast(mu_hat, ns_dtype) + else: + ns_input = mu_hat + # Apply Newton-schulz orthogonalization. + if batch_update: + updates = _batch_orthogonalize_tree( + ns_input, + resolved_weight_dim_nums, + state.ns_coeffs, + ns_steps, + preconditioning, + eps, + batch_update_size, + ns_unroll, + replicate_ns_matrix_axes, + replicate_ns_batch_axis, + ) + else: + # Save original leaf shardings for restore after NS. + if replicate_ns_matrix_axes: + saved_shardings = jax.tree.map(_get_array_sharding, ns_input, is_leaf=_is_weight_dim_nums) + updates = jax.tree.map( + lambda x, dim_num: orthogonalize_via_newton_schulz( + x, + state.ns_coeffs, + ns_steps, + preconditioning, + eps, + dim_num, + ns_unroll=ns_unroll, + replicate_ns_matrix_axes=replicate_ns_matrix_axes, + replicate_ns_batch_axis=replicate_ns_batch_axis, + ), + ns_input, + resolved_weight_dim_nums, + is_leaf=_is_weight_dim_nums, + ) + if replicate_ns_matrix_axes: + updates = jax.tree.map(_maybe_restore_sharding, updates, saved_shardings, is_leaf=_is_weight_dim_nums) + if adaptive: + # Scale the orthogonalized updates by the dual norm of the original + # updates. See https://arxiv.org/abs/2409.20325 for the derivation. + updates = jax.tree.map(lambda x, y: jnp.sum(x.conj() * y) * y, mu_hat, updates) + + mu = optax.tree.cast(mu, mu_dtype) + return updates, MuonState( + count=count_inc, + mu=mu, + ns_coeffs=state.ns_coeffs, + ) + + return base.GradientTransformation(init_fn, update_fn) + + +def muon( + learning_rate: base.ScalarOrSchedule, + ns_coeffs: Union[ + tuple[jax.typing.ArrayLike, jax.typing.ArrayLike, jax.typing.ArrayLike], + tuple[ + tuple[jax.typing.ArrayLike, jax.typing.ArrayLike, jax.typing.ArrayLike], + ..., + ], + str, + ] = _DEFAULT_NS_COEFFS, + ns_steps: jax.typing.ArrayLike = 5, + beta: jax.typing.ArrayLike = 0.95, + eps: jax.typing.ArrayLike = 1e-7, + weight_decay: jax.typing.ArrayLike = 0.0, + weight_decay_mask: Optional[Union[Any, Callable[[base.Params], Any]]] = None, + mu_dtype: Optional[jax.typing.DTypeLike] = None, + ns_dtype: Optional[jax.typing.DTypeLike] = jnp.float32, + *, + nesterov: bool = True, + nesterov_style: Literal["ema", "sgd"] = "ema", + adaptive: bool = False, + preconditioning: Literal["frobenius", "spectral", "aol", "schatten"] = "frobenius", + adam_b1: jax.typing.ArrayLike = 0.9, + adam_b2: jax.typing.ArrayLike = 0.999, + adam_eps: jax.typing.ArrayLike | None = None, + adam_eps_root: jax.typing.ArrayLike = 0.0, + adam_nesterov: bool = False, + adam_weight_decay: jax.typing.ArrayLike = 0.0, + adam_learning_rate: base.ScalarOrSchedule | None = None, + adam_weight_decay_mask: Optional[Union[Any, Callable[[base.Params], Any]]] = None, + muon_weight_dimension_numbers: WeightDimNumOrFn | None = None, + consistent_rms: jax.typing.ArrayLike | None = None, + batch_update: bool = True, + batch_update_size: int | None = 16, + ns_unroll: bool = False, + replicate_ns_matrix_axes: bool = False, + replicate_ns_batch_axis: str | None = None, +) -> base.GradientTransformation: + r"""Muon: Momentum Orthogonalized by Newton-schulz. + + Muon is a variant of Shampoo that uses the Newton-schulz method to + orthogonalize the momentum accumulated by the optimizer. Mathematically, it + does steepest descent under the Schatten-p norm, for some large p. With + p=infty, it is equivalent to Shampoo without accumulation, or steepest + descent under the Spectral norm. + + Note that Muon is currently only defined for 2D parameters, i.e. matrices. + This is because the Newton-Schulz iterator expects a matrix as input. + The non-2D parameters are instead passed through an AdamW optimizer + (using a weight decay of 0 as default). + + Args: + learning_rate: A global scaling factor, either fixed or evolving along + iterations with a scheduler, see :func:`optax.scale_by_learning_rate`. + ns_coeffs: Coefficients for the Newton-schulz method (can be a string + indicator for a preset). Existing presets: `muon`, `dion`. + ns_steps: Number of Newton-schulz iterations. + Ignored if `ns_coeffs` is a tuple of tuples. + beta: Decay rate for the exponentially weighted average of grads. + eps: Term added to the denominator to improve numerical stability. + Default 1e-7 matches Megatron-LM's hardcoded NS normalization epsilon. + weight_decay: Strength of the weight decay regularization. Note that this + weight decay is multiplied with the learning rate. This is consistent + with other frameworks such as PyTorch, but different from + (Loshchilov et al, 2019) where the weight decay is only multiplied with + the "schedule multiplier", but not the base learning rate. + weight_decay_mask: A tree with same structure as (or a prefix of) the params + PyTree, or a Callable that returns such a pytree given the params/updates. + The leaves should be booleans, `True` for leaves/subtrees you want to + apply the weight decay to, and `False` for those you want to skip. + mu_dtype: Data type of the momentum accumulator. + ns_dtype: Data type for Newton-Schulz orthogonalization. Input is cast to + this dtype before NS iterations. The output stays in ns_dtype through + LR scaling; JAX promotes to fp32 at param update. Default + ``jnp.float32`` keeps the NS polynomial in fp32 to avoid TPU/GPU bf16 + matmul-order drift. Set to ``jnp.bfloat16`` to match Megatron-LM's + bf16 NS, or ``None`` to skip casting. + nesterov: Whether to use Nesterov momentum. + nesterov_style: Style of momentum accumulation. 'ema' uses the optax + default (EMA accumulation with bias correction). 'sgd' uses classic + SGD-style momentum (buf = beta*buf + grad) with Nesterov look-ahead + (g = grad + beta*buf) and no bias correction, matching Megatron-LM. + adaptive: Whether to scale the updates by the dual norm of the + original updates. See + preconditioning: What type of preconditioning to use before NS iterations. + Available options are: + - 'frobenius' (default): Use Frobenius rescaling before NS: + safe, standard, but degrades orthogonalization quality when using + less than 5 NS steps. + - 'spectral' : Use Spectral norm rescaling before NS: + much more computationally intensive, but better orthogonalization + quality. + - 'aol': Use AOL rescalings to improve orthogonality with little to + no overhead, usually allows the user to remove one iterative NS step. + See . + - 'schatten': Use the Schatten-4 norm for rescaling, + allows for better performance with little to no extra cost. + See . + adam_b1: Exponential decay rate for Adam's first moment estimates. + adam_b2: Exponential decay rate for Adam's second moment estimates. + adam_nesterov: Whether to use Nesterov momentum in the Adam partition. + Default ``False`` to match Megatron-LM's standard Adam (no Nesterov). + adam_eps_root: Epsilon to stabilize division in Adam, square root version. + adam_weight_decay: Weight decay factor for Adam. + adam_learning_rate: Auxiliary learning rate for the Adam optimizer. + If `None`, the learning rate for Adam defaults to the same as Muon. + adam_weight_decay_mask: A tree with same structure as (or a prefix of) the + params PyTree, or a Callable that returns such a pytree given the params. + The leaves should be booleans, `True` for leaves/subtrees you want to + apply Adam weight decay to, and `False` for those you want to skip. + muon_weight_dimension_numbers: An optional tree of `MuonDimensionNumbers`s, + specifying how to reshape the parameters for orthogonalization otherwise + muon parameters are assumed to be 2D matrices. A `None` value indicates + that the parameter is not a muon parameter and will be optimized with + Adam. A callable takes as input the params and returns a possibly masked + pytree of specs, similar to `weight_decay_mask`. If not provided, muon is + applied to all 2D parameters. + consistent_rms: An optional float to activate consistent RMS scaling. + Scales updates by `sqrt(max(fan_in, fan_out)) * consistent_rms` to make + root mean square (RMS) shape-independent, like AdamW. `0.2` is recommended + to match AdamW's empirical RMS. See . + If `None`, uses width scaling `sqrt(max(1, fan_out / fan_in))`. + batch_update: If True, group same-shape Muon parameters and run + Newton-Schulz on stacked batches for efficiency. Numerically identical + to unbatched mode. + batch_update_size: Maximum batch size per NS call when batch_update is True. + If None, all same-shape parameters are stacked into one batch. + + Returns: + The corresponding `GradientTransformation`. + + References: + Jordan, `modded-nanogpt: Speedrunning the NanoGPT baseline + `_, 2024 + + Bernstein et al., `Old Optimizer, New Norm: An Anthology + `_, 2024 + + Liu et al., `Muon is Scalable for LLM Training`, + `_, 2025 + + Boissin et al., `Turbo-Muon: Accelerating Orthogonality-Based + Optimization with Pre-Conditioning`, + `_, 2025 + + Ahn et al., `Dion: Distributed Orthonormalized Updates`, + `_, 2025 + + Grishina et al., `Accelerating Newton-Schulz Iteration for Orthogonalization + via Chebyshev-type Polynomials`, + `_, 2025 + + Amsel et al., `The Polar Express: Optimal Matrix Sign Methods and Their + Application to the Muon Algorithm`, + `, 2025 + """ + + if adam_learning_rate is None: + adam_learning_rate = learning_rate + + if isinstance(ns_coeffs, str): + if ns_coeffs not in _NS_COEFFS_PRESET_DICT: + raise ValueError(f"Unknown ns_coeff preset string: {ns_coeffs}") + ns_coeffs_ = _NS_COEFFS_PRESET_DICT[ns_coeffs] + else: + ns_coeffs_ = ns_coeffs + + # None at root indicates the default 2D rule. + if muon_weight_dimension_numbers is None: + param_labels = lambda params: jax.tree.map(lambda x: "muon" if x.ndim == 2 else "adam", params) + muon_weight_dimension_numbers = MuonDimensionNumbers() + else: + + def param_labels(params): + dim_nums = ( + muon_weight_dimension_numbers(params) + if callable(muon_weight_dimension_numbers) + else muon_weight_dimension_numbers + ) + populate_subtree_ = lambda dim_num, x: jax.tree.map(lambda y: "muon" if dim_num is not None else "adam", x) + # Dimension numbers come first since they can be a prefix mask. + return jax.tree.map(populate_subtree_, dim_nums, params, is_leaf=lambda x: x is None or _is_weight_dim_nums(x)) + + # We need to normalize the dimension numbers because they have to match the + # tree structure of the masked muon state tree (see `combine.partition`). + def muon_weight_dim_nums_fn(params): + # if muon_weight_dimension_numbers is None: + # return None + # Normalize the dimension numbers for `combine.partition`. + # Insert MaskedNode() where muon state will be masked out. + dim_nums = ( + muon_weight_dimension_numbers(params) + if callable(muon_weight_dimension_numbers) + else muon_weight_dimension_numbers + ) + mask = jax.tree.map(lambda label: label == "muon", param_labels(params)) + is_leaf = lambda x: (x is None or _is_weight_dim_nums(x) or isinstance(x, _masking.MaskedNode)) + populate_subtree_ = lambda dim_nums, submask: jax.tree.map( + lambda m: dim_nums if m else _masking.MaskedNode(), submask + ) + return jax.tree.map(populate_subtree_, dim_nums, mask, is_leaf=is_leaf) + + return combine.partition( + transforms={ + "muon": combine.chain( + scale_by_muon( + ns_coeffs=ns_coeffs_, + ns_steps=ns_steps, + beta=beta, + eps=eps, + mu_dtype=mu_dtype, + ns_dtype=ns_dtype, + nesterov=nesterov, + nesterov_style=nesterov_style, + adaptive=adaptive, + preconditioning=preconditioning, + weight_dimension_numbers=muon_weight_dim_nums_fn, + batch_update=batch_update, + batch_update_size=batch_update_size, + ns_unroll=ns_unroll, + replicate_ns_matrix_axes=replicate_ns_matrix_axes, + replicate_ns_batch_axis=replicate_ns_batch_axis, + ), + scale_by_shape( + weight_dimension_numbers=muon_weight_dim_nums_fn, + consistent_rms=consistent_rms, + ), + transform.add_decayed_weights(weight_decay, weight_decay_mask), + transform.scale_by_learning_rate(learning_rate), + ), + "adam": alias.adamw( + learning_rate=adam_learning_rate, + b1=adam_b1, + b2=adam_b2, + eps=adam_eps if adam_eps is not None else eps, + eps_root=adam_eps_root, + weight_decay=adam_weight_decay, + mask=adam_weight_decay_mask, + mu_dtype=mu_dtype, + nesterov=adam_nesterov, + ), + }, + param_labels=param_labels, + ) From 056abb15b446ffe6c7c8df8394d061db89caac16 Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Fri, 7 Aug 2026 18:01:00 +0000 Subject: [PATCH 08/12] integrate muon --- src/maxtext/optimizers/optimizers.py | 4 ++-- src/maxtext/utils/muon_utils.py | 2 +- tests/unit/muon_utils_test.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/maxtext/optimizers/optimizers.py b/src/maxtext/optimizers/optimizers.py index 67e1f589ca..3c76e843f6 100644 --- a/src/maxtext/optimizers/optimizers.py +++ b/src/maxtext/optimizers/optimizers.py @@ -20,7 +20,7 @@ import jax.numpy as jnp import optax -from optax.contrib._muon import muon +from maxtext.optimizers.muon.muon import muon from maxtext.common.common_types import DecoderBlockType from maxtext.utils.muon_utils import get_muon_weight_dimension_numbers @@ -215,7 +215,7 @@ def get_optimizer(config, learning_rate_schedule, model=None): "beta": config.muon_beta, "weight_decay": config.muon_weight_decay, "muon_weight_dimension_numbers": muon_weight_dimension_numbers, - "consistent_rms": config.muon_consistent_rms, + "consistent_rms": getattr(config, "muon_consistent_rms", None), "ns_coeffs": ns_coeffs, "ns_steps": ns_steps, # AdamW-specific parameters diff --git a/src/maxtext/utils/muon_utils.py b/src/maxtext/utils/muon_utils.py index 9264a7f2a3..c1ce8489aa 100644 --- a/src/maxtext/utils/muon_utils.py +++ b/src/maxtext/utils/muon_utils.py @@ -36,7 +36,7 @@ from maxtext.layers import quantizations from maxtext.models import models from maxtext.utils import maxtext_utils, model_creation_utils -from optax.contrib._muon import MuonDimensionNumbers as mdn +from maxtext.optimizers.muon.muon import MuonDimensionNumbers as mdn def _is_path_contain_any(tuples, path): diff --git a/tests/unit/muon_utils_test.py b/tests/unit/muon_utils_test.py index a45f910ecf..b4b6995095 100644 --- a/tests/unit/muon_utils_test.py +++ b/tests/unit/muon_utils_test.py @@ -26,7 +26,7 @@ import jax.numpy as jnp from flax import linen as nn from flax import nnx -from optax.contrib._muon import MuonDimensionNumbers as mdn +from maxtext.optimizers.muon.muon import MuonDimensionNumbers as mdn from maxtext.utils import muon_utils From a9a6b33df725bc123336ec8574eeedc4b28acdda Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Fri, 7 Aug 2026 19:02:24 +0000 Subject: [PATCH 09/12] use adamw instead of muon for expert weights --- src/maxtext/utils/muon_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/maxtext/utils/muon_utils.py b/src/maxtext/utils/muon_utils.py index c1ce8489aa..b00fc62b5f 100644 --- a/src/maxtext/utils/muon_utils.py +++ b/src/maxtext/utils/muon_utils.py @@ -96,6 +96,7 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]: "conv1d", "gate", "shared_expert_gate", + "routed_experts", ) ) or segment == "bias" @@ -108,7 +109,7 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]: # L (optional) stands for layer when scan_layers=True if _is_path_contain_any(("MoeBlock_0", "routed_experts"), path): # exclude gate - if _is_path_contain_any(("wi_0", "wi_1", "wo"), path): + if False: return mdn((-2,), (-1,)) # 2.2 Special weights: Self attention From 6506ebaaae4e83d52726d873df2c23eb18f9ed1f Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Fri, 7 Aug 2026 19:02:47 +0000 Subject: [PATCH 10/12] explicitly_weight_ag for shard exp on fsdp --- src/maxtext/layers/moe.py | 20 ++++---------------- 1 file changed, 4 insertions(+), 16 deletions(-) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 6ae9abf2cc..70d85b056d 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -1554,11 +1554,7 @@ def is_batch_sharded_by_ep(input_activation): return input_activation.shape[0] > 1 def explicitly_weight_ag(shard_exp_on_fsdp): - if shard_exp_on_fsdp: - quantization_rule = qpl.get_current_rule("gmm") - if quantization_rule and quantization_rule.weight_calibration_method.startswith("fixed"): - return True - return False + return bool(shard_exp_on_fsdp) def maybe_aqt_partition(w0_kernel, w0_pspec, w1_kernel, w1_pspec, wo_kernel, wo_pspec): if isinstance(w0_kernel, aqt.QTensor): @@ -1596,17 +1592,9 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): # w0, w1, wo needs to be un sharded on fsdp / fsdp_transpose axis, so use # mlp_no_fsdp axis if self.config.shard_exp_on_fsdp: - quantization_rule = qpl.get_current_rule("gmm") - if quantization_rule and quantization_rule.weight_calibration_method.startswith("fixed"): - # special sharding when using static scaling for weights in quantization with shard_exp_on_fsdp - w0_pspec = self._logical_to_mesh_axes(self.wi_kernel_axes) - w1_pspec = self._logical_to_mesh_axes(self.wi_kernel_axes) - wo_pspec = self._logical_to_mesh_axes(self.wo_kernel_axes) - else: - # special sharding for dsv3 to remove overhead between gmm/AG - w0_pspec = self._logical_to_mesh_axes((None, None, "mlp_no_fsdp")) - w1_pspec = self._logical_to_mesh_axes((None, None, "mlp_no_fsdp")) - wo_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None)) + w0_pspec = self._logical_to_mesh_axes(self.wi_kernel_axes) + w1_pspec = self._logical_to_mesh_axes(self.wi_kernel_axes) + wo_pspec = self._logical_to_mesh_axes(self.wo_kernel_axes) elif self.config.use_2d_fsdp_sharding: w0_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None)) w1_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None)) From 36efee5c614252c4468c61b4102f7177a73e0797 Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Fri, 7 Aug 2026 21:26:00 +0000 Subject: [PATCH 11/12] update default nesterov style --- src/maxtext/optimizers/muon/muon.py | 4 ++-- src/maxtext/optimizers/optimizers.py | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/maxtext/optimizers/muon/muon.py b/src/maxtext/optimizers/muon/muon.py index 62de8597b6..b92152c76b 100644 --- a/src/maxtext/optimizers/muon/muon.py +++ b/src/maxtext/optimizers/muon/muon.py @@ -1031,7 +1031,7 @@ def scale_by_muon( ns_dtype: Optional[jax.typing.DTypeLike] = jnp.float32, *, nesterov: bool = True, - nesterov_style: Literal["ema", "sgd"] = "ema", + nesterov_style: Literal["ema", "sgd"] = "sgd", adaptive: bool = False, preconditioning: Literal["frobenius", "spectral", "aol", "schatten"] = "frobenius", weight_dimension_numbers: WeightDimNumOrFn | None = None, @@ -1252,7 +1252,7 @@ def muon( ns_dtype: Optional[jax.typing.DTypeLike] = jnp.float32, *, nesterov: bool = True, - nesterov_style: Literal["ema", "sgd"] = "ema", + nesterov_style: Literal["ema", "sgd"] = "sgd", adaptive: bool = False, preconditioning: Literal["frobenius", "spectral", "aol", "schatten"] = "frobenius", adam_b1: jax.typing.ArrayLike = 0.9, diff --git a/src/maxtext/optimizers/optimizers.py b/src/maxtext/optimizers/optimizers.py index 3c76e843f6..d728ef7a1a 100644 --- a/src/maxtext/optimizers/optimizers.py +++ b/src/maxtext/optimizers/optimizers.py @@ -216,6 +216,7 @@ def get_optimizer(config, learning_rate_schedule, model=None): "weight_decay": config.muon_weight_decay, "muon_weight_dimension_numbers": muon_weight_dimension_numbers, "consistent_rms": getattr(config, "muon_consistent_rms", None), + "nesterov_style": getattr(config, "muon_nesterov_style", "sgd"), "ns_coeffs": ns_coeffs, "ns_steps": ns_steps, # AdamW-specific parameters From 589cd744c7aa67dfda881decb39da9597f5484bf Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Sat, 8 Aug 2026 00:53:33 +0000 Subject: [PATCH 12/12] apply muon to expert weights and update optimizer config --- src/maxtext/optimizers/optimizers.py | 13 +++++++++++++ src/maxtext/utils/muon_utils.py | 5 ++--- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/src/maxtext/optimizers/optimizers.py b/src/maxtext/optimizers/optimizers.py index d728ef7a1a..6d1038306e 100644 --- a/src/maxtext/optimizers/optimizers.py +++ b/src/maxtext/optimizers/optimizers.py @@ -219,6 +219,19 @@ def get_optimizer(config, learning_rate_schedule, model=None): "nesterov_style": getattr(config, "muon_nesterov_style", "sgd"), "ns_coeffs": ns_coeffs, "ns_steps": ns_steps, + # Keeps the Newton-Schulz matmul batch sharded on "fsdp" (the physical + # axis shard_exp_on_fsdp uses for the expert dimension, see + # exp_with_fsdp in base.yml) instead of falling back to full + # replication, which is what blows up HBM for MoE expert weights. + "replicate_ns_matrix_axes": True, + "replicate_ns_batch_axis": "fsdp", + # muon()'s own default (16) doesn't divide the fsdp axis size on this + # mesh, so the auto-cap logic (which only runs when this is None) + # never fires and it silently falls back to full replication. + "batch_update_size": None, + # Halves the replicated NS working set vs the fp32 default; model + # already trains in bfloat16. + "ns_dtype": jnp.bfloat16, # AdamW-specific parameters "adam_b1": config.adam_b1, "adam_b2": config.adam_b2, diff --git a/src/maxtext/utils/muon_utils.py b/src/maxtext/utils/muon_utils.py index b00fc62b5f..4abeae75d5 100644 --- a/src/maxtext/utils/muon_utils.py +++ b/src/maxtext/utils/muon_utils.py @@ -96,7 +96,6 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]: "conv1d", "gate", "shared_expert_gate", - "routed_experts", ) ) or segment == "bias" @@ -107,9 +106,9 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]: # 2 Special weights # 2.1 Special weights: MoE, [0, L, -2, -1] # L (optional) stands for layer when scan_layers=True - if _is_path_contain_any(("MoeBlock_0", "routed_experts"), path): + if "MoeBlock_0" in path: # exclude gate - if False: + if _is_path_contain_any(("wi_0", "wi_1", "wo"), path): return mdn((-2,), (-1,)) # 2.2 Special weights: Self attention