From 75b04a338e4d5a0c22d7adbcaa9d2cefb487c020 Mon Sep 17 00:00:00 2001 From: Muskan Sharma Date: Fri, 4 Sep 2026 17:24:28 +0000 Subject: [PATCH 1/5] Apply changes from CL 974697834: Custom model definition based on qwen3-next --- .dockerignore | 2 + .../dockerfiles/maxtext_runner.Dockerfile | 3 + src/maxtext/configs/base.yml | 11 +- .../shard-exp-on-fsdp.yml | 12 + .../models/qwen3-next-80b-a3b-256e.yml | 87 ++ .../configs/models/qwen3-next-80b-a3b.yml | 65 +- src/maxtext/configs/types.py | 79 +- src/maxtext/kernels/attention/gdn_cp.py | 130 ++ src/maxtext/kernels/megablox/backend.py | 28 +- src/maxtext/kernels/mhc/api.py | 12 +- src/maxtext/kernels/mhc/mhc_kernels_bwd.py | 13 +- src/maxtext/kernels/mhc/mhc_kernels_fwd.py | 4 +- src/maxtext/kernels/ragged/ragged_gather.py | 2 +- .../kernels/ragged/ragged_gather_reduce_v2.py | 2 +- src/maxtext/layers/attention_op.py | 6 +- src/maxtext/layers/mhc.py | 176 ++- src/maxtext/layers/moe.py | 79 +- src/maxtext/layers/nnx_decoders.py | 157 +- src/maxtext/layers/quantizations.py | 6 +- src/maxtext/models/hybrid_gdn.py | 1042 ++++++++++++++ src/maxtext/models/qwen3.py | 1270 ++++++++++++----- src/maxtext/models/qwen3_5.py | 1 + src/maxtext/optimizers/optimizers.py | 16 +- src/maxtext/trainers/pre_train/train.py | 60 +- .../trainers/pre_train/train_compile.py | 4 +- src/maxtext/training_engine/maxtext_engine.py | 20 +- src/maxtext/utils/maxtext_utils.py | 21 +- src/maxtext/utils/maxtext_utils_nnx.py | 120 +- src/maxtext/utils/muon_utils.py | 325 +++-- src/maxtext/utils/train_utils.py | 15 +- tests/unit/gdn_cp_test.py | 289 ++++ tests/unit/mhc_test.py | 257 +++- tests/unit/muon_utils_test.py | 280 +++- tests/unit/nnx_decoders_test.py | 189 ++- tests/unit/optimizers_test.py | 367 +++-- tests/unit/qwen3_next_vs_reference_test.py | 45 + tests/unit/sharding_nnx_test.py | 1 - 37 files changed, 4384 insertions(+), 812 deletions(-) create mode 100644 src/maxtext/configs/models/qwen3-next-80b-a3b-256e.yml create mode 100644 src/maxtext/kernels/attention/gdn_cp.py create mode 100644 src/maxtext/models/hybrid_gdn.py create mode 100644 tests/unit/gdn_cp_test.py diff --git a/.dockerignore b/.dockerignore index e567ea2ff6..fe3d4214b5 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,2 +1,4 @@ .git maxtext_venv +.venv +venv13 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 269339e056..5ae5e38dca 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -304,6 +304,7 @@ use_2d_fsdp_sharding: false # deepseek moe first_num_dense_layers: 0 # number of initial dense layers in the model shared_experts: 0 +moe_shared_expert_gate: false routed_scaling_factor: 1.0 # scaling factor for routing scores routed_score_func: "" # scoring function for routing routed_bias: false # a flag if a learnable bias is added for routing @@ -317,6 +318,8 @@ topk_routing_group: -1 # number of top groups to route inputs. For EP, use_batch_split_schedule: false # a flag if splitting batch into micro-batches to hide communications that yields performance benefits. batch_split_factor: 1 # the factor by which to split the batch. Only used if use_batch_split_schedule is true. +full_attention_layer_offset: 0 + # For complex architectures like llama4 there are repeated sets of # inhomogeneous layers. E.g. maverick uses [dense+rope, moe+rope, dense+rope, moe+nope] # which can only be scanned together in one large block of inhomogeneous_layer_cycle_interval=4 layers. @@ -792,7 +795,7 @@ sft_train_on_completion_only: false # dataset_type must be synthetic, hf, grain, tfds # details in: https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_pipeline.md -dataset_type: tfds +dataset_type: synthetic # for TFDS input pipeline (dataset_type=tfds) dataset_path: "" # your path given as argument in download_dataset.sh, e.g. "gs://my-maxtext-dataset/" dataset_name: 'c4/en:3.1.0' @@ -1019,6 +1022,8 @@ mu_dtype: "" # data type to store "mu" of AdamW tracking the first moment. Inher muon_beta: 0.95 # Decay rate for the exponentially weighted average of grads. muon_weight_decay: 0 # Strength of the weight decay regularization. This is multiplied with the learning rate. muon_consistent_rms: None # If None, apply width scaling to updates. If float, apply consistent rms scaling (recommend 0.2). +muon_ns_steps: 5 # Number of Newton-Schulz iterations for Muon optimizer. +muon_use_all_to_all: true # Whether to use all-to-all communication during Newton-Schulz iterations in Muon. # Use iota operator in Embed @@ -1344,6 +1349,10 @@ gdn_chunk_size: 64 use_qk_norm_in_gdn: true # The ratio of dimension to apply ROPE on partial_rotary_factor: 1.0 +# 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 use_tokamax_splash: false # Setting this flag will use a non-pallas implementation. diff --git a/src/maxtext/configs/custom_mesh_and_rule/shard-exp-on-fsdp.yml b/src/maxtext/configs/custom_mesh_and_rule/shard-exp-on-fsdp.yml index 41e9487852..0d089f893c 100644 --- a/src/maxtext/configs/custom_mesh_and_rule/shard-exp-on-fsdp.yml +++ b/src/maxtext/configs/custom_mesh_and_rule/shard-exp-on-fsdp.yml @@ -38,6 +38,12 @@ logical_axis_rules: [ ['activation_q_length', ['context']], ['activation_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']], # Attention Weights + ['heads', ['fsdp_transpose', 'expert']], + ['q_heads', ['fsdp_transpose', 'expert']], + ['kv_heads', ['fsdp_transpose', 'expert']], + ['qkv', []], + ['kv', []], + ['kv_head_dim', []], ['q_lora', ['fsdp', 'fsdp_transpose', 'expert']], ["q_lora_up_proj", []], ['kv_lora', ['fsdp', 'fsdp_transpose', 'expert']], @@ -68,8 +74,14 @@ logical_axis_rules: [ ['activation_stage', 'stage'], # General Weights ['mlp', ['fsdp_transpose']], + ['gdn_head', ['fsdp_transpose', 'expert']], + ['embed', ['fsdp', 'fsdp_transpose', 'context', 'expert']], ['embed', ['fsdp', 'context', 'expert']], ['embed_attn', ['fsdp', 'context', 'expert']], ['norm', []], ['layers', 'stage'], + ['dense_layers', []], + ['moe_layers', []], + ['local_layers', []], + ['mhc', []], ] diff --git a/src/maxtext/configs/models/qwen3-next-80b-a3b-256e.yml b/src/maxtext/configs/models/qwen3-next-80b-a3b-256e.yml new file mode 100644 index 0000000000..c72bc1801b --- /dev/null +++ b/src/maxtext/configs/models/qwen3-next-80b-a3b-256e.yml @@ -0,0 +1,87 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# maxtext/configs/models/qwen3-next-80b-a3b-256e.yml + +# Set the decoder block to our new implementation +decoder_block: "qwen3_next" + +# Core Architectural Parameters +base_emb_dim: 3072 +base_num_decoder_layers: 40 +base_num_query_heads: 64 +base_num_kv_heads: 8 +head_dim: 64 +vocab_size: 128008 +normalization_layer_epsilon: 1.0e-6 + +# MoE Specific Parameters +# base_mlp_dim sizes the dense-prefix layer's MLP (see first_num_dense_layers below); +# base_moe_mlp_dim sizes every other (MoE) layer's routed + shared experts. +base_mlp_dim: 10240 +base_moe_mlp_dim: 1536 +num_experts: 256 +shared_experts: 1 +num_experts_per_tok: 8 +norm_topk_prob: true +# Router parity with reference: fp32 gate-logit matmul. +float32_gate_logits: true + +# DeepSeek-V3-style router: sigmoid scoring, aux-loss-free expert-bias balancing +# (primary mechanism) plus a small complementary aux loss (secondary safety net, +# matching DeepSeek-V3's own combined approach), and a top-k weight scaling factor. +routed_score_func: "sigmoid" +routed_bias: true +routed_bias_update_rate: 1.0e-3 +routed_scaling_factor: 2.5 +load_balance_loss_weight: 1.0e-3 + +# Explicit (diverges from the reference's False): keeps the shared-expert gate +# that real Qwen3-Next uses. +moe_shared_expert_gate: true + +# The first layer is a dense MLP (no MoE) and always uses full attention, +# mirroring DeepSeek V3's dense-prefix pattern. +first_num_dense_layers: 1 + +# Qwen3-Next Specific Parameters for Linear Attention (Gated Delta Net). +# Attention schedule (2 GDN : 1 full-attention, full-attention first in each cycle): +# [0,1,1,0,1,1,0,1,1,...] where 0=full attention, 1=Gated Delta Net. +inhomogeneous_layer_cycle_interval: 3 +full_attention_layer_offset: 0 +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_chunk_size: 64 +# L2-norm on Q/K inside the Gated Delta Rule is standard practice for linear-attention +# architectures (stabilizes the recurrent state, since linear attention lacks softmax's +# implicit boundedness) and matches real Qwen3-Next's own Gated Delta Net. This is also +# MaxText's own default (base.yml); kept explicit here for clarity. +use_qk_norm_in_gdn: true + +# RoPE Settings +rope_max_timescale: 10000 +partial_rotary_factor: 1 + +# Hyper-connections: mHC-lite enabled with Pallas kernel +mhc_expansion_rate: 4 +enable_mhc_lite: true +use_mhc_pallas_kernel: true +mhc_pallas_kernel_fwd_block_size: 256 +mhc_pallas_kernel_bwd_block_size: 256 + +# General Model Settings +enable_dropout: false diff --git a/src/maxtext/configs/models/qwen3-next-80b-a3b.yml b/src/maxtext/configs/models/qwen3-next-80b-a3b.yml index 765977f1b5..133ad9c04a 100644 --- a/src/maxtext/configs/models/qwen3-next-80b-a3b.yml +++ b/src/maxtext/configs/models/qwen3-next-80b-a3b.yml @@ -18,35 +18,70 @@ 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: 64 +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 (see first_num_dense_layers below); +# base_moe_mlp_dim sizes every other (MoE) layer's routed + shared experts. +base_mlp_dim: 10240 +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 +# Router parity with reference: fp32 gate-logit matmul. +float32_gate_logits: true -# Qwen3-Next Specific Parameters for Linear Attention (Gated Delta Net) -inhomogeneous_layer_cycle_interval: 4 +# DeepSeek-V3-style router: sigmoid scoring, aux-loss-free expert-bias balancing +# (primary mechanism) plus a small complementary aux loss (secondary safety net, +# matching DeepSeek-V3's own combined approach), and a top-k weight scaling factor. +routed_score_func: "sigmoid" +routed_bias: true +routed_bias_update_rate: 1.0e-3 +routed_scaling_factor: 2.5 +load_balance_loss_weight: 1.0e-3 + +# Explicit (diverges from the reference's False): keeps the shared-expert gate +# that real Qwen3-Next uses. +moe_shared_expert_gate: true + +# The first layer is a dense MLP (no MoE) and always uses full attention, +# mirroring DeepSeek V3's dense-prefix pattern. +first_num_dense_layers: 1 + +# Qwen3-Next Specific Parameters for Linear Attention (Gated Delta Net). +# Attention schedule (2 GDN : 1 full-attention, full-attention first in each cycle): +# [0,1,1,0,1,1,0,1,1,...] where 0=full attention, 1=Gated Delta Net. +inhomogeneous_layer_cycle_interval: 3 +full_attention_layer_offset: 0 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_chunk_size: 64 +# L2-norm on Q/K inside the Gated Delta Rule is standard practice for linear-attention +# architectures (stabilizes the recurrent state, since linear attention lacks softmax's +# implicit boundedness) and matches real Qwen3-Next's own Gated Delta Net. This is also +# MaxText's own default (base.yml); kept explicit here for clarity. +use_qk_norm_in_gdn: true # RoPE Settings -rope_max_timescale: 10000000 -partial_rotary_factor: 0.25 +rope_max_timescale: 10000 +partial_rotary_factor: 1 + +# Hyper-connections: mHC-lite enabled with Pallas kernel +mhc_expansion_rate: 4 +enable_mhc_lite: true +use_mhc_pallas_kernel: true +mhc_pallas_kernel_fwd_block_size: 256 +mhc_pallas_kernel_bwd_block_size: 256 # General Model Settings enable_dropout: false diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 0b135e52d1..dedb386bde 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -276,6 +276,7 @@ class ProfilerType(str, Enum): "qwen3-vl-4b", "qwen3-vl-30b-a3b", "qwen3-next-80b-a3b", + "qwen3-next-80b-a3b-256e", "qwen3-omni-30b-a3b", "qwen3-custom-30b-a3b", "qwen3.5-35b-a3b", @@ -1097,6 +1098,9 @@ class DeepSeekMoE(BaseModel): first_num_dense_layers: NonNegativeInt = Field(0, description="Number of initial dense layers in the model.") shared_experts: NonNegativeInt = Field(0, description="Number of shared experts.") + moe_shared_expert_gate: bool = Field( + False, description="Whether to use a gate on shared experts." + ) routed_scaling_factor: float = Field(1.0, description="Scaling factor for routing scores.") routed_score_func: str = Field("", description="Scoring function for routing (e.g., 'softmax', 'sigmoid').") routed_bias: bool = Field(False, description="Whether to add a bias term for routing.") @@ -1140,6 +1144,19 @@ class Qwen3Next(BaseModel): description="Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.", ) partial_rotary_factor: float = Field(1.0, description="The ratio of dimension to apply ROPE on") + 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." + ), + ) + full_attention_layer_offset: int = Field( + 0, description="Offset for full attention layer." + ) class HardwareAndMesh(BaseModel): @@ -2000,13 +2017,27 @@ class Muon(BaseModel): muon_beta: float = Field(0.95, description="Decay rate for the exponentially weighted average of grads.") muon_weight_decay: float = Field( - 0, - description="Strength of the weight decay regularization. This is multiplied with the learning rate.", + 0.0, + description=( + "Strength of the weight decay regularization. This is multiplied with" + " the learning rate." + ), ) muon_consistent_rms: float | None = Field( None, description="If None, apply width scaling to updates. If float, apply consistent rms scaling (recommend 0.2).", ) + muon_ns_steps: int = Field( + 5, + description="Number of Newton-Schulz iterations for Muon optimizer.", + ) + muon_use_all_to_all: bool = Field( + True, + description=( + "Whether to use all-to-all communication during Newton-Schulz" + " iterations in Muon optimizer." + ), + ) class PositionalEmbedding(BaseModel): @@ -3186,9 +3217,12 @@ def _validate_usp_context_parallelism(self): raise ValueError("TPU USP attention does not support sparse indexer masks.") if self.attention_type != "global": raise ValueError("TPU USP attention is initially supported only for global causal attention.") - if self.context_parallel_load_balance and usp_ring_size % 2 != 0: + if self.packing: + raise ValueError("TPU USP attention does not support packing yet.") + if self.context_parallel_load_balance: raise ValueError( - "TPU USP attention with context_parallel_load_balance=True requires an even ici_context_parallelism." + "TPU USP attention does not support" + " context_parallel_load_balance=True." ) if self.use_ragged_attention: raise ValueError("TPU USP attention does not support ragged attention.") @@ -4067,13 +4101,27 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de if ( self.routed_bias and self.routed_bias_update_rate > 0.0 - and self.decoder_block not in (DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4) + and self.decoder_block + not in ( + DecoderBlockType.DEEPSEEK, + DecoderBlockType.DEEPSEEK4, + DecoderBlockType.QWEN3_NEXT, + ) ): - raise ValueError("Loss-free load balancing is only supported for the DeepSeek decoder block.") - if not self.pure_nnx and self.routed_bias and self.decoder_block == DecoderBlockType.DEEPSEEK4: raise ValueError( - "Auxiliary-loss-free routed bias for DeepSeek V4 is only supported in pure NNX mode. " - "Please set pure_nnx=True or disable routed_bias." + "Loss-free load balancing is only supported for the DeepSeek," + " DeepSeek4, and Qwen3-Next decoder blocks." + ) + if ( + not self.pure_nnx + and self.routed_bias + and self.decoder_block + in (DecoderBlockType.DEEPSEEK4, DecoderBlockType.QWEN3_NEXT) + ): + raise ValueError( + "Auxiliary-loss-free routed bias for DeepSeek V4 and Qwen3-Next is" + " only supported in pure NNX mode. Please set pure_nnx=True or" + " disable routed_bias." ) if self.model_name.startswith("deepseek4") and self.first_num_hash_layers > 0 and self.use_ring_of_experts: raise ValueError("DeepSeek V4 hash routing is currently not supported with ring of experts.") @@ -4363,6 +4411,16 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de rotary_dim = int(self.head_dim * self.partial_rotary_factor) if rotary_dim % 2 != 0: raise ValueError(f"Calculated rotary dimension ({rotary_dim}) must be a multiple of 2.") + gdn_context_parallel_size = self.ici_context_parallelism * self.dcn_context_parallelism + if gdn_context_parallel_size > 1 and self.context_parallel_load_balance: + raise ValueError( + "GatedDeltaNet context parallelism requires context_parallel_load_balance=False. The GatedDeltaNet " + "layers carry a recurrence, so device order is sequence order: device i composes the state left by " + "device i-1. DUAL_CHUNK_SWAP hands device 0 the first and last chunks, device 1 the second and " + "second-to-last, and so on, which composes the segments out of order. Softmax attention tolerates the " + "reorder because it rebuilds the causal mask from positions; a recurrence cannot. The run still trains " + "and the loss still falls, so set this explicitly rather than relying on the failure being visible." + ) else: if self.partial_rotary_factor is not None and self.partial_rotary_factor != 1.0: raise ValueError("`partial_rotary_factor` is only effective when `decoder_block` is set to 'qwen3_next'.") @@ -4440,9 +4498,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.QWEN3_NEXT, DecoderBlockType.GPT_OSS, DecoderBlockType.GEMMA3, DecoderBlockType.LLAMA2, diff --git a/src/maxtext/kernels/attention/gdn_cp.py b/src/maxtext/kernels/attention/gdn_cp.py new file mode 100644 index 0000000000..c9ccf39c44 --- /dev/null +++ b/src/maxtext/kernels/attention/gdn_cp.py @@ -0,0 +1,130 @@ +# 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 +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Context-parallel evaluation of the GatedDeltaNet inter-chunk recurrence. + +The recurrence h_new = A @ h + B is affine in the state, so it composes +associatively and can be split across a sharded sequence: + + A_i = exp(g_last).I - k_g^T.w_i B_i = k_g^T.u_i + +Each device folds its local chunks into one (A, B) pair, the pairs compose +across devices with a prefix scan, and each device then replays its chunks from +the state that reaches it. Composition is associative but not commutative, so +device order is the sequence order. `context_parallel_load_balance` permutes +that order and `configs/types.py` rejects the combination. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +from jax import lax + +_PREC = jax.lax.Precision.HIGHEST + + +def compose(left, right): + """(A_r, B_r) . (A_l, B_l) = (A_r @ A_l, A_r @ B_l + B_r).""" + A_l, B_l = left + A_r, B_r = right + return ( + jnp.matmul(A_r, A_l, precision=_PREC), + jnp.matmul(A_r, B_l, precision=_PREC) + B_r, + ) + + +def compose_local(w, u, k, g): + """Fold this device's chunks into one affine map, in O(1) memory. + + lax.scan rather than associative_scan on purpose: associative_scan would + materialize (A, B) and their running composition for every chunk, which is + 17 GB per device at a million tokens and defeats the point. The parallelism + that matters here is across devices, not within one. + """ + k_dim = k.shape[-1] + eye = jnp.eye(k_dim, dtype=jnp.float32) + + # jax.checkpoint is required here. lax.scan keeps whatever + # the body computes as a backward residual, so A_i and B_i get stacked over + # every chunk even though the forward pass only ever needs one at a time. At + # sequence 262,144 with ctx=4 that was 103 GB of f32[1024,4,16,128,128] in the + # buffer dump. A_i and B_i are cheap to rebuild from w, u, k and g, which are + # already live, so recompute them in the backward pass instead of storing them. + @jax.checkpoint + def body(carry, x): + w_c, u_c, k_c, g_c = x + g_last = g_c[..., -1] + decay = jnp.exp(g_last)[..., None, None] + k_g = k_c.astype(jnp.float32) * jnp.exp(g_last[..., None] - g_c)[..., None] + k_g_T = k_g.swapaxes(-1, -2) + A_i = decay * eye - jnp.matmul(k_g_T, w_c.astype(jnp.float32), precision=_PREC) + B_i = jnp.matmul(k_g_T, u_c.astype(jnp.float32), precision=_PREC) + return compose(carry, (A_i, B_i)), None + + lead = w.shape[1:-2] + init = ( + jnp.broadcast_to(eye, lead + (k_dim, k_dim)).astype(jnp.float32), + jnp.zeros(lead + (k_dim, u.shape[-1]), jnp.float32), + ) + (A_loc, B_loc), _ = lax.scan(body, init, (w, u, k, g)) + return A_loc, B_loc + + +def incoming_state(A_loc, B_loc, h_init, cp_axis): + """State entering this device, plus the final state after all devices. + + log2(D) exchanges of one small matrix pair are the only cross-device traffic + in the scheme. Must be called inside a shard_map over `cp_axis`. + """ + # A prefix scan over the device axis, not an all-gather. Gathering the D pairs + # materializes [D, B, H, K, K] and [D, B, H, K, V] on every device, so the + # composition costs O(D) per device per layer. At ctx=256 on Qwen3.5-27B that + # is roughly 38 GB across the 48 GatedDeltaNet layers, and it is what makes + # per-device memory grow with the context axis instead of shrinking. + # + # Hillis-Steele instead: log2(D) ppermute exchanges, one (A, B) pair live at a + # time, O(1) in D, and log2(D) backward residuals rather than D. + D = lax.axis_size(cp_axis) + idx = lax.axis_index(cp_axis) + + # Inclusive prefix: each device ends holding the composition of ranks 0..idx. + # compose(left, right) applies left first, so the pair arriving from the + # earlier rank is the left operand. + a_run, b_run = A_loc, B_loc + step = 1 + while step < D: + fwd = [(i, i + step) for i in range(D - step)] + a_recv = lax.ppermute(a_run, cp_axis, fwd) + b_recv = lax.ppermute(b_run, cp_axis, fwd) + a_cmp, b_cmp = compose((a_recv, b_recv), (a_run, b_run)) + live = idx >= step + a_run = jnp.where(live, a_cmp, a_run) + b_run = jnp.where(live, b_cmp, b_run) + step *= 2 + + # A device starts from the exclusive prefix, which is the inclusive prefix of + # the device before it. One more shift by a single rank. + shift1 = [(i, i + 1) for i in range(D - 1)] + a_ex = lax.ppermute(a_run, cp_axis, shift1) + b_ex = lax.ppermute(b_run, cp_axis, shift1) + carried = jnp.matmul(a_ex, h_init, precision=_PREC) + b_ex + h_in = jnp.where(idx == 0, h_init, carried) + + # The final state is the last device's inclusive prefix, broadcast with a + # masked psum rather than a gather so this stays O(1) in D too. + last = idx == (D - 1) + a_tot = lax.psum(jnp.where(last, a_run, jnp.zeros_like(a_run)), cp_axis) + b_tot = lax.psum(jnp.where(last, b_run, jnp.zeros_like(b_run)), cp_axis) + final_h = jnp.matmul(a_tot, h_init, precision=_PREC) + b_tot + return h_in, final_h diff --git a/src/maxtext/kernels/megablox/backend.py b/src/maxtext/kernels/megablox/backend.py index 618965c840..bef18ea04a 100644 --- a/src/maxtext/kernels/megablox/backend.py +++ b/src/maxtext/kernels/megablox/backend.py @@ -30,6 +30,16 @@ import qwix.pallas as qpl +def _make_shape_dtype_struct(shape, dtype, varying_axes=()): + try: + manual_axis_type = jax.sharding.ManualAxisType( + varying=frozenset(varying_axes) + ) + return jax.ShapeDtypeStruct(shape, dtype, manual_axis_type=manual_axis_type) + except (TypeError, AttributeError): + return jax.ShapeDtypeStruct(shape, dtype) + + def _validate_args( *, lhs: jnp.ndarray, @@ -524,8 +534,8 @@ def out_transform_indices(n_i, grid_id, k_i, group_metadata, group_offset): } call_gmm = qpl.pallas_call( kernel, - out_shape=jax.ShapeDtypeStruct( - (m, n), preferred_element_type, manual_axis_type=jax.sharding.ManualAxisType(varying=frozenset(varying_axes)) + out_shape=_make_shape_dtype_struct( + (m, n), preferred_element_type, varying_axes ), grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=2, @@ -539,7 +549,9 @@ def out_transform_indices(n_i, grid_id, k_i, group_metadata, group_offset): scratch_shapes=[pltpu.VMEM((tm, tn), jnp.float32)], ), input_output_aliases=input_output_aliases, - compiler_params=pltpu.CompilerParams(dimension_semantics=("parallel", "arbitrary", "arbitrary")), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary") + ), interpret=interpret, cost_estimate=cost_estimate, metadata={"xprof_metadata": json.dumps(metadata)}, @@ -783,10 +795,8 @@ def out_transform_indices(n_i, k_i, grid_id, group_metadata, group_offset): } call_gmm = qpl.pallas_call( kernel, - out_shape=jax.ShapeDtypeStruct( - (num_actual_groups, k, n), - preferred_element_type, - manual_axis_type=jax.sharding.ManualAxisType(varying=frozenset(varying_axes)), + out_shape=_make_shape_dtype_struct( + (num_actual_groups, k, n), preferred_element_type, varying_axes ), grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=2, @@ -800,7 +810,9 @@ def out_transform_indices(n_i, k_i, grid_id, group_metadata, group_offset): scratch_shapes=[pltpu.VMEM((tk, tn), jnp.float32)], ), input_output_aliases=input_output_aliases, - compiler_params=pltpu.CompilerParams(dimension_semantics=("parallel", "arbitrary", "arbitrary")), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary") + ), interpret=interpret, cost_estimate=cost_estimate, metadata={"xprof_metadata": json.dumps(metadata)}, diff --git a/src/maxtext/kernels/mhc/api.py b/src/maxtext/kernels/mhc/api.py index f1f0b77c4c..624613459e 100644 --- a/src/maxtext/kernels/mhc/api.py +++ b/src/maxtext/kernels/mhc/api.py @@ -60,8 +60,8 @@ def pre( Args: x: Input streams of shape `(batch, sequence, streams, embedding)`. weights: Structured `MhcWeights` container with all layer parameters. - permutations: All permutation matrices of shape - `(num_permutations, streams, streams)`. + permutations: All permutation matrices of shape `(num_permutations, streams, + streams)`. config: Structured `MhcKernelConfig` tuning and compiler configuration. implementation: Preferred implementation (`"mosaic"` or `"mosaic_tpu"`). @@ -95,8 +95,8 @@ def post( """Runs the post-gate and residual stream mixing. Args: - layer_output: Output from the wrapped branch of shape - `(batch, sequence, embedding)`. + layer_output: Output from the wrapped branch of shape `(batch, sequence, + embedding)`. context: Opaque `MhcContext` returned by `pre`. config: Structured `MhcKernelConfig` tuning and compiler configuration. @@ -104,7 +104,9 @@ def post( Mixed output streams of shape `(batch, sequence, streams, embedding)`. """ if context.implementation not in ("mosaic", "mosaic_tpu"): - raise ValueError(f"Unsupported implementation in MhcContext: '{context.implementation}'") + raise ValueError( + f"Unsupported implementation in MhcContext: '{context.implementation}'" + ) kernel_context = (context.x, context.h_post, context.residual) return mhc_kernels_fwd.post( layer_output, diff --git a/src/maxtext/kernels/mhc/mhc_kernels_bwd.py b/src/maxtext/kernels/mhc/mhc_kernels_bwd.py index e92eda3351..dd6ed7ad9d 100644 --- a/src/maxtext/kernels/mhc/mhc_kernels_bwd.py +++ b/src/maxtext/kernels/mhc/mhc_kernels_bwd.py @@ -408,14 +408,16 @@ def post_bwd( def pre_op_bwd( config: common.MhcKernelConfig, - permutations: jax.Array, - residuals: tuple[tuple[jax.Array, jax.Array], tuple[jax.Array, common.MhcWeights]], + residuals: tuple[ + tuple[jax.Array, jax.Array], + tuple[jax.Array, jax.Array, common.MhcWeights], + ], cotangents: tuple[jax.Array, common.KernelContext], -) -> tuple[jax.Array, common.MhcWeights]: +) -> tuple[jax.Array, None, common.MhcWeights]: """Custom-VJP backward rule for the low-level pre-branch entry point.""" - saved, (x, weights) = residuals + saved, (x, permutations, weights) = residuals d_layer_input, (d_x, d_h_post, d_residual) = cotangents - return pre_bwd( + d_x_out, d_weights = pre_bwd( saved, (d_layer_input, d_x, d_h_post, d_residual), x, @@ -423,6 +425,7 @@ def pre_op_bwd( permutations, config=config, ) + return d_x_out, None, d_weights def post_op_bwd( diff --git a/src/maxtext/kernels/mhc/mhc_kernels_fwd.py b/src/maxtext/kernels/mhc/mhc_kernels_fwd.py index cd92e3750a..497e72c5f1 100644 --- a/src/maxtext/kernels/mhc/mhc_kernels_fwd.py +++ b/src/maxtext/kernels/mhc/mhc_kernels_fwd.py @@ -238,7 +238,7 @@ def post_fwd( return output.reshape(batch, sequence, streams, embedding) -@functools.partial(jax.custom_vjp, nondiff_argnums=(0, 2)) +@functools.partial(jax.custom_vjp, nondiff_argnums=(0,)) def _pre_op( config: common.MhcKernelConfig, x: jax.Array, @@ -258,7 +258,7 @@ def _pre_op_fwd( ): """Custom-VJP forward rule for the pre-branch operation.""" primals_out, saved = pre_fwd(x, weights, permutations, config) - return primals_out, (saved, (x, weights)) + return primals_out, (saved, (x, permutations, weights)) _pre_op.defvjp(_pre_op_fwd, mhc_kernels_bwd.pre_op_bwd) diff --git a/src/maxtext/kernels/ragged/ragged_gather.py b/src/maxtext/kernels/ragged/ragged_gather.py index 0c2aabb1c1..bf2a7c78fb 100644 --- a/src/maxtext/kernels/ragged/ragged_gather.py +++ b/src/maxtext/kernels/ragged/ragged_gather.py @@ -410,7 +410,7 @@ def ragged_gather( # Guard against eager initialization on non-TPU hardware (e.g. during CPU tests). # pltpu.get_tpu_info() expects TPU hardware and will crash if executed on CPU. - if enforce_fallback or jax.devices()[0].platform != "tpu": + if enforce_fallback: return _fallback_implementation(x, indices, weights, has_weights) sc_info = pltpu.get_tpu_info().sparse_core diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py index 0594676b74..a1cfb09d91 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py @@ -660,7 +660,7 @@ def ragged_gather_reduce( # Step 1: Choose the implementation (TensorCore fallback or SparseCore). # Guard against eager initialization on non-TPU hardware (e.g. during CPU tests). # pltpu.get_tpu_info() expects TPU hardware and will crash if executed on CPU. - if enforce_fallback or jax.devices()[0].platform != "tpu": + if enforce_fallback: return _fallback_implementation(x, indices, topk_weights, valid_rows_mask, reduce_group_size) sc_info = pltpu.get_tpu_info().sparse_core diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 84972e5d71..88e2e98e69 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -77,7 +77,11 @@ from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_pspec, get_logical_axis_rules import numpy as np from tokamax._src.ops.attention import base as tokamax_attention_base -from tokamax._src.ops.attention import pallas_triton as tokamax_pallas_triton + +try: + from tokamax._src.ops.attention import pallas_triton as tokamax_pallas_triton +except ImportError: + tokamax_pallas_triton = None from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_kernel as tokamax_splash_kernel from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_mask as tokamax_splash_mask # pylint: disable=line-too-long, g-doc-args, g-doc-return-or-yield, bad-continuation, g-inconsistent-quotes diff --git a/src/maxtext/layers/mhc.py b/src/maxtext/layers/mhc.py index a4ac46ecba..6bbed8f117 100644 --- a/src/maxtext/layers/mhc.py +++ b/src/maxtext/layers/mhc.py @@ -22,24 +22,28 @@ from flax import nnx import jax import jax.numpy as jnp -from jax.sharding import Mesh +from jax.sharding import Mesh, PartitionSpec as P from maxtext.common.common_types import Array, Config from maxtext.common.common_types import HyperConnectionType from maxtext.kernels.mhc import api as mhc_kernel from maxtext.layers import nnx_wrappers from maxtext.layers.initializers import default_bias_init, default_scalar_init, nd_dense_init, variable_to_logically_partitioned from maxtext.layers.normalizations import RMSNorm +from maxtext.utils.sharding import get_logical_axis_rules, logical_to_mesh_axes @functools.lru_cache(maxsize=None) +def _get_permutation_indices(k: int) -> tuple[tuple[int, ...], ...]: + return tuple(itertools.permutations(range(k))) + + def get_permutation_matrices(k: int) -> Array: """Generates all permutation matrices of size k. Reference: mHC-lite: https://openreview.net/pdf?id=5IJX6kvOif Shape: (k!, k, k) """ - perms = list(itertools.permutations(range(k))) - perms_array = jnp.array(perms) + perms_array = jnp.array(_get_permutation_indices(k)) return jnp.eye(k)[perms_array] @@ -106,8 +110,13 @@ def __init__( self.weight_dtype = self.config.weight_dtype self.matmul_precision = jax.lax.Precision(self.config.matmul_precision) - if getattr(self.config, "use_mhc_pallas_kernel", False) and not self.config.enable_mhc_lite: - raise ValueError("use_mhc_pallas_kernel=True requires enable_mhc_lite=True.") + if ( + getattr(self.config, "use_mhc_pallas_kernel", False) + and not self.config.enable_mhc_lite + ): + raise ValueError( + "use_mhc_pallas_kernel=True requires enable_mhc_lite=True." + ) # Norm layer self.mhc_norm = RMSNorm( @@ -269,22 +278,78 @@ def __call__( h_post = None h_res = None context = None - use_kernel = self.config.enable_mhc_lite and getattr(self.config, "use_mhc_pallas_kernel", False) + use_kernel = self.config.enable_mhc_lite and getattr( + self.config, "use_mhc_pallas_kernel", False + ) if use_kernel: - fwd_block_size = getattr(self.config, "mhc_pallas_kernel_fwd_block_size", 256) - bwd_block_size = getattr(self.config, "mhc_pallas_kernel_bwd_block_size", 128) + fwd_block_size = getattr( + self.config, "mhc_pallas_kernel_fwd_block_size", 256 + ) + bwd_block_size = getattr( + self.config, "mhc_pallas_kernel_bwd_block_size", 128 + ) + bwd_feature_block_size = getattr( + self.config, "mhc_pallas_kernel_bwd_feature_block_size", 1024 + ) kernel_config = mhc_kernel.MhcKernelConfig( block_size=fwd_block_size, bwd_block_size=bwd_block_size, + bwd_feature_block_size=bwd_feature_block_size, rms_epsilon=self.config.normalization_layer_epsilon, ) weights = self._get_mhc_weights() - layer_input, context = mhc_kernel.pre( - x, - weights, - jnp.asarray(get_permutation_matrices(self.k), self.dtype), - config=kernel_config, - ) + perms = jnp.asarray(get_permutation_matrices(self.k), self.dtype) + if self.mesh is not None: + rules = get_logical_axis_rules() + x_pspec = logical_to_mesh_axes( + ("activation_batch", "activation_length", None, "activation_embed"), + mesh=self.mesh, + rules=rules, + ) + layer_input_pspec = logical_to_mesh_axes( + ("activation_batch", "activation_length", "activation_embed"), + mesh=self.mesh, + rules=rules, + ) + h_post_pspec = logical_to_mesh_axes( + ("activation_batch", "activation_length", None), + mesh=self.mesh, + rules=rules, + ) + residual_pspec = logical_to_mesh_axes( + ("activation_batch", "activation_length", None, None), + mesh=self.mesh, + rules=rules, + ) + weights_pspec = jax.tree.map(lambda _: P(), weights) + perms_pspec = P() + context_pspec = mhc_kernel.MhcContext( + x=x_pspec, + h_post=h_post_pspec, + residual=residual_pspec, + implementation="mosaic", + ) + + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=(x_pspec, weights_pspec, perms_pspec), + out_specs=(layer_input_pspec, context_pspec), + check_vma=False, + ) + def shard_mapped_pre(x_val, weights_val, perms_val): + return mhc_kernel.pre( + x_val, weights_val, perms_val, config=kernel_config + ) + + layer_input, context = shard_mapped_pre(x, weights, perms) + else: + layer_input, context = mhc_kernel.pre( + x, + weights, + perms, + config=kernel_config, + ) else: with jax.named_scope("mhc_norm"): # 1. Flatten the tensor, and RMS normalization @@ -295,10 +360,14 @@ def __call__( post_alpha = jnp.asarray(self.post_alpha[...], self.dtype) res_alpha = jnp.asarray(self.res_alpha[...], self.dtype) - alpha_concat = jnp.concatenate([pre_alpha, post_alpha, res_alpha], axis=-1) + alpha_concat = jnp.concatenate( + [pre_alpha, post_alpha, res_alpha], axis=-1 + ) # MatMul on normalized input - h_concat = jnp.einsum("bsm,mn -> bsn", norm_x, alpha_concat, precision=self.matmul_precision) + h_concat = jnp.einsum( + "bsm,mn -> bsn", norm_x, alpha_concat, precision=self.matmul_precision + ) h_pre = h_concat[..., : self.k] h_post = h_concat[..., self.k : 2 * self.k] h_res = h_concat[..., 2 * self.k :] @@ -311,13 +380,9 @@ def __call__( 1.0, eps=1e-6, ) - # bskd, bsk -> bsd (fused contracted GEMM) - layer_input = jnp.einsum( - "bsk,bskd->bsd", - pre_mapping, - x, - precision=self.matmul_precision, - ) + # Moving away from einsum seems to allow XLA to perform better fusions + # bskd, bsk -> bsd + layer_input = jnp.sum(x * jnp.expand_dims(pre_mapping, axis=3), axis=2) # 3. Pre-norm layer_input = norm_fn(layer_input) @@ -336,17 +401,68 @@ def __call__( raise ValueError(f"Unsupported type: {mhc_type}") if use_kernel: - fwd_block_size = getattr(self.config, "mhc_pallas_kernel_fwd_block_size", 256) - bwd_block_size = getattr(self.config, "mhc_pallas_kernel_bwd_block_size", 128) + fwd_block_size = getattr( + self.config, "mhc_pallas_kernel_fwd_block_size", 256 + ) + bwd_block_size = getattr( + self.config, "mhc_pallas_kernel_bwd_block_size", 128 + ) + bwd_feature_block_size = getattr( + self.config, "mhc_pallas_kernel_bwd_feature_block_size", 1024 + ) kernel_config = mhc_kernel.MhcKernelConfig( block_size=fwd_block_size, bwd_block_size=bwd_block_size, + bwd_feature_block_size=bwd_feature_block_size, ) - output = mhc_kernel.post( - layer_out, - context, - config=kernel_config, - ) + if self.mesh is not None: + rules = get_logical_axis_rules() + x_pspec = logical_to_mesh_axes( + ("activation_batch", "activation_length", None, "activation_embed"), + mesh=self.mesh, + rules=rules, + ) + layer_out_pspec = logical_to_mesh_axes( + ("activation_batch", "activation_length", "activation_embed"), + mesh=self.mesh, + rules=rules, + ) + h_post_pspec = logical_to_mesh_axes( + ("activation_batch", "activation_length", None), + mesh=self.mesh, + rules=rules, + ) + residual_pspec = logical_to_mesh_axes( + ("activation_batch", "activation_length", None, None), + mesh=self.mesh, + rules=rules, + ) + context_pspec = mhc_kernel.MhcContext( + x=x_pspec, + h_post=h_post_pspec, + residual=residual_pspec, + implementation="mosaic", + ) + + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=(layer_out_pspec, context_pspec), + out_specs=x_pspec, + check_vma=False, + ) + def shard_mapped_post(layer_out_val, context_val): + return mhc_kernel.post( + layer_out_val, context_val, config=kernel_config + ) + + output = shard_mapped_post(layer_out, context) + else: + output = mhc_kernel.post( + layer_out, + context, + config=kernel_config, + ) return output, metadata # 5. Post mapping diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 82d4b01e47..892e93f473 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -55,7 +55,11 @@ ) import numpy as np import qwix -from qwix.contrib.sparsity import sparsity_module + +try: + from qwix.contrib.sparsity import sparsity_module +except ImportError: + sparsity_module = None import qwix.pallas as qpl import tokamax @@ -401,14 +405,14 @@ def __call__(self, inputs: jax.Array, _initializing: bool = False) -> Tuple[jax. _initializing, out_sharding=output_sharding, ) - pre_bias_logits = None - if self.score_func: output = linears._convert_to_activation_function(self.score_func)(output) - # NOTE: deepseek2 has a different pattern - if self.model_name.startswith(("deepseek3", "deepseek4", "kimi-k2")): - pre_bias_logits = output + # Snapshot pre-bias logits unconditionally (cheap, already-computed array). Only + # consumed by callers that need it (DeepSeek V3/V4, Kimi-K2, or Qwen3-Next opting into + # bias/grouped routing via RoutedMoE.uses_grouped_or_bias_routing()); harmless + # no-op for every other decoder block. + pre_bias_logits = output if self.use_bias: bias = jnp.asarray(self.bias[...], self.dtype) @@ -748,6 +752,23 @@ def should_update_load_balance(self): """ return self.config.routed_bias and self.config.routed_bias_update_rate > 0.0 and not self.is_hash_routing + def uses_grouped_or_bias_routing(self) -> bool: + """Whether expert selection uses post-bias/grouped logits, weighted by pre-bias logits. + + True for DeepSeek V3/V4 and Kimi-K2 (kept as a model_name check, not + decoder_block, + since DeepSeek V2 + and V3 share DecoderBlockType.DEEPSEEK), and additively for Qwen3-Next once + it opts in via + routed_bias (aux-loss-free bias) or n_routing_groups (expert grouping). + """ + return self.config.model_name.startswith( + ("deepseek3", "deepseek4", "kimi-k2") + ) or ( + self.config.decoder_block == ctypes.DecoderBlockType.QWEN3_NEXT + and (self.config.routed_bias or self.config.n_routing_groups != -1) + ) + def get_topk( self, gate_logits, @@ -791,8 +812,7 @@ def get_topk( # Cast input_ids to int32 to safely index the hash routing table top_k_indices = tid2eid_int[input_ids.astype(jnp.int32)] top_k_weights = jnp.take_along_axis(pre_bias_logits, top_k_indices, axis=-1) - # NOTE: deepseek2 has a different pattern - elif self.config.model_name.startswith(("deepseek3", "deepseek4", "kimi-k2")): + elif self.uses_grouped_or_bias_routing(): top_k_weights, top_k_indices = self.deepseek_routing(gate_logits, pre_bias_logits) elif self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4: router_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1) @@ -801,7 +821,22 @@ def get_topk( else: top_k_weights, top_k_indices = jax.lax.top_k(gate_logits, self.num_experts_per_tok) - if self.config.decoder_block in (ctypes.DecoderBlockType.DEEPSEEK, ctypes.DecoderBlockType.DEEPSEEK4): + # Qwen3-Next takes the DeepSeek-style scaling path once it opts into a non-default + # routed_score_func / routed_scaling_factor / grouped-or-bias routing; a fully-default + # qwen3_next config falls through to the plain softmax branch below, identical to + # before this option existed. + qwen3_next_needs_deepseek_style_scaling = ( + self.config.decoder_block == ctypes.DecoderBlockType.QWEN3_NEXT + and ( + bool(self.config.routed_score_func) + or self.config.routed_scaling_factor != 1.0 + or self.uses_grouped_or_bias_routing() + ) + ) + if ( + self.config.decoder_block == ctypes.DecoderBlockType.DEEPSEEK + or qwen3_next_needs_deepseek_style_scaling + ): top_k_weights = self.deepseek_scale_weights(top_k_weights) if valid_token_mask is not None: top_k_weights = top_k_weights * valid_token_mask @@ -1563,9 +1598,11 @@ def get_tokamax_group_sizes(group_sizes, inputs, _kernel): elif self.config.attention in ("vllm_rpa", "vllm_batched_rpa"): return group_sizes else: + num_groups = group_sizes.shape[0] + avg_size = inputs.shape[0] // num_groups return tokamax.RaggedDotGroupSizes( group_sizes, - inputs.shape[0], + (avg_size,) * num_groups, ) def get_quantization_dtypes(): @@ -1726,8 +1763,7 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): wo_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_embed")) gate_logits_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) - # NOTE: deepseek2 has a different pattern - if self.config.model_name.startswith(("deepseek3", "deepseek4", "kimi-k2")): + if self.uses_grouped_or_bias_routing(): pre_bias_logits_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) else: # pre_bias_logits is None for non-deepseek3/4 models, including deepseek2 @@ -2599,7 +2635,7 @@ def sparse_matmul_route_and_compute( gate_logits_logical_axes = (batch_logical_axis, "activation_norm_length", None) pre_bias_logits_logical_axes = ( (batch_logical_axis, "activation_norm_length", None) - if self.config.model_name.startswith(("deepseek3", "deepseek4", "kimi-k2")) + if self.uses_grouped_or_bias_routing() else None ) inputs = self._maybe_shard_with_pspec(inputs, input_partition_pspec, logical_axes=input_logical_axes) @@ -2956,9 +2992,7 @@ def dense_matmul( """Dense matrix multiplication.""" # gate_logits: batch, length, expert gate_logits = self._maybe_shard_with_logical(gate_logits, ("activation_batch_moe", "activation_length_moe", None)) - # NOTE: deepseek2 has a different pattern - if self.config.model_name.startswith(("deepseek3", "deepseek4", "kimi-k2")): - # pre_bias_logits is None for non-deepseek3/4 models, including deepseek2 + if self.uses_grouped_or_bias_routing(): pre_bias_logits = self._maybe_shard_with_logical( pre_bias_logits, ("activation_batch_moe", "activation_length_moe", None) ) @@ -3290,6 +3324,7 @@ def fused_moe_matmul( # pylint: disable=import-outside-toplevel # pytype: disable=import-error from tpu_inference.layers.common.fused_moe_gmm import fused_moe_func + from tpu_inference import envs as tpu_inference_envs except ImportError as e: raise ImportError("fused_moe_matmul requires the tpu-inference package.") from e @@ -3333,6 +3368,18 @@ def fused_moe_matmul( use_ep=use_ep, activation=activation, scoring_fn=scoring_fn, + # Forward the same environment-backed kernel knobs that tpu-inference passes on its + # own serving path (tpu_inference/layers/common/moe.py). Without these, env vars such + # as ONEHOT_MOE_PERMUTE_THRESHOLD and VLLM_MOE_CHUNK_SIZE are silently ignored when a + # MaxText model is served through vLLM, so the two paths run the same kernel with + # different configurations. In particular, ONEHOT_MOE_PERMUTE_THRESHOLD selects the + # one-hot permute path instead of the SparseCore ragged_gather_reduce kernel, which + # is what makes expert parallelism usable on TPU generations whose SparseCore has + # fewer SIMD lanes than the kernel requires (e.g. v5p). + enable_rs_kernel=tpu_inference_envs.ENABLE_RS_KERNEL, + use_gmm_fused_rs_kernel=tpu_inference_envs.USE_GMM_FUSED_RS_KERNEL, + onehot_moe_permute_threshold=tpu_inference_envs.ONEHOT_MOE_PERMUTE_THRESHOLD, + moe_chunk_size=tpu_inference_envs.VLLM_MOE_CHUNK_SIZE, ) # Reshape output 2D [T, D] -> 3D [B, S, D] diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index e0347b8f1e..ddb518f44b 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -480,6 +480,15 @@ 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_with_dense = ( + self.config.decoder_block == DecoderBlockType.QWEN3_NEXT + and self.config.first_num_dense_layers > 0 + and self.config.scan_layers + ) + self.is_qwen3_next = ( + self.config.decoder_block == DecoderBlockType.QWEN3_NEXT + and not self.is_qwen3_next_with_dense + ) if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4: self.hc_head = mhc.DeepSeek4HyperHead( @@ -590,6 +599,10 @@ 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(rngs, mesh) + elif self.is_qwen3_next_with_dense: + self._init_scanned_qwen3_next_with_dense(rngs) else: self._init_scanned_generic(decoder_block_classes, rngs) @@ -692,6 +705,73 @@ def _init_scanned_deepseek_standard(self, dense_cls, moe_cls, rngs): num_moe = config.num_decoder_layers - config.first_num_dense_layers self.moe_layers = self._create_scanned_layers(moe_cls, length=num_moe, metadata_axis_name="moe_layers", rngs=rngs) + def _init_scanned_qwen3_next_with_dense(self, rngs): + """Initializes scanned Qwen3-Next layers with a `first_num_dense_layers` dense prefix. + + Splits the stack into three pieces: an unscanned dense prefix (real global + `layer_idx`, forced full attention), a scanned middle of uniform + `Qwen3NextScannableBlock` repeats, and β€” since `num_decoder_layers - + first_num_dense_layers` need not be a multiple of + `inhomogeneous_layer_cycle_interval` + β€” an unscanned remainder tail, mirroring `_init_scanned_gemma3`'s + scan+remainder + split. `layer_idx_offset` keeps the attention-type cycle globally correct + across + all three pieces despite the dense prefix breaking its natural period. + """ + config = self.config + n_dense = config.first_num_dense_layers + cycle = config.inhomogeneous_layer_cycle_interval + remaining = config.num_decoder_layers - n_dense + scan_length = remaining // cycle + num_remaining = remaining % cycle + + self.dense_layers = self._create_scanned_layers( + qwen3.Qwen3NextDecoderLayer, + length=n_dense, + metadata_axis_name="dense_layers", + rngs=rngs, + layer_idx=0, + is_dense_layer=True, + ) + + policy = self.get_remat_policy() + layer_kwargs = { + "num_of_layers": cycle, + "layer_idx_offset": n_dense, + "remat_policy_fn": policy, + "apply_internal_remat": True, + } + rem_layer_kwargs = { + "num_of_layers": num_remaining, + "layer_idx_offset": n_dense + scan_length * cycle, + "remat_policy_fn": policy, + "apply_internal_remat": True, + } + + if scan_length > 0: + self.layers = self._create_scanned_layers( + qwen3.Qwen3NextScannableBlock, + length=scan_length, + metadata_axis_name="layers", + rngs=rngs, + **layer_kwargs, + ) + else: + self.layers = nnx.List([]) + + if num_remaining > 0: + self.layers_remainder = qwen3.Qwen3NextScannableBlock( + config=config, + mesh=self.mesh, + model_mode=self.model_mode, + quant=self.quant, + rngs=rngs, + **rem_layer_kwargs, + ) + else: + self.layers_remainder = None + def _init_scanned_gemma3(self, decoder_block_classes, rngs, mesh): """Initializes scanned Gemma3 layers.""" config = self.config @@ -1389,14 +1469,24 @@ def _apply_embedding( deterministic, model_mode, multimodal_input=None, + decoder_input_embeddings=None, ): """Applies token and positional embeddings to the input tokens.""" + cfg = self.config - y = shared_embedding(decoder_input_tokens.astype("int32"), model_mode=model_mode) + # vLLM passes token IDs for text-only requests, but passes complete, premerged + # text and multimodal embeddings for multimodal requests. The latter enter + # through `decoder_input_embeddings` to avoid embedding and merging them again. + y = ( + decoder_input_embeddings + if decoder_input_embeddings is not None + else shared_embedding(decoder_input_tokens.astype("int32"), model_mode=model_mode) + ) - # Merge the image embeddings with the text embeddings for multimodal models - if multimodal_input is not None: + # Precomputed embeddings are complete (including any multimodal replacements), + # so only merge modality embeddings when token embeddings were created here. + if decoder_input_embeddings is None and multimodal_input is not None: image_embeddings = multimodal_input.image_embeddings bidirectional_mask = multimodal_input.bidirectional_mask image_masks = multimodal_input.image_masks @@ -1664,6 +1754,7 @@ def __call__( deepstack_visual_embeds: None | list[jnp.ndarray] = None, multimodal_input: None | MultimodalInput = None, forced_routed_experts: jnp.ndarray | None = None, + decoder_input_embeddings=None, ): cfg = self.config assert decoder_input_tokens.ndim == 2 # [batch, len] @@ -1705,6 +1796,7 @@ def __call__( deterministic, model_mode, multimodal_input=multimodal_input, + decoder_input_embeddings=decoder_input_embeddings, ) mhc_reduce = None @@ -1962,6 +2054,19 @@ 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, + ) + elif self.is_qwen3_next_with_dense: + y = self._apply_qwen3_next_dense_scanned_blocks( + y, + layer_args, + layer_kwargs, + ) else: cycle_interval = cfg.inhomogeneous_layer_cycle_interval scan_length = int(cfg.num_decoder_layers / cycle_interval) @@ -2418,6 +2523,51 @@ def pure_gemma_fn(graphdef, state_in, y_in, kv_in): return y + def _apply_qwen3_next_dense_scanned_blocks( + self, + y, + layer_args, + layer_kwargs, + ): + """Applies Qwen3-Next's dense prefix, scanned middle, and remainder tail (see + + `_init_scanned_qwen3_next_with_dense`), mirroring + `_apply_gemma3_scanned_blocks`. + """ + cfg = self.config + n_dense = cfg.first_num_dense_layers + if n_dense > 0: + y, self.dense_layers, _ = self._apply_layers_sequentially( + self.dense_layers, + y, + *layer_args, + length=n_dense, + metadata_axis_name="dense_layers", + ) + + remaining = cfg.num_decoder_layers - n_dense + cycle = cfg.inhomogeneous_layer_cycle_interval + scan_length = remaining // cycle + if scan_length > 0: + y, self.layers, _ = self._apply_layers_sequentially( + self.layers, + y, + *layer_args, + length=scan_length, + **layer_kwargs, + ) + + num_remaining = remaining % cycle + if num_remaining > 0: + out = self.layers_remainder( + y, + *layer_args, + **layer_kwargs, + ) + y = out[0] if isinstance(out, tuple) else out + + return y + def _apply_gemma4_small_layers( self, y, @@ -2558,6 +2708,7 @@ def _append_scanned(name): if hasattr(self, "pipeline_module"): _add(getattr(self.pipeline_module, "layers", None)) + _append_scanned("dense_layers") _append_scanned("scanned_blocks") # Gemma 4 _append_scanned("layers") _append_unscanned("layers") diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index 0ef3aeca55..5cf82d5d4e 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -30,7 +30,11 @@ import qwix from qwix._src.core import numerics from qwix._src.core import dot_general_qt -from qwix._src.core import sparsity + +try: + from qwix._src.core import sparsity +except ImportError: + sparsity = None import jax import jax.numpy as jnp diff --git a/src/maxtext/models/hybrid_gdn.py b/src/maxtext/models/hybrid_gdn.py new file mode 100644 index 0000000000..42a10ae535 --- /dev/null +++ b/src/maxtext/models/hybrid_gdn.py @@ -0,0 +1,1042 @@ +# 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 + Pallas Custom VJP backward.""" + +import functools +from typing import Any, Optional, Tuple + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp +from maxtext.layers.normalizations import l2norm +try: + from tokamax._src.ops.experimental.causal_conv1d_gated_delta_rule import ( + wrapper as tokamax_gdn_wrapper, + ) +except (ImportError, ModuleNotFoundError): + try: + from tokamax._src.ops.causal_conv1d_gated_delta_rule import ( + wrapper as tokamax_gdn_wrapper, + ) + except (ImportError, ModuleNotFoundError): + tokamax_gdn_wrapper = None + + +def _pallas_gdn_bwd_kernel( + padded_pre_conv_qkv_ref, + qkv_ref, + b_ref, + a_ref, + a_log_ref, + dt_bias_ref, + do_ref, + chunk_states_ref, + conv_weight_ref, + seq_lens_ref, + d_qkv_ref, + d_b_ref, + d_a_ref, + d_conv_weight_ref, + d_a_log_ref, + d_dt_bias_ref, + padded_pre_conv_qkv_vmem, + qkv_vmem, + b_vmem, + a_vmem, + a_log_vmem, + dt_bias_vmem, + do_vmem, + chunk_states_vmem, + d_qkv_vmem, + d_b_vmem, + d_a_vmem, + d_conv_weight_scratch, + d_a_log_scratch, + d_dt_bias_scratch, + sem_pre_conv_qkv, + sem_qkv, + sem_b, + sem_a, + sem_a_log, + sem_dt_bias, + sem_do, + sem_chunk_states, + sem_d_qkv, + sem_d_b, + sem_d_a, + sem_d_conv_weight, + sem_d_a_log, + sem_d_dt_bias, + *, + batch_size: int, + num_chunks: int, + chunk_size: int, + dim_size: int, + num_v_heads: int, + kq_head_dim: int, + v_head_dim: int, + kernel_size: int, + pad_len: int = 8, + use_qk_norm_in_gdn: bool = False, +): + seq_idx = pl.program_id(0) + + d_conv_weight_scratch[...] = jnp.zeros_like(d_conv_weight_scratch) + + num_kq_heads = (dim_size - num_v_heads * v_head_dim) // (kq_head_dim * 2) + q_size = num_kq_heads * kq_head_dim + k_size = num_kq_heads * kq_head_dim + v_size = num_v_heads * v_head_dim + repeats = num_v_heads // num_kq_heads + + def chunk_forward(q, k, v, b_val, a_val, a_log_val, dt_bias_val, state_prev): + q = q.astype(jnp.float32) + k = k.astype(jnp.float32) + v = v.astype(jnp.float32) + if use_qk_norm_in_gdn: + q = l2norm(q, dim=-1, eps=1e-6) + k = l2norm(k, dim=-1, eps=1e-6) + scale = 1.0 / jnp.sqrt(kq_head_dim) + q = q * scale + b_val = b_val.astype(jnp.float32) + a_val = a_val.astype(jnp.float32) + a_log_val = a_log_val.astype(jnp.float32) + dt_bias_val = dt_bias_val.astype(jnp.float32) + state_prev = state_prev.astype(jnp.float32) + q_rep = jnp.repeat(q, repeats, axis=1) + k_rep = jnp.repeat(k, repeats, axis=1) + + beta = jax.nn.sigmoid(b_val) + + # EXACT GDN v3 gating formula + log_g = -jnp.exp(a_log_val) * jax.nn.softplus(a_val + dt_bias_val) + v_beta = v * beta[:, :, None] + + # Fast MXU cumsum replacement + mask_cumsum = jnp.tril( + jnp.ones((chunk_size, chunk_size), dtype=log_g.dtype) + ) + cumsum_log_g = jnp.dot(mask_cumsum, log_g) + + diff = cumsum_log_g[:, None, :] - cumsum_log_g[None, :, :] + mask = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=diff.dtype)) + safe_diff = jnp.where(mask[:, :, None] == 1.0, diff, -1e4) + G = jnp.exp(safe_diff) * mask[:, :, None] + G = jnp.transpose(G, (2, 0, 1)) + + q_h = jnp.transpose(q_rep, (1, 0, 2)) + k_h = jnp.transpose(k_rep, (1, 0, 2)) + v_h = jnp.transpose(v_beta, (1, 0, 2)) + + # Dual Causal Attention + attn = jnp.einsum("hck,hdk->hcd", q_h, k_h) + attn_causal = attn * G + + out_intra = jnp.einsum("hcd,hdv->hcv", attn_causal, v_h) + out_intra = jnp.transpose(out_intra, (1, 0, 2)) + + cross_decay = jnp.exp(cumsum_log_g) + q_scaled = q_rep * cross_decay[:, :, None] + out_cross = jnp.einsum("chk,hkv->chv", q_scaled, state_prev) + + out = out_intra + out_cross + + state_decay_end = G[:, chunk_size - 1, :] + state_prev_decayed = state_prev * cross_decay[-1, :, None, None] + + k_scaled = k_h * state_decay_end[:, :, None] + state_new_intra = jnp.einsum("hck,hcv->hkv", k_scaled, v_h) + + state_new = state_prev_decayed + state_new_intra + + return out, state_new + + def fetch_inputs(chunk_idx, slot): + start_idx = chunk_idx * chunk_size + copy_len = chunk_size + pad_len + pltpu.make_async_copy( + padded_pre_conv_qkv_ref.at[seq_idx, pl.ds(start_idx, copy_len)], + padded_pre_conv_qkv_vmem.at[slot], + sem_pre_conv_qkv.at[slot], + ).start() + pltpu.make_async_copy( + qkv_ref.at[seq_idx, chunk_idx], qkv_vmem.at[slot], sem_qkv.at[slot] + ).start() + pltpu.make_async_copy( + b_ref.at[seq_idx, chunk_idx], b_vmem.at[slot], sem_b.at[slot] + ).start() + pltpu.make_async_copy( + a_ref.at[seq_idx, chunk_idx], a_vmem.at[slot], sem_a.at[slot] + ).start() + pltpu.make_async_copy( + do_ref.at[seq_idx, chunk_idx], do_vmem.at[slot], sem_do.at[slot] + ).start() + pltpu.make_async_copy( + chunk_states_ref.at[seq_idx, chunk_idx], + chunk_states_vmem.at[slot], + sem_chunk_states.at[slot], + ).start() + + def wait_inputs(slot): + pltpu.make_async_copy( + padded_pre_conv_qkv_vmem.at[slot], + padded_pre_conv_qkv_vmem.at[slot], + sem_pre_conv_qkv.at[slot], + ).wait() + pltpu.make_async_copy( + qkv_vmem.at[slot], qkv_vmem.at[slot], sem_qkv.at[slot] + ).wait() + pltpu.make_async_copy( + b_vmem.at[slot], b_vmem.at[slot], sem_b.at[slot] + ).wait() + pltpu.make_async_copy( + a_vmem.at[slot], a_vmem.at[slot], sem_a.at[slot] + ).wait() + pltpu.make_async_copy( + do_vmem.at[slot], do_vmem.at[slot], sem_do.at[slot] + ).wait() + pltpu.make_async_copy( + chunk_states_vmem.at[slot], + chunk_states_vmem.at[slot], + sem_chunk_states.at[slot], + ).wait() + + def store_outputs(chunk_idx, slot): + pltpu.make_async_copy( + d_qkv_vmem.at[slot], + d_qkv_ref.at[seq_idx, chunk_idx], + sem_d_qkv.at[slot], + ).start() + pltpu.make_async_copy( + d_b_vmem.at[slot], d_b_ref.at[seq_idx, chunk_idx], sem_d_b.at[slot] + ).start() + pltpu.make_async_copy( + d_a_vmem.at[slot], d_a_ref.at[seq_idx, chunk_idx], sem_d_a.at[slot] + ).start() + + def wait_outputs(slot): + pltpu.make_async_copy( + d_qkv_vmem.at[slot], d_qkv_vmem.at[slot], sem_d_qkv.at[slot] + ).wait() + pltpu.make_async_copy( + d_b_vmem.at[slot], d_b_vmem.at[slot], sem_d_b.at[slot] + ).wait() + pltpu.make_async_copy( + d_a_vmem.at[slot], d_a_vmem.at[slot], sem_d_a.at[slot] + ).wait() + + # Prologue + fetch_inputs(num_chunks - 1, 0) + + # Fetch static inputs (a_log and dt_bias) + pltpu.make_async_copy( + a_log_ref.at[seq_idx], a_log_vmem, sem_a_log.at[0] + ).start() + pltpu.make_async_copy( + dt_bias_ref.at[seq_idx], dt_bias_vmem, sem_dt_bias.at[0] + ).start() + pltpu.make_async_copy(a_log_vmem, a_log_vmem, sem_a_log.at[0]).wait() + pltpu.make_async_copy(dt_bias_vmem, dt_bias_vmem, sem_dt_bias.at[0]).wait() + + d_state = jnp.zeros((num_v_heads, kq_head_dim, v_head_dim), dtype=jnp.float32) + d_conv_weight_acc = jnp.zeros((kernel_size, dim_size), dtype=jnp.float32) + d_a_log_acc = jnp.zeros((num_v_heads,), dtype=jnp.float32) + d_dt_bias_acc = jnp.zeros((num_v_heads,), dtype=jnp.float32) + + def loop_body(i, carry): + d_state, d_conv_weight_acc, d_a_log_acc, d_dt_bias_acc = carry + + chunk_idx = num_chunks - 1 - i + slot = i % 2 + next_slot = (i + 1) % 2 + + wait_inputs(slot) + + def fetch_next(): + fetch_inputs(chunk_idx - 1, next_slot) + + jax.lax.cond(i < num_chunks - 1, fetch_next, lambda: None) + + def wait_prev_out(): + wait_outputs(slot) + + jax.lax.cond(i > 1, wait_prev_out, lambda: None) + + padded_pre_conv_qkv_val = padded_pre_conv_qkv_vmem[slot, ...] + qkv_val = qkv_vmem[slot, ...] + b_val = b_vmem[slot, ...] + a_val = a_vmem[slot, ...] + do_val = do_vmem[slot, ...] + state_prev_val = chunk_states_vmem[slot, ...] + a_log_val = a_log_vmem[...] + dt_bias_val = dt_bias_vmem[...] + + q = qkv_val[:, :q_size].reshape((chunk_size, num_kq_heads, kq_head_dim)) + k = qkv_val[:, q_size : q_size + k_size].reshape( + (chunk_size, num_kq_heads, kq_head_dim) + ) + v = qkv_val[:, q_size + k_size :].reshape( + (chunk_size, num_v_heads, v_head_dim) + ) + + _, vjp_fn = jax.vjp( + chunk_forward, + q, + k, + v, + b_val, + a_val, + a_log_val, + dt_bias_val, + state_prev_val, + ) + ( + d_q, + d_k, + d_v, + d_b_val, + d_a_val, + d_a_log_val, + d_dt_bias_val, + d_state_prev, + ) = vjp_fn((do_val.astype(jnp.float32), d_state)) + + d_qkv = jnp.concatenate( + [ + d_q.reshape(chunk_size, q_size), + d_k.reshape(chunk_size, k_size), + d_v.reshape(chunk_size, v_size), + ], + axis=-1, + ) + + d_cw_rows = [] + for k_idx in range(kernel_size): + start_slice = pad_len - k_idx + shifted = padded_pre_conv_qkv_val[start_slice : start_slice + chunk_size] + d_cw_rows.append(jnp.sum(d_qkv * shifted.astype(jnp.float32), axis=0)) + d_cw = jnp.stack(d_cw_rows, axis=0) + + d_conv_weight_acc += d_cw.astype(jnp.float32) + d_a_log_acc += d_a_log_val.astype(jnp.float32) + d_dt_bias_acc += d_dt_bias_val.astype(jnp.float32) + d_state = d_state_prev.astype(jnp.float32) + + d_qkv_vmem[slot, ...] = d_qkv.astype(d_qkv_vmem.dtype) + d_b_vmem[slot, ...] = d_b_val.astype(d_b_vmem.dtype) + d_a_vmem[slot, ...] = d_a_val.astype(d_a_vmem.dtype) + + store_outputs(chunk_idx, slot) + + return d_state, d_conv_weight_acc, d_a_log_acc, d_dt_bias_acc + + d_state, d_conv_weight_acc, d_a_log_acc, d_dt_bias_acc = jax.lax.fori_loop( + 0, + num_chunks, + loop_body, + (d_state, d_conv_weight_acc, d_a_log_acc, d_dt_bias_acc), + ) + + def wait_last_out(): + wait_outputs((num_chunks - 1) % 2) + + jax.lax.cond(num_chunks > 0, wait_last_out, lambda: None) + + def wait_prev_last_out(): + wait_outputs((num_chunks - 2) % 2) + + jax.lax.cond(num_chunks > 1, wait_prev_last_out, lambda: None) + + d_conv_weight_scratch[...] = d_conv_weight_acc.astype( + d_conv_weight_scratch.dtype + ) + d_a_log_scratch[...] = d_a_log_acc.astype(d_a_log_scratch.dtype) + d_dt_bias_scratch[...] = d_dt_bias_acc.astype(d_dt_bias_scratch.dtype) + + pltpu.make_async_copy( + d_conv_weight_scratch, + d_conv_weight_ref.at[seq_idx, ...], + sem_d_conv_weight.at[0], + ).start() + pltpu.make_async_copy( + d_a_log_scratch, d_a_log_ref.at[seq_idx, ...], sem_d_a_log.at[0] + ).start() + pltpu.make_async_copy( + d_dt_bias_scratch, d_dt_bias_ref.at[seq_idx, ...], sem_d_dt_bias.at[0] + ).start() + + pltpu.make_async_copy( + d_conv_weight_scratch, d_conv_weight_scratch, sem_d_conv_weight.at[0] + ).wait() + pltpu.make_async_copy( + d_a_log_scratch, d_a_log_scratch, sem_d_a_log.at[0] + ).wait() + pltpu.make_async_copy( + d_dt_bias_scratch, d_dt_bias_scratch, sem_d_dt_bias.at[0] + ).wait() + + +def pallas_fused_conv1d_gdn_bwd_computation( + pre_conv_qkv: jax.Array, + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + do: jax.Array, + chunk_states: jax.Array, + conv_weight: jax.Array, + seq_lens: Optional[jax.Array] = None, + *, + num_v_heads: int, + kq_head_dim: int, + v_head_dim: int, + kernel_size: int, + chunk_size: int = 64, + use_qk_norm_in_gdn: bool = False, +) -> Tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + """Executes the Pallas reverse-chunk GDNv3 backward kernel.""" + batch_size, seq_len, dim_size = pre_conv_qkv.shape + num_chunks = seq_len // chunk_size + + pre_conv_qkv_4d = pre_conv_qkv.reshape( + batch_size, num_chunks, chunk_size, dim_size + ) + qkv_4d = qkv.reshape(batch_size, num_chunks, chunk_size, dim_size) + b_4d = b.reshape(batch_size, num_chunks, chunk_size, num_v_heads) + a_4d = a.reshape(batch_size, num_chunks, chunk_size, num_v_heads) + do_4d = do.reshape( + batch_size, num_chunks, chunk_size, num_v_heads, v_head_dim + ) + + if a_log.ndim == 1: + a_log_2d = jnp.broadcast_to(a_log[None, :], (batch_size, num_v_heads)) + else: + a_log_2d = a_log + if dt_bias.ndim == 1: + dt_bias_2d = jnp.broadcast_to(dt_bias[None, :], (batch_size, num_v_heads)) + else: + dt_bias_2d = dt_bias + + if conv_weight.ndim == 3: + conv_weight_2d = conv_weight.squeeze(1) + else: + conv_weight_2d = conv_weight + + if seq_lens is None: + seq_lens = jnp.full((batch_size,), seq_len, dtype=jnp.int32) + + pad_len = ((kernel_size - 1 + 7) // 8) * 8 + pre_conv_pad = jnp.zeros( + (batch_size, pad_len, dim_size), dtype=pre_conv_qkv.dtype + ) + padded_pre_conv_qkv = jnp.concatenate([pre_conv_pad, pre_conv_qkv], axis=1) + + grid = (batch_size,) + hbm_spec = pl.BlockSpec(memory_space=pl.ANY) + + d_qkv_shape = jax.ShapeDtypeStruct(qkv_4d.shape, qkv_4d.dtype) + d_b_shape = jax.ShapeDtypeStruct(b_4d.shape, b_4d.dtype) + d_a_shape = jax.ShapeDtypeStruct(a_4d.shape, a_4d.dtype) + d_conv_weight_shape = jax.ShapeDtypeStruct( + (batch_size, kernel_size, dim_size), conv_weight_2d.dtype + ) + d_a_log_shape = jax.ShapeDtypeStruct( + (batch_size, num_v_heads), a_log_2d.dtype + ) + d_dt_bias_shape = jax.ShapeDtypeStruct( + (batch_size, num_v_heads), dt_bias_2d.dtype + ) + + out_shapes = ( + d_qkv_shape, + d_b_shape, + d_a_shape, + d_conv_weight_shape, + d_a_log_shape, + d_dt_bias_shape, + ) + + scratch_shapes = ( + pltpu.VMEM( + (2, chunk_size + pad_len, dim_size), padded_pre_conv_qkv.dtype + ), + pltpu.VMEM((2, chunk_size, dim_size), qkv_4d.dtype), + pltpu.VMEM((2, chunk_size, num_v_heads), b_4d.dtype), + pltpu.VMEM((2, chunk_size, num_v_heads), a_4d.dtype), + pltpu.VMEM((num_v_heads,), a_log_2d.dtype), + pltpu.VMEM((num_v_heads,), dt_bias_2d.dtype), + pltpu.VMEM((2, chunk_size, num_v_heads, v_head_dim), do_4d.dtype), + pltpu.VMEM((2, num_v_heads, kq_head_dim, v_head_dim), chunk_states.dtype), + pltpu.VMEM((2, chunk_size, dim_size), d_qkv_shape.dtype), + pltpu.VMEM((2, chunk_size, num_v_heads), d_b_shape.dtype), + pltpu.VMEM((2, chunk_size, num_v_heads), d_a_shape.dtype), + pltpu.VMEM((kernel_size, dim_size), d_conv_weight_shape.dtype), + pltpu.VMEM((num_v_heads,), d_a_log_shape.dtype), + pltpu.VMEM((num_v_heads,), d_dt_bias_shape.dtype), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((1,)), + pltpu.SemaphoreType.DMA((1,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((1,)), + pltpu.SemaphoreType.DMA((1,)), + pltpu.SemaphoreType.DMA((1,)), + ) + + d_qkv, d_b, d_a, d_conv_weight, d_a_log, d_dt_bias = pl.pallas_call( + functools.partial( + _pallas_gdn_bwd_kernel, + batch_size=batch_size, + num_chunks=num_chunks, + chunk_size=chunk_size, + dim_size=dim_size, + num_v_heads=num_v_heads, + kq_head_dim=kq_head_dim, + v_head_dim=v_head_dim, + kernel_size=kernel_size, + pad_len=pad_len, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ), + out_shape=out_shapes, + grid=grid, + in_specs=[hbm_spec] * 10, + out_specs=[hbm_spec] * 6, + scratch_shapes=scratch_shapes, + compiler_params=pltpu.CompilerParams( + disable_bounds_checks=True, + ), + )( + padded_pre_conv_qkv, + qkv_4d, + b_4d, + a_4d, + a_log_2d, + dt_bias_2d, + do_4d, + chunk_states, + conv_weight_2d, + seq_lens, + ) + + d_conv_weight_reduced = jnp.sum(d_conv_weight, axis=0) + d_a_log_reduced = jnp.sum(d_a_log, axis=0) + d_dt_bias_reduced = jnp.sum(d_dt_bias, axis=0) + + d_qkv_flat = d_qkv.reshape(batch_size, seq_len, dim_size) + d_b_flat = d_b.reshape(batch_size, seq_len, num_v_heads) + d_a_flat = d_a.reshape(batch_size, seq_len, num_v_heads) + + if conv_weight.ndim == 3: + d_conv_weight_out = d_conv_weight_reduced[:, None, :] + else: + d_conv_weight_out = d_conv_weight_reduced + + return ( + d_qkv_flat, + d_b_flat, + d_a_flat, + d_conv_weight_out, + d_a_log_reduced, + d_dt_bias_reduced, + ) + + +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 _compute_forward_conv_and_states( + 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, + 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 = False, + compute_dtype: jnp.dtype, +) -> Tuple[jax.Array, jax.Array]: + """Computes convolved QKV and inter-chunk states during forward pass.""" + batch_size, seq_len, dim_size = qkv.shape + num_chunks = seq_len // chunk_size + + # Conv1D + conv_input = jnp.pad(qkv, ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) + conv_out = jax.lax.conv_general_dilated( + lhs=conv_input, + rhs=conv_weight.astype(qkv.dtype), + window_strides=(1,), + padding="VALID", + dimension_numbers=("NWC", "WIO", "NWC"), + feature_group_count=dim_size, + ) + 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) + + # Chunk states progression + num_kq_heads = num_k_heads + q_size = num_kq_heads * head_k_dim + k_size = num_kq_heads * head_k_dim + repeats = num_v_heads // num_kq_heads + + q = qkv_conv[:, :, :q_size].reshape( + batch_size, num_chunks, chunk_size, num_kq_heads, head_k_dim + ) + k = qkv_conv[:, :, q_size : q_size + k_size].reshape( + batch_size, num_chunks, chunk_size, num_kq_heads, head_k_dim + ) + v = qkv_conv[:, :, q_size + k_size :].reshape( + batch_size, num_chunks, chunk_size, num_v_heads, head_v_dim + ) + + if use_qk_norm_in_gdn: + q = l2norm(q, dim=-1, eps=1e-6) + k = l2norm(k, dim=-1, eps=1e-6) + + scale = jax.lax.rsqrt(jnp.array(head_k_dim, dtype=jnp.float32)).astype( + compute_dtype + ) + q = q * scale + + b_4d = b.reshape(batch_size, num_chunks, chunk_size, num_v_heads) + a_4d = a.reshape(batch_size, num_chunks, chunk_size, num_v_heads) + + if recurrent_state is None: + init_state = jnp.zeros( + (batch_size, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32 + ) + else: + init_state = recurrent_state.astype(jnp.float32) + + def scan_fn(carry_state, chunk_inputs): + q_i, k_i, v_i, b_i, a_i = chunk_inputs + q_rep = jnp.repeat(q_i, repeats, axis=2) + k_rep = jnp.repeat(k_i, repeats, axis=2) + + beta = jax.nn.sigmoid(b_i) + log_g = -jnp.exp(a_log) * jax.nn.softplus(a_i + dt_bias) + v_beta = v_i * beta[:, :, :, None] + + mask_cumsum = jnp.tril( + jnp.ones((chunk_size, chunk_size), dtype=log_g.dtype) + ) + cumsum_log_g = jnp.einsum( + "cd,bhd->bhc", mask_cumsum, log_g.swapaxes(1, 2) + ).swapaxes(1, 2) + + diff = cumsum_log_g[:, :, None, :] - cumsum_log_g[:, None, :, :] + mask = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=diff.dtype)) + safe_diff = jnp.where(mask[None, :, :, None] == 1.0, diff, -1e4) + G = jnp.exp(safe_diff) * mask[None, :, :, None] + + cross_decay = jnp.exp(cumsum_log_g) + state_decay_end = G[:, chunk_size - 1, :, :] + state_prev_decayed = carry_state * cross_decay[:, -1, :, None, None] + + k_scaled = k_rep * state_decay_end[:, :, :, None] + state_new_intra = jnp.einsum("bchk,bchv->bhkv", k_scaled, v_beta) + + new_state = state_prev_decayed + state_new_intra + return new_state, carry_state + + q_chunks = q.swapaxes(0, 1) + k_chunks = k.swapaxes(0, 1) + v_chunks = v.swapaxes(0, 1) + b_chunks = b_4d.swapaxes(0, 1) + a_chunks = a_4d.swapaxes(0, 1) + + _, chunk_states = jax.lax.scan( + scan_fn, init_state, (q_chunks, k_chunks, v_chunks, b_chunks, a_chunks) + ) + chunk_states = chunk_states.swapaxes(0, 1) + + return qkv_conv, chunk_states + + +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, +): + 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 + Pallas 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, + ) + qkv_conv, chunk_states = _compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_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, + qkv_conv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + chunk_states, + 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, +): + ( + pre_conv_qkv, + qkv_conv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + chunk_states, + conv_state, + recurrent_state, + ) = residuals + d_out, d_states = cotangents + d_conv_state, d_recurrent_state = d_states + + d_qkv, d_b, d_a, d_conv_weight, d_a_log, d_dt_bias = ( + pallas_fused_conv1d_gdn_bwd_computation( + pre_conv_qkv=pre_conv_qkv, + qkv=qkv_conv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=d_out, + chunk_states=chunk_states, + conv_weight=conv_weight, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ) + ) + d_conv_bias = None if conv_bias is None else jnp.zeros_like(conv_bias) + d_conv_state = None if conv_state is None else jnp.zeros_like(conv_state) + d_recurrent_state = ( + None if recurrent_state is None else jnp.zeros_like(recurrent_state) + ) + return ( + d_qkv, + d_b, + d_a, + d_conv_weight, + d_conv_bias, + d_a_log, + d_dt_bias, + 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 9ccd18c2f0..7a0ca4ae3d 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -17,50 +17,83 @@ # pylint: disable=no-name-in-module import functools -from typing import Any, cast import math import os +from typing import Any, cast +from flax import linen as nn +from flax import linen as nn +from flax import nnx import jax -import jax.nn from jax import lax from jax.ad_checkpoint import checkpoint_name -from jax.sharding import Mesh +from jax.experimental import xla_metadata +import jax.nn import jax.numpy as jnp - -from flax import linen as nn -from flax import nnx - -from maxtext.common.common_types import AttentionType, Config, DType, Array, BATCH, EMBED, MODEL_MODE_TRAIN, LENGTH, MODEL_MODE_AUTOREGRESSIVE -from maxtext.common.common_types import KV_BATCH, KV_HEAD -from maxtext.utils.sharding import ( - create_sharding, - get_logical_axis_rules, - logical_to_mesh_axes, - maybe_shard_with_logical, - remove_incompatible_mesh_axes_from_partition_spec, +from jax.sharding import Mesh +from maxtext.common.common_types import ( + Array, + AttentionType, + BATCH, + Config, + DType, + EMBED, + HyperConnectionType, + KV_BATCH, + KV_HEAD, + LENGTH, + MODEL_MODE_AUTOREGRESSIVE, + MODEL_MODE_TRAIN, + ShardMode, ) +from maxtext.inference import kvcache +from maxtext.kernels.attention import gdn_cp from maxtext.layers import attentions from maxtext.layers import initializers as max_initializers +from maxtext.layers import mhc from maxtext.layers import moe +from maxtext.layers import nnx_scan +from maxtext.layers import nnx_scan from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations -from maxtext.layers.embeddings import Qwen3OmniMoeVisionPosEmbedInterpolate, PositionalEmbedding -from maxtext.layers.normalizations import RMSNorm, l2norm, Qwen3NextRMSNorm, Qwen3NextRMSNormGated -from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.layers.attentions import Attention +from maxtext.layers.embeddings import PositionalEmbedding, Qwen3OmniMoeVisionPosEmbedInterpolate +from maxtext.layers.initializers import nd_dense_init, variable_to_logically_partitioned 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.layers.normalizations import Qwen3NextRMSNorm, Qwen3NextRMSNormGated, RMSNorm, l2norm +from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_utils -from maxtext.inference import kvcache - +from maxtext.utils import maxtext_utils +from maxtext.utils.sharding import ( + create_sharding, + get_logical_axis_rules, + logical_to_mesh_axes, + maybe_shard_with_logical, + remove_incompatible_mesh_axes_from_partition_spec, +) # ----------------------------------------- # Qwen3-Next Layer Implementations # ----------------------------------------- +def gdn_context_axes(cfg) -> tuple[str, ...]: + """Mesh axes carrying the GatedDeltaNet sequence, empty when it is replicated. + + Either context knob can carry it, and both can be live at once, so this + returns a tuple that `lax` collectives accept directly. + """ + return tuple( + name + for name, size in ( + ("context", cfg.ici_context_parallelism), + ("context_usp_ulysses", getattr(cfg, "ici_context_usp_ulysses_parallelism", 1)), + ) + if size > 1 + ) + + def naive_jax_chunk_gated_delta_rule( query, key, value, g, beta, chunk_size=64, initial_state=None, use_qk_norm_in_gdn=False ): @@ -195,6 +228,7 @@ def jax_chunk_gated_delta_rule( chunk_size: int = 64, initial_state: None | Array = None, use_qk_norm_in_gdn: bool = False, + cp_axis: None | str = None, compute_dtype: jnp.dtype = jnp.bfloat16, ) -> tuple[Array, None | Array]: """Optimized JAX implementation of Gated Delta Rule.""" @@ -350,7 +384,15 @@ def scan_body(h, args): return h_new, o_c - final_h, o_chunks = lax.scan(scan_body, h_init, xs) + if cp_axis is None: + final_h, o_chunks = lax.scan(scan_body, h_init, xs) + else: + # Sequence is sharded over cp_axis, so a sequential scan over chunks is not + # available. Fold the local chunks into one affine map, exchange those, then + # replay locally from the correct incoming state. See kernels/attention/gdn_cp.py. + A_loc, B_loc = gdn_cp.compose_local(w_scan, u_scan, k_scan, g_scan) + h_in, final_h = gdn_cp.incoming_state(A_loc, B_loc, h_init, cp_axis) + _, o_chunks = lax.scan(scan_body, h_in, xs) # ========================================================================= # STAGE 4: FINALIZATION @@ -618,8 +660,19 @@ def __call__( # mixed_qkvz: (B, S, H_k, 2*D_k + 2*D_v*V_per_K) mixed_qkvz = qkvz.reshape(new_shape_qkvz) if self.mesh is not None: - logical_rules = get_logical_axis_rules() - qkvz_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD, None), mesh=self.mesh, rules=logical_rules) + logical_rules = ( + None + if self.config.using_pipeline_parallelism + else get_logical_axis_rules() + ) + # LENGTH, not None. This with_sharding_constraint told XLA to gather the + # full sequence onto every device for the qkvz projection. With ctx=2 at + # seq 32,768 that is six live bf16[2, 32768, 16, 512] buffers of 1.07 GB + # each -- found by diffing XLA buffer assignment between ctx=1 and ctx=2, + # and the reason the first version of this patch cut GDN memory but still + # lost overall. + cp_len = LENGTH if gdn_context_axes(cfg) else None + qkvz_pspec = logical_to_mesh_axes((KV_BATCH, cp_len, KV_HEAD, None), mesh=self.mesh, rules=logical_rules) # Training microbatches can be smaller than the physical KV_BATCH mesh partition. qkvz_pspec = remove_incompatible_mesh_axes_from_partition_spec( qkvz_pspec, @@ -672,14 +725,17 @@ def __call__( # vLLM PAGED STATE PATH: use tpu_inference fused conv + ragged delta-rule. # ========================================================================= try: - from tpu_inference.layers.common.gdn_attention import run_jax_gdn_attention # pylint: disable=import-outside-toplevel # pytype: disable=import-error - from tpu_inference.layers.common.sharding import ShardingAxisName # pylint: disable=import-outside-toplevel # pytype: disable=import-error - from tpu_inference.layers.common.utils import ( # pylint: disable=import-outside-toplevel # pytype: disable=import-error + # pylint: disable=import-outside-toplevel + # pytype: disable=import-error + from tpu_inference.layers.common.gdn_attention import GdnAttentionConfig, run_jax_gdn_attention + from tpu_inference.layers.common.ragged_gated_delta_rule_wrapper import RaggedGatedDeltaRuleImpl + from tpu_inference.layers.common.sharding import ShardingAxisName + from tpu_inference.layers.common.utils import ( reorder_concatenated_tensor_for_sharding, truncate_sharded_tensor, ) - from tpu_inference.utils import get_mesh_shape_product # pylint: disable=import-outside-toplevel # pytype: disable=import-error - from jax.sharding import PartitionSpec as P_spec # pylint: disable=import-outside-toplevel # pytype: disable=import-error + from tpu_inference.utils import get_mesh_shape_product + from jax.sharding import PartitionSpec as P_spec except ImportError as e: raise ImportError( "GDN attention kernel require the vllm-tpu package. Please install it with `pip install vllm-tpu`." @@ -719,6 +775,11 @@ def __call__( conv_state_paged, recurrent_state_paged = kv_cache + # Use REF impl (pure JAX) to avoid Mosaic kernel compilation issues. + gdn_config = GdnAttentionConfig( + ragged_gated_delta_rule_impl=RaggedGatedDeltaRuleImpl.REF + ) + # Compile against the active request bucket rather than the runner's # maximum-size metadata buffers. dp_size = get_mesh_shape_product(self.mesh, attn_data) @@ -739,30 +800,34 @@ def __call__( dp_size, ) - (new_conv_state_paged, new_recurrent_state_paged), gdn_output = run_jax_gdn_attention( - mixed_qkv, - b_flat, - a_flat, - conv_state_paged, - recurrent_state_paged, - conv_weight, - None, # conv_bias: MaxText conv1d uses use_bias=False. - jnp.asarray(self.A_log[...], dtype=cfg.dtype), - jnp.asarray(self.dt_bias[...], dtype=cfg.dtype), - state_indices, - query_start_loc, - attention_metadata.request_distribution, # pyrefly: ignore[missing-attribute] - seq_lens, - self.num_k_heads, - self.num_v_heads, - self.head_k_dim, - self.head_v_dim, - cfg.gdn_conv_kernel_dim, - mesh=self.mesh, + (new_conv_state_paged, new_recurrent_state_paged), gdn_output = ( + run_jax_gdn_attention( + mixed_qkv, + b_flat, + a_flat, + conv_state_paged, + recurrent_state_paged, + conv_weight, + None, # conv_bias: MaxText conv1d uses use_bias=False. + jnp.asarray(self.A_log[...], dtype=cfg.dtype), + jnp.asarray(self.dt_bias[...], dtype=cfg.dtype), + state_indices, + query_start_loc, + attention_metadata.request_distribution, # pyrefly: ignore[missing-attribute] + seq_lens, + self.num_k_heads, + self.num_v_heads, + self.head_k_dim, + self.head_v_dim, + cfg.gdn_conv_kernel_dim, + mesh=self.mesh, + config=gdn_config, + ) ) # Reshape GDN output and apply gated norm + out projection. gdn_output = gdn_output.reshape(batch, seq_len, self.num_v_heads, self.head_v_dim) + gdn_output = checkpoint_name(gdn_output, "context") gated_output = self.norm(gdn_output, z) gated_output = gated_output.reshape(batch, seq_len, -1) output = self.out_proj(gated_output) @@ -778,7 +843,7 @@ def __call__( v = value.reshape(batch, seq_len, -1) # ========================================================================= - # STEP B: 1D Convolution + # STEP B & C: 1D Convolution & Gated Delta Rule Recurrence # ========================================================================= qkv = jnp.concatenate([q, k, v], axis=-1) batch, seq_len, _ = qkv.shape @@ -791,7 +856,6 @@ def __call__( recurrent_state, conv_state = active_cache.get_gdn_states() orig_cache_batch = conv_state.shape[0] - # 1. Safely shrink/expand conv_state to match incoming qkv (e.g. 16 -> 1) if conv_state.shape[0] != batch: if conv_state.shape[0] == 1: conv_state = jnp.broadcast_to(conv_state, (batch,) + conv_state.shape[1:]) @@ -801,7 +865,6 @@ def __call__( else: conv_state = conv_state[:batch] - # 2. Safely shrink/expand recurrent_state to match incoming qkv if recurrent_state.shape[0] != batch: if recurrent_state.shape[0] == 1: recurrent_state = jnp.broadcast_to(recurrent_state, (batch,) + recurrent_state.shape[1:]) @@ -811,152 +874,308 @@ def __call__( else: recurrent_state = recurrent_state[:batch] - conv_input = jnp.concatenate([conv_state, qkv], axis=1) - - if decoder_segment_ids is not None: - valid_lens = jnp.sum(decoder_segment_ids != 0, axis=1) - - def extract_state(c_in, v_len): - return jax.lax.dynamic_slice_in_dim(c_in, v_len, conv_kernel_size - 1, axis=0) - - next_conv_state = jax.vmap(extract_state)(conv_input, valid_lens) - else: - next_conv_state = conv_input[:, -(conv_kernel_size - 1) :, :] - else: - conv_input = jnp.pad(qkv, ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) - - # Perform the convolution. - conv_out = self.conv1d(conv_input) - # Slice the output to match the original input sequence length. - conv_out = conv_out[:, -seq_len:, :] - qkv_conv = jax.nn.silu(conv_out.astype(jnp.float32)).astype(cfg.dtype) - # q_conv shape: (B, S, key_dim), k_conv shape: (B, S, key_dim), v_conv shape: (B, S, value_dim) - q_conv, k_conv, v_conv = jnp.split(qkv_conv, [self.key_dim, 2 * self.key_dim], axis=-1) - - # Reshape for multi-head processing - # query shape: (B, S, H_k, D_k) - query = q_conv.reshape(batch, seq_len, self.num_k_heads, self.head_k_dim) - # key shape: (B, S, H_k, D_k) - key = k_conv.reshape(batch, seq_len, self.num_k_heads, self.head_k_dim) - # value shape: (B, S, H_v, D_v) - value = v_conv.reshape(batch, seq_len, self.num_v_heads, self.head_v_dim) - - # ========================================================================= - # STEP C: Gated Delta Rule Recurrence - # ========================================================================= - A_log = jnp.asarray(self.A_log[...], dtype=cfg.dtype) - dt_bias = jnp.asarray(self.dt_bias[...], dtype=cfg.dtype) - # beta shape: (B, S, H_v) - beta = jax.nn.sigmoid(b) - # g shape: (B, S, H_v) - g = -jnp.exp(A_log) * jax.nn.softplus(a + dt_bias) - - if decoder_segment_ids is not None: - mask = decoder_segment_ids != 0 - # Apply mask by broadcasting to respective shapes - key = jnp.where(mask[..., None, None], key, 0.0) - value = jnp.where(mask[..., None, None], value, 0.0) - g = jnp.where(mask[..., None], g, 0.0) - - if self.num_v_heads > self.num_k_heads and self.num_v_heads % self.num_k_heads == 0: - repeats = self.num_v_heads // self.num_k_heads - # query shape after repeat: (B, S, H_v, D_k) - query = jnp.repeat(query, repeats, axis=2) - # key shape after repeat: (B, S, H_v, D_k) - key = jnp.repeat(key, repeats, axis=2) - - if seq_len == 1 and model_mode == MODEL_MODE_AUTOREGRESSIVE: - core_attn_out, next_recurrent_state = jax_ar_gated_delta_rule( - query, - key, - value, - g, - beta, - initial_state=recurrent_state, # pyrefly: ignore[bad-argument-type] - use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, - compute_dtype=cfg.dtype, + if getattr(cfg, "use_gdn_kernel", False) and getattr( + cfg, "use_hybrid_gdn", False + ): + from maxtext.models.hybrid_gdn import hybrid_fused_conv1d_gdn + + 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, + ) ) - elif self.mesh is not None: - logical_rules = get_logical_axis_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) - ) - qkv_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD, None), mesh=self.mesh, rules=logical_rules) - g_beta_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD), mesh=self.mesh, rules=logical_rules) - state_pspec = logical_to_mesh_axes((KV_BATCH, KV_HEAD, None, None), mesh=self.mesh, rules=logical_rules) - # Keep every shard_map input/output batch spec consistent when replication is required. - qkv_pspec = remove_incompatible_mesh_axes_from_partition_spec( - qkv_pspec, - query.shape, - self.mesh, - dims=(0,), - allow_remove_axes=True, - ) - g_beta_pspec = remove_incompatible_mesh_axes_from_partition_spec( - g_beta_pspec, - g.shape, - self.mesh, - dims=(0,), - allow_remove_axes=True, + else jnp.zeros( + (batch, self.num_v_heads, self.head_k_dim, self.head_v_dim), + dtype=cfg.dtype, + ) ) - state_pspec = remove_incompatible_mesh_axes_from_partition_spec( - state_pspec, - recurrent_state_arg.shape, - self.mesh, - dims=(0,), - allow_remove_axes=True, + 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) ) + 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 + ) - @functools.partial( - jax.shard_map, - mesh=self.mesh, - in_specs=( - qkv_pspec, # query - qkv_pspec, # key - qkv_pspec, # value - g_beta_pspec, # g - g_beta_pspec, # beta - state_pspec, # initial_state - ), - out_specs=( - qkv_pspec, # core_attn_out - state_pspec, # final_state - ), - check_vma=False, + @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_arg, + recurrent_state=recurrent_state_arg, + 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, + ) + ) + else: + if conv_state is not None: + conv_input = jnp.concatenate([conv_state, qkv], axis=1) + if decoder_segment_ids is not None: + valid_lens = jnp.sum(decoder_segment_ids != 0, axis=1) + + def extract_state(c_in, v_len): + return jax.lax.dynamic_slice_in_dim( + c_in, v_len, conv_kernel_size - 1, axis=0 + ) + + next_conv_state = jax.vmap(extract_state)(conv_input, valid_lens) + else: + next_conv_state = conv_input[:, -(conv_kernel_size - 1) :, :] + else: + conv_input = jnp.pad(qkv, ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) + + conv_out = self.conv1d(conv_input) + conv_out = conv_out[:, -seq_len:, :] + qkv_conv = jax.nn.silu(conv_out.astype(jnp.float32)).astype(cfg.dtype) + q_conv, k_conv, v_conv = jnp.split( + qkv_conv, [self.key_dim, 2 * self.key_dim], axis=-1 ) - def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): - return jax_chunk_gated_delta_rule( - query=q, - key=k, - value=v, - g=g_val, - beta=beta_val, + + query = q_conv.reshape(batch, seq_len, self.num_k_heads, self.head_k_dim) + key = k_conv.reshape(batch, seq_len, self.num_k_heads, self.head_k_dim) + value = v_conv.reshape(batch, seq_len, self.num_v_heads, self.head_v_dim) + + A_log = jnp.asarray(self.A_log[...], dtype=cfg.dtype) + dt_bias = jnp.asarray(self.dt_bias[...], dtype=cfg.dtype) + beta = jax.nn.sigmoid(b) + g = -jnp.exp(A_log) * jax.nn.softplus(a + dt_bias) + + if decoder_segment_ids is not None: + mask = decoder_segment_ids != 0 + key = jnp.where(mask[..., None, None], key, 0.0) + value = jnp.where(mask[..., None, None], value, 0.0) + g = jnp.where(mask[..., None], g, 0.0) + + if ( + self.num_v_heads > self.num_k_heads + and self.num_v_heads % self.num_k_heads == 0 + ): + repeats = self.num_v_heads // self.num_k_heads + query = jnp.repeat(query, repeats, axis=2) + key = jnp.repeat(key, repeats, axis=2) + + if seq_len == 1 and model_mode == MODEL_MODE_AUTOREGRESSIVE: + core_attn_out, next_recurrent_state = jax_ar_gated_delta_rule( + query, + key, + value, + g, + beta, + initial_state=recurrent_state, + use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, + compute_dtype=cfg.dtype, + cp_axis=cp_axes or None, + ) + 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=init_h, + 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 = ( + None + if self.config.using_pipeline_parallelism + else get_logical_axis_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, + ) + ) + cp_axes = gdn_context_axes(cfg) + cp_len = LENGTH if cp_axes else None + qkv_pspec = logical_to_mesh_axes( + (KV_BATCH, cp_len, KV_HEAD, None), + mesh=self.mesh, + rules=logical_rules, + ) + g_beta_pspec = logical_to_mesh_axes( + (KV_BATCH, cp_len, KV_HEAD), mesh=self.mesh, rules=logical_rules + ) + state_pspec = logical_to_mesh_axes( + (KV_BATCH, KV_HEAD, None, None), mesh=self.mesh, rules=logical_rules + ) + # Keep every shard_map input/output batch spec consistent when replication is required. + qkv_pspec = remove_incompatible_mesh_axes_from_partition_spec( + qkv_pspec, + query.shape, + self.mesh, + dims=(0,), + allow_remove_axes=True, + ) + g_beta_pspec = remove_incompatible_mesh_axes_from_partition_spec( + g_beta_pspec, + g.shape, + self.mesh, + dims=(0,), + allow_remove_axes=True, + ) + state_pspec = remove_incompatible_mesh_axes_from_partition_spec( + state_pspec, + recurrent_state_arg.shape, + self.mesh, + dims=(0,), + allow_remove_axes=True, + ) - core_attn_out, next_recurrent_state = shard_mapped_delta_rule(query, key, value, g, beta, recurrent_state_arg) - else: - 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, - ) + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=( + qkv_pspec, # query + qkv_pspec, # key + qkv_pspec, # value + g_beta_pspec, # g + g_beta_pspec, # beta + state_pspec, # initial_state + ), + out_specs=( + qkv_pspec, # core_attn_out + state_pspec, # final_state + ), + check_vma=False, + ) + def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): + return jax_chunk_gated_delta_rule( + query=q, + key=k, + value=v, + g=g_val, + beta=beta_val, + chunk_size=cfg.gdn_chunk_size, + initial_state=init_h, + use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, + compute_dtype=cfg.dtype, + cp_axis=cp_axes or None, + ) + + core_attn_out, next_recurrent_state = shard_mapped_delta_rule( + query, key, value, g, beta, recurrent_state_arg + ) + else: + 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, + ) if model_mode != MODEL_MODE_TRAIN and active_cache is not None: assert next_conv_state is not None assert next_recurrent_state is not None - if next_conv_state.shape[0] != orig_cache_batch: # pyrefly: ignore[unbound-name] + if next_conv_state.shape[0] != orig_cache_batch: if next_conv_state.shape[0] == 1: next_conv_state = jnp.broadcast_to(next_conv_state, (orig_cache_batch,) + next_conv_state.shape[1:]) next_recurrent_state = jnp.broadcast_to( @@ -971,7 +1190,9 @@ def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): next_recurrent_state = next_recurrent_state[:orig_cache_batch] if model_mode != MODEL_MODE_TRAIN and active_cache is not None: - active_cache.update_gdn_states(next_recurrent_state, next_conv_state) # pyrefly: ignore[bad-argument-type] + active_cache.update_gdn_states(next_recurrent_state, next_conv_state) + + core_attn_out = checkpoint_name(core_attn_out, "context") # ========================================================================= # STEP D: Final Output Stage @@ -1007,12 +1228,12 @@ def init_kv_caches(self, batch_size: int): value_heads=self.num_v_heads, key_head_size=self.head_k_dim, value_head_size=self.head_v_dim, - dtype=self.dtype, # pyrefly: ignore[missing-attribute] + dtype=self.dtype, is_gdn=True, conv_kernel_size=conv_kernel_size, conv_dim=conv_dim, - model_mode=self.model_mode, # pyrefly: ignore[missing-attribute] - rngs=self.rngs, # pyrefly: ignore[missing-attribute] + model_mode=self.model_mode, + rngs=self.rngs, ) @@ -1140,16 +1361,17 @@ def __init__(self, config: Config, mesh: Mesh, quant: None | Quant = None, *, rn rngs=rngs, ) - # 2. Instantiate and apply the shared expert, if the config asks for one. + # 2. Instantiate and apply the shared expert(s), if the config asks for one. # cfg.shared_experts defaults to 0. Building the expert every time puts an # extra full-size MLP in every layer of a dense configuration. self.use_shared_expert = cfg.shared_experts > 0 if self.use_shared_expert: + shared_expert_mlp_dim = maxtext_utils.get_shared_expert_mlp_dim(cfg) self.shared_expert = MlpBlock( config=cfg, mesh=mesh, in_features=cfg.emb_dim, - intermediate_dim=cfg.moe_mlp_dim, + intermediate_dim=cfg.shared_experts * shared_expert_mlp_dim, activations=cfg.mlp_activations, intermediate_dropout_rate=cfg.dropout_rate, dtype=cfg.dtype, @@ -1159,17 +1381,22 @@ def __init__(self, config: Config, mesh: Mesh, quant: None | Quant = None, *, rn rngs=rngs, ) - # 3. Instantiate and apply the gate for the shared expert. - self.shared_expert_gate = DenseGeneral( - in_features_shape=cfg.emb_dim, - out_features_shape=1, - use_bias=False, # Qwen3-Next shared_expert_gate does not have a bias - dtype=cfg.dtype, - kernel_init=max_initializers.nd_dense_init(cfg.dense_init_scale, "fan_in", "truncated_normal"), - kernel_axes=("embed", None), - matmul_precision=cfg.matmul_precision, - rngs=rngs, - ) + # 3. Instantiate the (optional) gate for the shared expert. + if cfg.moe_shared_expert_gate: + self.shared_expert_gate = DenseGeneral( + in_features_shape=cfg.emb_dim, + out_features_shape=1, + use_bias=False, # Qwen3-Next shared_expert_gate does not have a bias + dtype=cfg.dtype, + kernel_init=max_initializers.nd_dense_init( + cfg.dense_init_scale, "fan_in", "truncated_normal" + ), + kernel_axes=("embed", None), + matmul_precision=cfg.matmul_precision, + rngs=rngs, + ) + else: + self.shared_expert_gate = None else: self.shared_expert = None self.shared_expert_gate = None @@ -1179,44 +1406,59 @@ def __call__( hidden_states: Array, deterministic: bool, forced_routed_experts: jnp.ndarray | None = None, - ) -> tuple[Array, Array | None]: - """ - Applies the sparse MoE block to the input hidden states. + ) -> tuple[Array, Array | None, Array | None]: + """Applies the sparse MoE block to the input hidden states. Args: - hidden_states: The input array from the previous layer. Shape: (batch, seq, embed_dim) + hidden_states: The input array from the previous layer. Shape: (batch, + seq, embed_dim) deterministic: If True, disables dropout. Returns: A tuple containing: - The output array of the MoE block. - - The load balancing loss from the routed experts, if applicable during training. + - The load balancing loss from the routed experts, if applicable during + training. + - The aux-loss-free expert-bias updates from the routed experts, if + applicable. """ # 1. Apply the routed experts block. - routed_output, load_balance_loss, _ = self.routed_experts(hidden_states, forced_routed_experts=forced_routed_experts) + routed_output, load_balance_loss, moe_bias_updates = self.routed_experts( + hidden_states, forced_routed_experts=forced_routed_experts + ) if not self.use_shared_expert: - return routed_output, load_balance_loss + return routed_output, load_balance_loss, moe_bias_updates # 2. Apply the shared expert. shared_expert_output = self.shared_expert(hidden_states, deterministic=deterministic) - # 3. Apply the gate for the shared expert. - shared_gate_output = self.shared_expert_gate(hidden_states) + # 3. Apply the (optional) gate for the shared expert. + if self.shared_expert_gate is not None: + shared_gate_output = self.shared_expert_gate(hidden_states) + shared_expert_output = ( + jax.nn.sigmoid(shared_gate_output) * shared_expert_output + ) # 4. Combine the outputs. - final_output = routed_output + jax.nn.sigmoid(shared_gate_output) * shared_expert_output + final_output = routed_output + shared_expert_output - return final_output, load_balance_loss + return final_output, load_balance_loss, moe_bias_updates class Qwen3NextScannableBlock(nnx.Module): - """A scannable block of Qwen3-Next decoder layers. + """A scannable block of Qwen3-Next decoder layers with hierarchical nested scans. + + One block covers a single period of the attention pattern defined by + `config.inhomogeneous_layer_cycle_interval`: several linear-attention + (GatedDeltaNet) layers plus one full-attention layer. The linear-attention + layers are homogeneous, so they are stacked and run through + `nnx_scan.apply_scanned_layers`; the lone full-attention layer runs inside a + trip-count-one `jax.lax.scan` that acts as an XLA scheduling barrier. - 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. + Nesting the scans this way lets each sub-layer be rematerialized on its own + (`apply_internal_remat`) instead of rematerializing the whole block, so only + one sub-layer's activations are live at a time. Attributes: config: The model configuration object. @@ -1225,83 +1467,254 @@ class Qwen3NextScannableBlock(nnx.Module): quant: Optional quantization configuration. """ - def __init__(self, config: Config, mesh: Mesh, model_mode: str, quant: None | Quant = None, *, rngs: nnx.Rngs): + def __init__( + self, + config: Config, + mesh: Mesh, + model_mode: str, + quant: None | Quant = None, + *, + num_of_layers: int | None = None, + layer_idx_offset: int = 0, + remat_policy_fn: Any | None = None, + apply_internal_remat: bool = False, + rngs: nnx.Rngs, + ): self.config = config self.mesh = mesh self.model_mode = model_mode self.quant = quant self.rngs = rngs + self.remat_policy_fn = remat_policy_fn + self.apply_internal_remat = apply_internal_remat cfg = self.config + self.is_mhc_enabled = cfg.mhc_expansion_rate > 1 + if self.is_mhc_enabled: + self.activation_axis_names = ( + "activation_batch", + "activation_norm_length", + "mhc", + "activation_embed", + ) + else: + self.activation_axis_names = ( + "activation_batch", + "activation_norm_length", + "activation_embed", + ) + if num_of_layers is None: + num_of_layers = cfg.inhomogeneous_layer_cycle_interval + self.num_of_layers = num_of_layers + self.layer_idx_offset = layer_idx_offset + + cycle_interval = cfg.inhomogeneous_layer_cycle_interval + # Qwen3-Next puts the full-attention layer last in every cycle, which is what + # Qwen3NextDecoderLayer derives from layer_idx when it is not told explicitly. + full_attention_offset = ( + getattr(cfg, "full_attention_layer_offset", -1) % cycle_interval + ) - # 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( + positions = [(layer_idx_offset + i) % cycle_interval for i in range(num_of_layers)] + self.num_local = sum(1 for p in positions if p != full_attention_offset) + self.num_global = sum(1 for p in positions if p == full_attention_offset) + if self.num_global > 1: + raise ValueError( + f"A Qwen3-Next scannable block spans {num_of_layers} layers starting at offset {layer_idx_offset}, which " + f"covers {self.num_global} full-attention layers; the block supports at most one." + ) + # The local scan runs before the global layer, so the block only reproduces the + # model's layer order when the full-attention layer is last in the period. + if self.num_global == 1 and positions[-1] != full_attention_offset: + raise ValueError( + f"Qwen3-Next scannable block expects the full-attention layer last in the block, but a block of " + f"{num_of_layers} layers starting at layer_idx_offset={layer_idx_offset} lands it at block position " + f"{positions.index(full_attention_offset)}. Blocks must start on a cycle boundary." + ) + + 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, + is_dense_layer=False, + is_full_attention_layer=False, + 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, + layer_idx=full_attention_offset, + is_dense_layer=False, + is_full_attention_layer=True, + rngs=self.rngs, + ) + 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, ) - setattr(self, layer_name, layer) + + 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. + + Inference KV caches are a Python list of per-layer entries, so this path + unrolls the local layers statically rather than scanning them. + """ + 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) + # Collect only non-Param state: parameters are read-only here, so stacking + # them back would allocate a second copy of every layer weight. Non-Param + # state is stacked on axis 0, matching nnx_scan.create_scanned_layers. + _, _, updated_state = nnx.split(layer, nnx.Param, ...) + per_layer_states.append(updated_state) + + nnx.update(self.local_layers, jax.tree.map(lambda *xs: jnp.stack(xs), *per_layer_states)) + + 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, - decoder_segment_ids: None | jnp.ndarray, - decoder_positions: None | jnp.ndarray, - deterministic: bool, - model_mode: str, + decoder_segment_ids: None | jnp.ndarray = None, + decoder_positions: None | jnp.ndarray = None, + deterministic: bool = False, + model_mode: str = "train", previous_chunk=None, slot: None | int = None, kv_cache=None, attention_metadata=None, ) -> tuple[Array, None]: - """Applies the block of decoder layers to the input carry. + cfg = self.config + inputs = carry + inputs = nn.with_logical_constraint( + inputs, + self.activation_axis_names, + ) - 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): - """ - This layer is a hybrid, capable of functioning as either: - 1. A standard attention + MoE layer. - 2. A linear attention + MoE layer. + """This layer is a hybrid, capable of functioning as either: 1. + + A standard attention + MoE layer. 2. A linear attention + MoE layer. - NOTE: This implementation assumes every layer contains a MoE block, which is true for - models like Qwen3-Next-80B-A3B where `decoder_sparse_step=1`. For models that - interleave dense and sparse MLP layers, conditional logic would be needed here. + The first `config.first_num_dense_layers` layers (by `layer_idx`) use a plain + dense MLP instead of MoE, and always use full attention, mirroring DeepSeek + V3's + dense-prefix pattern (see `models/deepseek.py::DeepSeekDenseLayer`). Attributes: config: The model configuration object. @@ -1312,7 +1725,16 @@ class Qwen3NextDecoderLayer(nnx.Module): """ def __init__( - self, config: Config, mesh: Mesh, model_mode: str, layer_idx: int, quant: None | Quant = None, *, rngs: nnx.Rngs + self, + config: Config, + mesh: Mesh, + model_mode: str, + layer_idx: int, + quant: None | Quant = None, + *, + is_dense_layer: bool | None = None, + is_full_attention_layer: bool | None = None, + rngs: nnx.Rngs, ): self.config = config self.mesh = mesh @@ -1320,7 +1742,24 @@ def __init__( self.layer_idx = layer_idx self.quant = quant cfg = self.config - self.activation_axis_names = ("activation_batch", "activation_norm_length", "activation_embed") + self.is_mhc_enabled = cfg.mhc_expansion_rate > 1 + if self.is_mhc_enabled: + self.activation_axis_names = ( + "activation_batch", + "activation_norm_length", + "mhc", + "activation_embed", + ) + else: + self.activation_axis_names = ( + "activation_batch", + "activation_norm_length", + "activation_embed", + ) + + if is_dense_layer is None: + is_dense_layer = layer_idx < cfg.first_num_dense_layers + self.is_dense_layer = is_dense_layer # First LayerNorm, applied before the attention block. self.input_layernorm = Qwen3NextRMSNorm( @@ -1331,8 +1770,21 @@ def __init__( rngs=rngs, ) - # Determine the type of attention mechanism for the current layer. - is_full_attention_layer = (self.layer_idx + 1) % cfg.inhomogeneous_layer_cycle_interval == 0 + # Determine the type of attention mechanism for the current layer. Dense layers + # always use full attention (see class docstring). `full_attention_layer_offset` + # picks which position in the cycle is full attention; -1 (Python's negative-modulo + # wraps to cycle-1) reproduces the original "last position in the cycle" schedule. + full_attention_offset = ( + getattr(cfg, "full_attention_layer_offset", -1) + % cfg.inhomogeneous_layer_cycle_interval + ) + if is_full_attention_layer is None: + is_full_attention_layer = ( + self.is_dense_layer + or self.layer_idx % cfg.inhomogeneous_layer_cycle_interval + == full_attention_offset + ) + self.is_full_attention_layer = is_full_attention_layer # Conditionally instantiate either the Linear Attention or Full Attention block. if is_full_attention_layer: @@ -1360,8 +1812,92 @@ def __init__( rngs=rngs, ) - # Instantiate our `Qwen3NextSparseMoeBlock`. - self.mlp = Qwen3NextSparseMoeBlock(config=cfg, mesh=self.mesh, quant=self.quant, rngs=rngs) + # Dense layers use a plain MLP; all other layers use `Qwen3NextSparseMoeBlock`. + if self.is_dense_layer: + self.mlp = MlpBlock( + in_features=cfg.emb_dim, + intermediate_dim=cfg.mlp_dim, + activations=cfg.mlp_activations, + intermediate_dropout_rate=cfg.dropout_rate, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + config=cfg, + mesh=self.mesh, + quant=self.quant, + model_mode=model_mode, + rngs=rngs, + ) + else: + self.mlp = Qwen3NextSparseMoeBlock( + config=cfg, mesh=self.mesh, quant=self.quant, rngs=rngs + ) + + # Manifold-Constrained Hyper Connections: replaces the plain residual add around the + # attention and MoE branches with a learned multi-stream mixing. See maxtext/layers/mhc.py + # and models/deepseek4.py::DeepSeek4DecoderLayer for the reference implementation. + if self.is_mhc_enabled: + self.mhc_attention = mhc.ManifoldConstrainedHyperConnections( + cfg, cfg.emb_dim, self.mesh, rngs + ) + self.mhc_mlp = mhc.ManifoldConstrainedHyperConnections( + cfg, cfg.emb_dim, self.mesh, rngs + ) + + def pre_attention_norm_op(self, x): + normed = self.input_layernorm(x) + return nn.with_logical_constraint( + normed, + ("activation_batch", "activation_norm_length", "activation_embed"), + ) + + def post_attention_norm_op(self, x): + normed = self.post_attention_layernorm(x) + return nn.with_logical_constraint( + normed, + ("activation_batch", "activation_norm_length", "activation_embed"), + ) + + def attention_branch( + self, + inputs_q, + inputs_kv=None, + decoder_segment_ids=None, + inputs_positions=None, + deterministic=None, + model_mode=None, + kv_cache=None, + attention_metadata=None, + **kwargs, + ): + """Adapts Qwen3-Next's two attention variants to mHC's inputs_q/inputs_kv branch_fn convention.""" + del inputs_kv, kwargs + if isinstance(self.attention, Qwen3NextFullAttention): + return self.attention( + inputs_q, + decoder_segment_ids, + inputs_positions, + deterministic, + model_mode, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + ) + return self.attention( + inputs_q, + model_mode=model_mode, + kv_cache=kv_cache, + decoder_segment_ids=decoder_segment_ids, + attention_metadata=attention_metadata, + ) + + def mlp_op(self, inputs, deterministic, *args, **kwargs): + """Adapts the dense/MoE MLP's return shape to mHC's MLP_MOE 3-tuple convention.""" + del args, kwargs + if self.is_dense_layer: + return self.mlp(inputs, deterministic=deterministic), None, None + mlp_out, load_balance_loss, moe_bias_updates = self.mlp( + inputs, deterministic=deterministic + ) + return mlp_out, load_balance_loss, moe_bias_updates def __call__( self, @@ -1378,6 +1914,50 @@ 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] + + inputs = nn.with_logical_constraint(inputs, self.activation_axis_names) + # inputs = checkpoint_name(inputs, "decoder_layer_input") + + if self.is_mhc_enabled: + intermediate_inputs, _ = self.mhc_attention( + self.pre_attention_norm_op, + self.attention_branch, + x=inputs, + mhc_type=HyperConnectionType.ATTENTION, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=deterministic, + model_mode=model_mode, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + ) + + layer_output, metadata = self.mhc_mlp( + self.post_attention_norm_op, + self.mlp_op, + x=intermediate_inputs, + mhc_type=HyperConnectionType.MLP_MOE, + deterministic=deterministic, + ) + load_balance_loss = metadata.get("load_balance_loss", None) + 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) + moe_bias_updates = metadata.get("moe_bias_updates", None) + if ( + self.config.routed_bias + and self.config.routed_bias_update_rate > 0.0 + and moe_bias_updates is not None + ): + self.moe_bias_updates = nnx.Intermediate(moe_bias_updates) + + layer_output = nn.with_logical_constraint( + layer_output, self.activation_axis_names + ) + return layer_output, kv_cache + residual = inputs # First LayerNorm, applied before the attention block. @@ -1386,13 +1966,15 @@ def __call__( # Conditionally apply either the Linear Attention or Full Attention block. if isinstance(self.attention, Qwen3NextFullAttention): - attention_output, new_kv_cache = cast(Qwen3NextFullAttention, self.attention)( + attention_output, new_kv_cache = cast( + Qwen3NextFullAttention, self.attention + )( hidden_states, decoder_segment_ids, decoder_positions, deterministic, model_mode, - kv_cache=kv_cache, # pyrefly: ignore[bad-argument-type] + kv_cache=kv_cache, attention_metadata=attention_metadata, ) else: @@ -1415,13 +1997,26 @@ def __call__( hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = nn.with_logical_constraint(hidden_states, self.activation_axis_names) - # Instantiate and call our `Qwen3NextSparseMoeBlock`. - mlp_output, load_balance_loss = self.mlp(hidden_states, deterministic=deterministic) - - # We sow the load balancing loss so it can be collected and added to the total loss - # during training. - 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) + # Apply the dense MLP or `Qwen3NextSparseMoeBlock`. + if self.is_dense_layer: + mlp_output = self.mlp(hidden_states, deterministic=deterministic) + else: + mlp_output, load_balance_loss, moe_bias_updates = self.mlp( + hidden_states, deterministic=deterministic + ) + # We sow the load balancing loss so it can be collected and added to the total loss + # during training. + 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) + if ( + self.config.routed_bias + and self.config.routed_bias_update_rate > 0.0 + and moe_bias_updates is not None + ): + self.moe_bias_updates = nnx.Intermediate(moe_bias_updates) # Final residual connection (after the MoE block) layer_output = residual + mlp_output @@ -1457,15 +2052,22 @@ def __init__( # Physical shardings used to pin sublayer outputs under ShardMode.EXPLICIT. In # ShardMode.AUTO the callees ignore these and let GSPMD infer the layout. - self.out_sharding = create_sharding(mesh, self.activation_axis_names, rules=get_logical_axis_rules()) - self.mlp_intermediate_sharding = create_sharding(mesh, self.mlp_activation_axis_names, rules=get_logical_axis_rules()) - self._maybe_shard_with_logical = functools.partial( - maybe_shard_with_logical, - mesh=mesh, - shard_mode=config.shard_mode, - debug_sharding=config.debug_sharding, - extra_stack_level=1, - ) + if config.shard_mode == ShardMode.EXPLICIT: + self.out_sharding = create_sharding(mesh, self.activation_axis_names, rules=get_logical_axis_rules()) + self.mlp_intermediate_sharding = create_sharding( + mesh, self.mlp_activation_axis_names, rules=get_logical_axis_rules() + ) + self._maybe_shard_with_logical = functools.partial( + maybe_shard_with_logical, + mesh=mesh, + shard_mode=config.shard_mode, + debug_sharding=config.debug_sharding, + extra_stack_level=1, + ) + else: + self.out_sharding = None + self.mlp_intermediate_sharding = None + self._maybe_shard_with_logical = lambda inputs, *args, **kwargs: inputs # Corresponds to Qwen3's `input_layernorm` self.pre_self_attention_layer_norm = RMSNorm( @@ -1701,8 +2303,10 @@ def update_cache(cache, val): return cache.at[layer_idx].set(val) return cache - stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache) # pyrefly: ignore[unbound-name] - return (layer_output, stacked_kv_cache, layer_idx + 1), None # pyrefly: ignore[unbound-name] + stacked_kv_cache = jax.tree_util.tree_map( + update_cache, stacked_kv_cache, kv_cache + ) + return (layer_output, stacked_kv_cache, layer_idx + 1), None else: return layer_output, kv_cache @@ -1729,8 +2333,10 @@ def __init__( use_postshuffle_norm: bool = False, dtype: DType = jnp.float32, weight_dtype: DType = jnp.float32, - kernel_init: max_initializers.NdInitializer = max_initializers.nd_dense_init(1.0, "fan_in", "normal"), - rngs: nnx.Rngs = None, # pyrefly: ignore[bad-function-definition] + kernel_init: max_initializers.NdInitializer = max_initializers.nd_dense_init( + 1.0, "fan_in", "normal" + ), + rngs: nnx.Rngs = None, ): """Initializes the Qwen3Omni vision patch merger. @@ -1843,8 +2449,10 @@ def __init__( config: Config, dtype: DType = jnp.float32, weight_dtype: DType = jnp.float32, - kernel_init: max_initializers.NdInitializer = max_initializers.nd_dense_init(1.0, "fan_in", "normal"), - rngs: nnx.Rngs = None, # pyrefly: ignore[bad-function-definition] + kernel_init: max_initializers.NdInitializer = max_initializers.nd_dense_init( + 1.0, "fan_in", "normal" + ), + rngs: nnx.Rngs = None, ): """Initializes the Qwen3Omni vision MLP. @@ -1921,7 +2529,7 @@ def __init__( # Default to float32 for numerical stability in 3D convolutions on image/video inputs dtype: DType = jnp.float32, weight_dtype: DType = jnp.float32, - rngs: nnx.Rngs = None, # pyrefly: ignore[bad-function-definition] + rngs: nnx.Rngs = None, ): """Initializes the Qwen3Omni vision patch embedding. @@ -1973,8 +2581,16 @@ def __call__(self, hidden_states: Array, video_mask: Array | None = None) -> tup attention_mask = None if video_mask is not None: - mask_patch_elements = self.temporal_patch_size * self.patch_size * self.patch_size - attention_mask = video_mask.reshape(video_mask.shape[0], -1, mask_patch_elements).max(axis=-1).astype(jnp.int32) + patch_mask = video_mask[ + :, + 0, + :: self.temporal_patch_size, + :: self.patch_size, + :: self.patch_size, + ] + attention_mask = patch_mask.reshape(video_mask.shape[0], -1).astype( + jnp.int32 + ) return hidden_states, attention_mask @@ -1987,7 +2603,7 @@ class Qwen3OmniMoeVisionAttention(nnx.Module): attn: Underlying attention module """ - def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): """Initializes the Qwen3Omni vision attention layer. Args: @@ -2011,7 +2627,7 @@ def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): # pyre float32_logits=self.config.float32_logits, dtype=self.config.dtype_mm, weight_dtype=self.config.weight_dtype, - mesh=mesh, # pyrefly: ignore[bad-argument-type] + mesh=mesh, dropout_rate=0.0, attention_type=AttentionType.FULL, is_nope_layer=False, @@ -2029,8 +2645,6 @@ def __call__( num_frames: int, height: int, width: int, - attention_mask: Array | None = None, - valid_grid: tuple[int, int, int] | None = None, deterministic: bool = True, ) -> Array: """ @@ -2039,8 +2653,6 @@ def __call__( num_frames: Number of temporal frames (static) height: Height in patches (static) width: Width in patches (static) - attention_mask: Optional mask identifying valid tokens in the padded sequence. - valid_grid: Optional unpadded `(frames, height, width)` grid used for vision RoPE. deterministic: Whether to use deterministic mode (disable dropout) Returns: @@ -2051,14 +2663,11 @@ def __call__( "num_frames": num_frames, "height": height, "width": width, - "token_mask": attention_mask, - "valid_grid": valid_grid, } output, _ = self.attn( inputs_q=hidden_states, inputs_kv=hidden_states, deterministic=deterministic, - decoder_segment_ids=attention_mask, rope_kwargs=rope_kwargs, ) @@ -2077,7 +2686,7 @@ class Qwen3OmniMoeVisionBlock(nnx.Module): mlp_out: Second MLP layer """ - def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): """Initializes the Qwen3Omni vision transformer block. Args: @@ -2111,8 +2720,6 @@ def __call__( num_frames: int, height: int, width: int, - attention_mask: Array | None = None, - valid_grid: tuple[int, int, int] | None = None, ) -> Array: """ Args: @@ -2125,12 +2732,7 @@ def __call__( Output tensor of shape (batch, T*H*W, hidden_size) """ x = x + self.attn( - self.ln1(x), - num_frames=num_frames, - height=height, - width=width, - attention_mask=attention_mask, - valid_grid=valid_grid, + self.ln1(x), num_frames=num_frames, height=height, width=width ) y = self.ln2(x) y = self.mlp(y) @@ -2152,7 +2754,7 @@ class Qwen3OmniMoeVisionEncoder(nnx.Module): deep_idx: Indices of layers to extract deep features from """ - def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): """Initializes the Qwen3Omni vision encoder. Args: @@ -2192,8 +2794,6 @@ def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None): # pyre def __call__( self, hidden_states: Array, - video_mask: Array | None = None, - video_grid_thw: Array | tuple[int, int, int] | None = None, deterministic: bool = True, ): """ @@ -2210,12 +2810,6 @@ def __call__( num_frames = num_frames // self.config.temporal_patch_size_for_vit height = height // self.config.patch_size_for_vit width = width // self.config.patch_size_for_vit - attention_mask = None - if video_mask is not None: - mask_patch_elements = ( - self.config.temporal_patch_size_for_vit * self.config.patch_size_for_vit * self.config.patch_size_for_vit - ) - attention_mask = video_mask.reshape(batch_size, -1, mask_patch_elements).max(axis=-1).astype(jnp.int32) hidden_states = hidden_states.reshape( -1, self.config.num_channels_for_vit, @@ -2226,30 +2820,16 @@ def __call__( x, _ = self.patch_embed(hidden_states) x = x.reshape(batch_size, -1, self.config.hidden_size_for_vit) - if attention_mask is not None and video_grid_thw is None: - raise ValueError("video_grid_thw is required when video_mask is provided.") - pos = self.pos_embed_interpolate( - num_frames, - height, - width, - video_grid_thw=video_grid_thw, # pyrefly: ignore[bad-argument-type] - attention_mask=attention_mask, - ) + pos = self.pos_embed_interpolate(num_frames, height, width) + + pos = pos[jnp.newaxis, :, :] x = x + pos - valid_grid = video_grid_thw h_traj = [] for i in range(self.depth): block_name = f"blocks_{i}" blk = getattr(self, block_name) - x = blk( - x, - num_frames=num_frames, - height=height, - width=width, - attention_mask=attention_mask, - valid_grid=valid_grid, - ) + x = blk(x, num_frames=num_frames, height=height, width=width) h_traj.append(x) deep_feats = [] @@ -2271,7 +2851,7 @@ class Qwen3OmniMoeVisionProjector(nnx.Module): merger: Patch merger for spatial reduction """ - def __init__(self, config: Config, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, *, rngs: nnx.Rngs = None): """Initializes the Qwen3Omni vision projector. Args: @@ -2319,7 +2899,7 @@ def qwen3omni_visionprojector_as_linen(config: Config, mesh: Mesh) -> nn.Module: class Qwen3OmniAudioEncoderLayer(nnx.Module): """Transformer encoder layer for audio model.""" - def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): self.config = config self.mesh = mesh self.rngs = rngs @@ -2424,7 +3004,7 @@ class Qwen3OmniAudioEncoder(nnx.Module): mesh: Mesh, JAX device mesh (used for sharding) """ - def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): self.config = config self.mesh = mesh self.rngs = rngs @@ -2575,7 +3155,7 @@ def __call__( class Qwen3OmniAudioProjector(nnx.Module): """Projection layer that converts audio encoder output to model embedding space.""" - def __init__(self, config: Config, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] + def __init__(self, config: Config, *, rngs: nnx.Rngs = None): self.config = config self.proj1 = DenseGeneral( in_features_shape=config.d_model_for_audio, diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py index dc24527558..0f4ccb4a50 100644 --- a/src/maxtext/models/qwen3_5.py +++ b/src/maxtext/models/qwen3_5.py @@ -150,6 +150,7 @@ def __init__( # Determine the type of attention mechanism for the current layer. is_full_attention_layer = (self.layer_idx + 1) % cfg.inhomogeneous_layer_cycle_interval == 0 + self.is_full_attention_layer = is_full_attention_layer # Conditionally instantiate either the Linear Attention or Full Attention block. if is_full_attention_layer: diff --git a/src/maxtext/optimizers/optimizers.py b/src/maxtext/optimizers/optimizers.py index 67e1f589ca..ddf71507ad 100644 --- a/src/maxtext/optimizers/optimizers.py +++ b/src/maxtext/optimizers/optimizers.py @@ -18,11 +18,10 @@ import re import jax import jax.numpy as jnp - -import optax -from optax.contrib._muon import muon from maxtext.common.common_types import DecoderBlockType +from maxtext.optimizers.muon import muon from maxtext.utils.muon_utils import get_muon_weight_dimension_numbers +import optax def _get_path_mask_fn(patterns, match_returns_true=True): @@ -166,7 +165,7 @@ def skip_update(): return optax.GradientTransformationExtraArgs(init_fn, update_fn) -def get_optimizer(config, learning_rate_schedule, model=None): +def get_optimizer(config, learning_rate_schedule, model=None, mesh=None): """Create optimizer.""" if config.opt_type == "adamw": # Create AdamW Optimizer following Llama2's training details, see https://arxiv.org/pdf/2307.09288.pdf section 2.2 @@ -195,7 +194,9 @@ def get_optimizer(config, learning_rate_schedule, model=None): elif config.opt_type == "muon": # extract muon dimension number from model structure if model is not None: - muon_weight_dimension_numbers = get_muon_weight_dimension_numbers(model, config) + muon_weight_dimension_numbers = get_muon_weight_dimension_numbers( + model, config, mesh=mesh + ) else: raise ValueError("Please specify model to extract muon dimension number.") @@ -204,7 +205,9 @@ def get_optimizer(config, learning_rate_schedule, model=None): ns_steps = 10 else: ns_coeffs = (3.4445, -4.7750, 2.0315) - ns_steps = 5 + ns_steps = getattr(config, "muon_ns_steps", 5) + + use_all_to_all = getattr(config, "muon_use_all_to_all", True) muon_kwargs = { # Shared parameters: "nesterov" uses default @@ -218,6 +221,7 @@ def get_optimizer(config, learning_rate_schedule, model=None): "consistent_rms": config.muon_consistent_rms, "ns_coeffs": ns_coeffs, "ns_steps": ns_steps, + "use_all_to_all": use_all_to_all, # AdamW-specific parameters "adam_b1": config.adam_b1, "adam_b2": config.adam_b2, diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 692690db76..f3b4506994 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -404,20 +404,29 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr nested_key = ("intermediates", "decoder", "moe_layers", "moe_bias_updates") moe_bias_updates = maxtext_utils.get_nested_value(intermediate_outputs, nested_key, None) else: - # NNX intermediates are model-rooted (no "intermediates" prefix), - # so match by suffix instead. Unlike collect_intermediates_by_suffix - # we must not ravel: the decoder update is a 2-D matrix that's - # transposed and MTP update is 1-D matrix. - for path, val in jax.tree_util.tree_leaves_with_path(intermediate_outputs): - keys = tuple(k.key for k in path if hasattr(k, "key")) - if not keys or keys[-1] != "moe_bias_updates": - continue - if "decoder" in keys: - moe_bias_updates = (val,) - elif "mtp_block" in keys: - if mtp_moe_bias_updates is None: - mtp_moe_bias_updates = [] - mtp_moe_bias_updates.append(val) + # NNX intermediates are model-rooted (no "intermediates" prefix), so match by + # suffix instead. A decoder block may sow more than one moe_bias_updates leaf + # (e.g. Qwen3-Next sows one per layer position inside each scanned block, since + # its MoE gates aren't a single homogeneous scanned collection like DeepSeek's), + # so collect every match keyed by its path, joined into a single string -- the + # path is used downstream (maxtext_utils_nnx.apply_moe_bias_updates) to locate + # the matching gate.bias parameter generically. This must be a dict (path string + # -> array), not a list of (path, array) pairs: this aux value flows through + # jax.lax.scan under gradient accumulation (gradient_accumulation.py), which + # requires every leaf to be a valid JAX array -- a plain string leaf (as part of + # a path tuple) would break that, whereas dict keys are pytree structure, not + # leaves, so they pass through untouched. Unlike collect_intermediates_by_suffix + # we must not ravel: each update is a 2-D matrix transposed at the apply site. + moe_bias_updates = { + "/".join(tuple(k.key for k in path if hasattr(k, "key"))): val + for path, val in jax.tree_util.tree_leaves_with_path( + intermediate_outputs + ) + if tuple(k.key for k in path if hasattr(k, "key"))[-1:] + == ("moe_bias_updates",) + } + if not moe_bias_updates: + moe_bias_updates = None # Add the model's primary output to the intermediates dict so it can be used # by the acceptance rate calculation in eval_step. @@ -671,7 +680,7 @@ def move(path, value): state.apply_gradients(grads) new_state = state - # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family + # Apply updates for Auxiliary-Loss-Free load balancing # pylint: disable=too-many-nested-blocks if config.routed_bias and config.routed_bias_update_rate > 0.0: if getattr(config, "model_name", "").startswith("deepseek4"): @@ -701,23 +710,10 @@ def move(path, value): node.bias.value = node.bias.value + jnp.array(update_val) if getattr(config, "log_moe_bias_norms", False): bias_metrics[f"learning/moe_bias_update_norm_{name_prefix}"] = jnp.linalg.norm(jnp.array(update_val)) - else: - # 1. Update main decoder scanned MoE layers. - # The update from the scan is (num_moe_layers, num_experts) and must be transposed. - decoder_layer = getattr(new_state.model.decoder, "moe_layers", new_state.model.decoder) - decoder_bias = _find_gate_bias(decoder_layer) - if decoder_bias is not None: - decoder_bias.value = decoder_bias.value + jnp.array(moe_bias_updates[0]) - - # 2. Update auxiliary MTP MoE layers (if enabled). - # Unlike the main decoder, each MTP layer is an individual un-scanned layer - # with a 1D bias of shape (num_experts,). - if mtp_moe_bias_updates is not None and hasattr(new_state.model, "mtp_block"): - for i, update in enumerate(mtp_moe_bias_updates): - mtp_layer = getattr(new_state.model.mtp_block, f"mtp_layer_{i + 1}", None) - mtp_bias = _find_gate_bias(mtp_layer) - if mtp_bias is not None: - mtp_bias.value = mtp_bias.value + jnp.array(update) + elif moe_bias_updates is not None: + maxtext_utils_nnx.apply_moe_bias_updates( + new_state.model, moe_bias_updates + ) lm_loss = xent_sum / (total_weights + EPS) scalar_metrics = { diff --git a/src/maxtext/trainers/pre_train/train_compile.py b/src/maxtext/trainers/pre_train/train_compile.py index b705bef562..46ed4880a2 100644 --- a/src/maxtext/trainers/pre_train/train_compile.py +++ b/src/maxtext/trainers/pre_train/train_compile.py @@ -138,7 +138,9 @@ def get_shaped_inputs(topology_mesh, config): # The learning_rate_schedule is baked into the compiled object. learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) # pass in model for muon - tx = optimizers.get_optimizer(config, learning_rate_schedule, model) + tx = optimizers.get_optimizer( + config, learning_rate_schedule, model, mesh=topology_mesh + ) # Shaped RNG keys _, example_rng = jax.random.split(jax.random.PRNGKey(0), 2) diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 8069e951af..b330aca01f 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -298,6 +298,8 @@ def __init__( TypeError: If training_config is not a pyconfig.HyperParameters instance. ValueError: If training_config.model_name is not specified or empty, or if `wrap_with_tunix_adapter` is requested without `tokenizer_pad_id` or without a `mesh`. + NotImplementedError: If `training_config.lora.enable_lora` is True. This engine has no LoRA + path and would otherwise full-finetune the base model while the config claims LoRA. """ if not isinstance(training_config, pyconfig.HyperParameters): raise TypeError( @@ -311,6 +313,13 @@ def __init__( ) if mesh is None: raise ValueError("wrap_with_tunix_adapter=True requires a mesh; the adapter is built under it.") + + # This engine has no LoRA path + if getattr(getattr(training_config, "lora", None), "enable_lora", False): + raise NotImplementedError( + "MaxTextTrainingEngine does not support LoRA, but lora.enable_lora=True was set. " + "This engine trains all parameters, set lora.enable_lora=False to train with this engine." + ) self._config = training_config self._mesh = mesh self._init_rng = jax.random.PRNGKey(training_config.init_weights_seed) @@ -354,9 +363,11 @@ def __init__( self._cached_losses: list[abstract_engine.WeightedMetric | jax.Array] = [] # `create_training_optimizer` returns a raw optax GradientTransformation. `TrainStateNNX.apply_gradients` # calls `optimizer.update(model, grads)`, which is the nnx.Optimizer signature, and - # `checkpointing.CheckpointState` expects an nnx.Optimizer too, so wrap it here. This engine is only - # driven via tunix's GRPO+Raiden integration, which never enables LoRA, so `wrt` is always `nnx.Param`. - self._learning_rate_schedule, tx = train_utils.create_training_optimizer(self._config, self._model) + # `checkpointing.CheckpointState` expects an nnx.Optimizer too, so wrap it here. `wrt=nnx.Param` + # covers every parameter, which is correct only because LoRA is rejected above. + self._learning_rate_schedule, tx = train_utils.create_training_optimizer( + self._config, self._model, mesh=self._mesh + ) self._optimizer = nnx.Optimizer(self._model, tx, wrt=nnx.Param) self._train_step: int = 0 @@ -565,12 +576,13 @@ def _update_kernel(self, state_pure, accumulated_grads, micro_step_count, mean_l lambda g: g / micro_step_count, accumulated_grads, ) + grad_norm = max_utils.l2norm_pytree(grads) if self._config.gradient_clipping_threshold > 0: grads = maxtext_utils.apply_gradient_clipping(grads, None, self._config.gradient_clipping_threshold) + local_state = nnx.merge(self._state_graphdef, state_pure, copy=True) if hasattr(local_state, "apply_gradients"): if self._config.skip_step_on_spikes: - grad_norm = max_utils.l2norm_pytree(grads) local_state.apply_gradients(grads, loss=mean_loss, grad_norm=grad_norm) opt_obj = getattr(local_state, "optimizer", self._optimizer) if opt_obj is not None: diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 513594ad57..476300badb 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -768,7 +768,14 @@ def get_dense_moe_layers(config): num_moe_layers = config.num_decoder_layers // config.interleave_moe_layer_step num_dense_layers = config.num_decoder_layers - num_moe_layers return num_dense_layers, num_moe_layers - elif config.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, DecoderBlockType.DEEPSEEK4): + elif config.decoder_block in ( + DecoderBlockType.QWEN3_NEXT, + DecoderBlockType.QWEN3_5, + ): + num_dense_layers = config.first_num_dense_layers + num_moe_layers = config.num_decoder_layers - config.first_num_dense_layers + return num_dense_layers, num_moe_layers + elif config.decoder_block == DecoderBlockType.DEEPSEEK4: return 0, config.num_decoder_layers elif config.decoder_block == DecoderBlockType.DEFAULT: raise ValueError("Unsupported decoder block for dense/MoE layer calculation") @@ -1970,10 +1977,10 @@ def get_abstract_state_nnx(config, mesh, nnx_init_trainstate_fn, is_training=Tru named_sharding_state = sharding.nnx_construct_named_sharding(abs_var_state, mesh) def _to_abstract_var(a_var, s_var): - a_val = a_var.get_value() + a_val = a_var.get_value() if hasattr(a_var, "get_value") else a_var s_val = getattr(s_var, "sharding", None) if isinstance(s_var, nnx.Variable) else s_var if s_val is None and isinstance(s_var, nnx.Variable): - s_val = s_var.get_value() + s_val = s_var.get_value() if hasattr(s_var, "get_value") else s_var def _extract_primary_sharding(s): if isinstance(s, (jax.sharding.Sharding, jax.sharding.PartitionSpec)): @@ -1998,7 +2005,9 @@ def _make_abstract_leaf(leaf_a, leaf_s): else: s_tree = ( jax.tree.map(lambda _: s_val, a_val) - if isinstance(s_val, (jax.sharding.Sharding, jax.sharding.PartitionSpec)) + if isinstance( + s_val, (jax.sharding.Sharding, jax.sharding.PartitionSpec) + ) else s_val ) new_val = jax.tree.map( @@ -2007,7 +2016,9 @@ def _make_abstract_leaf(leaf_a, leaf_s): s_tree, is_leaf=lambda x: hasattr(x, "shape") and hasattr(x, "dtype"), ) - return a_var.replace(value=new_val) + if hasattr(a_var, "replace"): + return a_var.replace(value=new_val) + return new_val abstract_state = jax.tree.map( _to_abstract_var, diff --git a/src/maxtext/utils/maxtext_utils_nnx.py b/src/maxtext/utils/maxtext_utils_nnx.py index 1b67c2592c..1c50a72a0e 100644 --- a/src/maxtext/utils/maxtext_utils_nnx.py +++ b/src/maxtext/utils/maxtext_utils_nnx.py @@ -164,7 +164,7 @@ def create_nnx_sharded_model( named_sharding = nnx_extract_named_sharding(abstract_state) if mesh is None: - mesh = abstract_model.mesh # pyrefly: ignore[missing-attribute] + mesh = abstract_model.mesh # JIT a function that creates the model state with proper sharding from the start. # By providing out_shardings, we instruct JAX to produce sharded output directly, @@ -218,6 +218,37 @@ def nnx_update_sharding_meta(variable, transform_fn): return variable +def nnx_sync_moveaxis(tree, from_axis, to_axis): + """Moves an axis in both values and sharding metadata of nnx.Variables.""" + if from_axis == to_axis: + return tree + + def _op(x): + is_var = isinstance(x, nnx.Variable) + val = x.get_value() if is_var else x + if not hasattr(val, "shape"): + return x + + new_val = jnp.moveaxis(val, from_axis, to_axis) + if not is_var: + return new_val + + def move_fn(l): + while len(l) < val.ndim: + l.append(None) + if len(l) > max(from_axis, to_axis): + l.insert(to_axis, l.pop(from_axis)) + return l + + return nnx_update_sharding_meta(x.replace(value=new_val), move_fn) + + return jax.tree.map( + _op, + tree, + is_leaf=lambda x: isinstance(x, nnx.Variable) or hasattr(x, "shape"), + ) + + def nnx_remove_scan_axis(tree, name="layers"): """Removes the given scan axis from the PartitionSpec.""" @@ -284,3 +315,90 @@ def add_fn(l): return nnx_update_sharding_meta(x, add_fn) return jax.tree.map(_op, tree, is_leaf=lambda x: isinstance(x, nnx.Variable)) + + +def nnx_add_and_sync_scan_axis(tree, name="layers", scan_axis=0): + """Adds the given scan axis to PartitionSpec and moves axis if scan_axis != 0.""" + tree = nnx_add_scan_axis(tree, name, 0) + if scan_axis != 0: + new_params, new_rest = tree.split(nnx.Param, ...) + new_params = nnx_sync_moveaxis(new_params, 0, scan_axis) + tree = nnx.merge_state(new_params, new_rest) + return tree + + +def apply_moe_bias_updates( + model: nnx.Module, moe_bias_updates: dict[str, jax.Array] +) -> None: + """Applies aux-loss-free load-balancing bias updates to each MoE gate's bias, in place. + + Different decoder blocks sow `moe_bias_updates` intermediates at different + depths: DeepSeek + sows once per homogeneous scanned `moe_layers` collection, while Qwen3-Next + sows once per + layer position inside each scanned block (its MoE gates aren't a single + uniform scanned + collection). Rather than hardcoding an absolute attribute path for one decoder + block, this + locates each target `gate.bias` parameter generically: it's the unique `Param` + leaf whose + path ends in `("gate", "bias")` and lives under the same parent module that + produced the + corresponding sow, found by matching path prefixes against the live parameter + tree. + + Args: + model: The NNX model whose gate-bias parameters will be updated in place. + moe_bias_updates: A dict mapping each sow path (its components joined with + "/") to its update array, as collected in `train.py`'s `train_step` from + the model's sown `moe_bias_updates` intermediates. A dict (rather than a + list of (path, array) pairs) because this value flows through + `jax.lax.scan` under gradient accumulation, which requires every leaf to + be an array -- dict keys are pytree structure, not leaves, so the string + paths survive that unscathed. + """ + param_paths = [ + tuple(k.key for k in path if hasattr(k, "key")) + for path, _ in jax.tree_util.tree_leaves_with_path( + nnx.state(model, nnx.Param).to_pure_dict() + ) + ] + gate_bias_paths = [p for p in param_paths if p[-2:] == ("gate", "bias")] + + for sow_path_str, update in moe_bias_updates.items(): + sow_path = tuple(sow_path_str.split("/")) + parent_path = sow_path[:-1] + matches = [ + p for p in gate_bias_paths if p[: len(parent_path)] == parent_path + ] + if len(matches) != 1: + raise ValueError( + f"Expected exactly one gate.bias parameter under {parent_path} for" + f" the moe_bias_updates sown at {sow_path}, found {len(matches)}:" + f" {matches}" + ) + target_bias = model + for key in matches[0]: + target_bias = getattr(target_bias, key) + update_arr = jnp.array(update) + if target_bias.value.shape == update_arr.shape: + target_bias.value = target_bias.value + update_arr + elif target_bias.value.ndim == update_arr.ndim and sorted( + target_bias.value.shape + ) == sorted(update_arr.shape): + target_shape = target_bias.value.shape + update_shape = update_arr.shape + used = [False] * len(update_shape) + perm = [] + for s in target_shape: + for i, us in enumerate(update_shape): + if us == s and not used[i]: + perm.append(i) + used[i] = True + break + if len(perm) == len(target_shape): + target_bias.value = target_bias.value + jnp.transpose(update_arr, perm) + else: + target_bias.value = target_bias.value + update_arr + else: + target_bias.value = target_bias.value + update_arr diff --git a/src/maxtext/utils/muon_utils.py b/src/maxtext/utils/muon_utils.py index 91d7d781c9..e6f2971494 100644 --- a/src/maxtext/utils/muon_utils.py +++ b/src/maxtext/utils/muon_utils.py @@ -24,19 +24,22 @@ python3 -m maxtext.utils.muon_utils qwen3-4b True """ +import collections.abc import os import sys from typing import Optional, Tuple -import flax.linen as nn from flax import nnx +import flax.linen as nn +from flax.linen import partitioning as nn_partitioning import jax from maxtext.configs import pyconfig -from maxtext.utils.globals import MAXTEXT_PKG_DIR from maxtext.layers import quantizations from maxtext.models import models +from maxtext.optimizers.muon import MuonDimensionNumbers as mdn from maxtext.utils import maxtext_utils, model_creation_utils -from optax.contrib._muon import MuonDimensionNumbers as mdn +from maxtext.utils import sharding as sharding_lib +from maxtext.utils.globals import MAXTEXT_PKG_DIR def _is_path_contain_any(tuples, path): @@ -44,128 +47,252 @@ def _is_path_contain_any(tuples, path): return any(x in path for x in tuples) -def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]: - """ - Determines Muon dimension numbers based on the parameter's hierarchical path. - - This function defines the mapping from a parameter's logical path within the model - to its corresponding MuonDimensionNumbers (mdn). The strategy is applied in - a specific order to handle general cases and then more specific ones, allowing - for fall-through logic in nested structures. +# Parameters excluded from Muon updates (e.g. 1D norms, embeddings, scalars) +EXCLUDED_SUBSTRINGS = ( + "scale", + "embedding", + "logits_dense", + "post_beta", + "pre_beta", + "res_beta", + "hc_base", + "sinks", + "tid2eid", + "A_log", + "dt_bias", + "conv1d", + "shared_expert_gate", +) + +EXCLUDED_EXACT_SEGMENTS = { + "bias", +} + +# Attention module identifiers and tensor projection names that require +# head-aware (3D/4D) dimension specifications. +ATTENTION_BLOCK_NAMES = ( + "self_attention", + "full_attention", + "attention", + "self_attn", + "attn", + "attention_mla", + "GptOssAttention", +) + +ATTENTION_QKV_NAMES = ( + "query", + "key", + "value", + "wq_b", + "wkv_b", + "wkv", + "q_proj", + "k_proj", + "v_proj", +) + +ATTENTION_OUT_NAMES = ("out", "o_proj") + +MOE_BLOCK_NAMES = ( + "MoeBlock_0", + "moe_block", + "routed_experts", + "GptOssMlp", +) + + +def transform_logic(path: Tuple[str, ...], shape: Optional[Tuple[int, ...]] = None) -> Optional[mdn]: + """Determines Muon dimension numbers based on parameter path and shape. + + This function maps a parameter's hierarchical path within the model + to its corresponding MuonDimensionNumbers (mdn) specifying the reduction + and output axes for 2D matrix orthogonalization. + + In MaxText, layer scanning places the layer scan axis at index 1 + (`param_scan_axis = 1`), resulting in shapes: + - Standard weights / MLPs: [in_features, num_layers, out_features] + - Attention QKV: [in_features, num_layers, num_heads, head_dim] + - Attention Out: [num_heads, num_layers, head_dim, out_features] + - MoE routed experts: [num_experts, num_layers, in_features, out_features] + - Grouped linear (o_a_proj): [o_groups, num_layers, in_features_per_group, out_features_per_group] Strategy: - 1. Exclusions: Parameters not suitable for Muon (e.g., scalars, embeddings, - unembedding) are explicitly returned as `None`. - 2. Special Weights: - 2.1 MoE Block Specific Weights - 2.2 Self-Attention Specific Weights - 3. Standard Weights: Default mapping for most other 3D weight shapes. + 1. Exclusions: Non-matrix, 1D, scalar, embedding, or state-space parameters + are excluded (returns None) and optimized via AdamW. + 2. MoE Routed Experts: [num_experts, ..., in_features, out_features] + map to reduction axis (-2,) and output axis (-1,). + Note: gate.kernel is [in_features, (num_layers), num_experts] and uses standard (0,) and (-1,). + 3. Grouped Linear: [n_groups, ..., in_features_per_group, out_features_per_group] + maps to reduction axis (-2,) and output axis (-1,). + 4. Head-expanded Attention: QKV and Output projections with 3D/4D shapes + map to head-aware reduction and output axes: + - QKV: reduction (0,), output (-2, -1) + - Out (unflattened): reduction (0, -2), output (-1,) + 5. Standard Weights: Default 2D matrix mapping (0,) and (-1,) for MLPs, + GDN, shared experts, MHC alpha projections, and dense projections. Args: - path: A tuple of strings representing the hierarchical path of the parameter. + path: Tuple of strings representing the parameter's hierarchical path. + shape: Optional shape tuple of the parameter tensor. Returns: - An instance of `MuonDimensionNumbers` if a specific mapping is found, - `None` for excluded parameters, or a default `mdn` for standard weights. + An instance of `optax.contrib.MuonDimensionNumbers` if a valid mapping is + found, or `None` if the parameter is excluded from Muon updates. """ + # Exclude 1D / scalar parameters + if shape is not None and len(shape) < 2: + return None - # 1 Exclude parameters not suitable for Muon (scalar, embeddings, unembedding) - # "embedding": embedding - # "logits_dense": output embedding - # "tid2eid": lookup table in hash routing moe - # "scale": scalar, common module - # "sinks": scalar, attention sink - # "bias": scalar, common module - # "hc_base": scalar, in mhc head - # "post_beta", "pre_beta", "res_beta": scalar, in mhc - # "A_log": scalar / 1D per head, in gdn linear attention - # "conv1d": depthwise 1D convolution - # "shared_expert_gate": scalar projection (output_dim = 1) + # Exclude non-matrix parameters, embeddings, biases, and normalization if any( - any( - x in segment - for x in ( - "scale", - "embedding", - "logits_dense", - "post_beta", - "pre_beta", - "res_beta", - "hc_base", - "sinks", - "tid2eid", - "A_log", - "conv1d", - "shared_expert_gate", - ) - ) + segment in EXCLUDED_EXACT_SEGMENTS or (segment.endswith("bias") and segment != "position_bias") + or any(x in segment for x in EXCLUDED_SUBSTRINGS) for segment in path ): return None - # 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", "moe_block", "GptOssMlp"), path): - # exclude gate - if _is_path_contain_any(("wi", "wi_0", "wi_1", "wo"), path): + # MoE routed expert weights: [num_experts, (num_layers), in_features, out_features] + if _is_path_contain_any(MOE_BLOCK_NAMES, path): + if _is_path_contain_any(("wi", "wi_0", "wi_1", "wo", "gate_up_proj"), path): return mdn((-2,), (-1,)) - # 2.2 Special weights: Self attention / Attention - elif _is_path_contain_any(("self_attention", "GptOssAttention", "attention"), path): - # Attention output projection: [0, L, -2, -1] - # For standard attention (e.g. self_attention, GptOssAttention), out projection reduces over (0, -2). - # Note: Qwen3-Next full attention flattens heads into a 2D projection, so it uses standard weights mdn((0,), (-1,)). - if "out" in path and _is_path_contain_any(("self_attention", "GptOssAttention"), path): - return mdn((0, -2), (-1,)) - # Block-diagonal grouped linear layer: [n_groups, L, in_features_per_group, out_features_per_group] - elif "o_a_proj" in path: - return mdn((-2,), (-1,)) - # Attention qkv projection: [0, L, -2, -1] - # MLA, exclude wq_a / wkv_a - elif _is_path_contain_any(("query", "key", "value", "wq_b", "wkv_b", "wkv"), path): - return mdn((0,), (-2, -1)) + # Block-diagonal grouped linear layer (e.g. DeepSeek-V4 attention output projection): + # [n_groups, (num_layers), in_features_per_group, out_features_per_group] -> reduce (-2,), output (-1,) + if "o_a_proj" in path: + return mdn((-2,), (-1,)) - # 3 Standard weights, [0, L, -1] + # Head-expanded attention projections (3D unscanned or 4D scanned) + if _is_path_contain_any(ATTENTION_BLOCK_NAMES, path) and (shape is None or len(shape) > 2): + if _is_path_contain_any(ATTENTION_QKV_NAMES, path): + # [in_features, (num_layers), num_heads, head_dim] -> reduce (0,), output (-2, -1) + return mdn((0,), (-2, -1)) + if _is_path_contain_any(ATTENTION_OUT_NAMES, path): + # Standard attention out projection: [num_heads, (num_layers), head_dim, out_features] -> reduce (0, -2), output (-1,) + # Note: Qwen3-Next flattens heads into a 2D projection (in_features = num_heads * head_dim), + # so its out projection is a standard 2D matrix [in_features, (num_layers), out_features] -> reduce (0,), output (-1,) + if _is_path_contain_any(("self_attention", "GptOssAttention"), path) or (shape is not None and len(shape) == 4): + return mdn((0, -2), (-1,)) + + # Standard 2D matrix weights (dense MLPs, shared experts, GDN, MHC alpha, dense projections, router weights) + # [in_features, (num_layers), out_features] -> reduce (0,), output (-1,) return mdn((0,), (-1,)) -def get_transform_tree(tree, path=()): - """Extraction utility via recursion.""" - if isinstance(tree, dict): - return {k: get_transform_tree(v, path + (k,)) for k, v in tree.items()} +def get_transform_tree(tree, mesh=None, path=()): + """Recursively extracts `MuonDimensionNumbers` and shardings for Linen abstract parameters.""" + if isinstance(tree, (dict, collections.abc.Mapping)) or hasattr(tree, "items"): + return { + k: get_transform_tree(v, mesh=mesh, path=path + (k,)) + for k, v in tree.items() + } else: - return transform_logic(path) + val = getattr(tree, "value", tree) + val_shape = getattr(val, "shape", None) + dim_num = transform_logic(path, shape=val_shape) + if dim_num is not None: + names = getattr(tree, "names", None) + if names is not None: + # Prepend None for leading scanned layer dimensions if names length < tensor rank + if val_shape is not None and len(names) < len(val_shape): + diff = len(val_shape) - len(names) + names = (None,) * diff + names + # Resolve logical axis names to physical NamedSharding when a device mesh is available + if mesh is not None: + sharding = sharding_lib.create_sharding(mesh, names) + else: + sharding = names + else: + sharding = None + return mdn( + reduction_axis=dim_num.reduction_axis, + output_axis=dim_num.output_axis, + sharding=sharding, + ) + return None -def get_muon_weight_dimension_numbers(model, config, verbose=False): - """Extract muon dimension number from model structure.""" +def get_muon_weight_dimension_numbers( + model, config=None, mesh=None, verbose=False +): + """Extracts a matching pytree of MuonDimensionNumbers with physical shardings from a model. - if isinstance(model, nnx.Module): - _, abstract_param, _ = nnx.split(model, nnx.Param, ...) + Evaluates within an active `nn_partitioning.axis_rules` context to map logical + partition axes to physical mesh axes. Supports both NNX and Linen models. + """ + if mesh is None and config is not None and hasattr(config, "mesh_axes"): + devices_array = maxtext_utils.create_device_mesh(config) + mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) - def apply_transform_nnx(path: Tuple[jax.tree_util.KeyEntry, ...], leaf): - # Convert jax.tree_util.KeyEntry path to Tuple[str, ...] - path_strings = tuple(p.key for p in path if isinstance(p, jax.tree_util.DictKey)) - return transform_logic(path_strings) + logical_rules = ( + getattr(config, "logical_axis_rules", ()) if config is not None else () + ) + # Populate logical axis rules in Flax context for automatic NamedSharding resolution + with nn_partitioning.axis_rules(logical_rules): + if isinstance(model, nnx.Module): + # Extract abstract parameters from the NNX model hierarchy + _, abstract_param, _ = nnx.split(model, nnx.Param, ...) + + # Resolve physical NamedSharding for each parameter under the active axis rules + named_sharding_state = ( + sharding_lib.nnx_construct_named_sharding(abstract_param, mesh) + if mesh is not None + else None + ) + abstract_dict = nnx.to_pure_dict(abstract_param) + named_sharding_dict = ( + nnx.to_pure_dict(named_sharding_state) + if named_sharding_state is not None + else abstract_dict + ) - # NNX abstract_param is an nnx.State (not Linen's dict of LogicallyPartitioned leaves); - # tree_map_with_path round-trips that structure so each Param.value holds the mdn result. - muon_weight_dimension_numbers = jax.tree_util.tree_map_with_path( - apply_transform_nnx, nnx.to_pure_dict(abstract_param) - ) - muon_weight_dimension_numbers = nnx.State(muon_weight_dimension_numbers) + def apply_transform_nnx( + path: Tuple[jax.tree_util.KeyEntry, ...], leaf, abs_leaf + ): + path_strings = tuple( + p.key for p in path if isinstance(p, jax.tree_util.DictKey) + ) + abs_val = ( + abs_leaf.get_value() if hasattr(abs_leaf, "get_value") else abs_leaf + ) + val_shape = getattr(abs_val, "shape", None) + dim_num = transform_logic(path_strings, shape=val_shape) + if dim_num is not None: + val = leaf.get_value() if hasattr(leaf, "get_value") else leaf + sharding = ( + val + if isinstance( + val, (jax.sharding.NamedSharding, jax.sharding.PartitionSpec) + ) + else getattr(leaf, "sharding", None) + ) + if isinstance(sharding, jax.ShapeDtypeStruct) or not isinstance( + sharding, (jax.sharding.NamedSharding, jax.sharding.PartitionSpec) + ): + sharding = getattr(leaf, "sharding", None) + return mdn( + reduction_axis=dim_num.reduction_axis, + output_axis=dim_num.output_axis, + sharding=sharding, + ) + return None - else: # Linen - # quickly get param structure without materialization - abstract_param = maxtext_utils.get_abstract_param(model, config) - # get muon dimension number from param - muon_weight_dimension_numbers = get_transform_tree(abstract_param) + # Walk the parameter tree to produce a matching nnx.State of MuonDimensionNumbers + muon_weight_dimension_numbers = jax.tree_util.tree_map_with_path( + apply_transform_nnx, named_sharding_dict, abstract_dict + ) + muon_weight_dimension_numbers = nnx.State(muon_weight_dimension_numbers) - if verbose: - _print_structure_debug(abstract_param, muon_weight_dimension_numbers) - return muon_weight_dimension_numbers + else: # Linen + abstract_param = maxtext_utils.get_abstract_param(model, config) + muon_weight_dimension_numbers = get_transform_tree( + abstract_param, mesh=mesh + ) + + if verbose: + _print_structure_debug(abstract_param, muon_weight_dimension_numbers) + return muon_weight_dimension_numbers def _print_structure_debug(abstract_param, muon_weight_dimension_numbers): @@ -237,7 +364,9 @@ def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False): else: model = models.transformer_as_linen(config, mesh=mesh, quant=quant) # Get dimension number - muon_weight_dimension_numbers = get_muon_weight_dimension_numbers(model, config, verbose=verbose) + muon_weight_dimension_numbers = get_muon_weight_dimension_numbers( + model, config, mesh=mesh, verbose=verbose + ) if pure_nnx: muon_weight_dimension_numbers = {"params": nnx.to_pure_dict(muon_weight_dimension_numbers)} return muon_weight_dimension_numbers diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index 41f74a21b6..768a6ce824 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -43,18 +43,22 @@ from maxtext.utils.rampup_batch import create_rampup_manager -def create_training_optimizer(config, model): +def create_training_optimizer(config, model, mesh=None): """Creates the optimizer and learning rate schedule.""" learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) # pass in model for muon - tx = optimizers.get_optimizer(config, learning_rate_schedule, model) + tx = optimizers.get_optimizer( + config, learning_rate_schedule, model, mesh=mesh + ) return learning_rate_schedule, tx def create_checkpoint_manager(config, mesh, init_state_fn): """Creates the init_rng, optimizer, learning rate schedule, and checkpoint manager.""" # pass in model for muon - logger = checkpointing.setup_checkpoint_logger(config) + # `setup_checkpoint_logger` only emits a deprecation warning now (Orbax v1 logs + # internally) and always returns None; we still pass it through for API parity. + logger = checkpointing.setup_checkpoint_logger(config) # pylint: disable=assignment-from-no-return if config.enable_multi_tier_checkpointing: checkpoint_manager = emergency_checkpointing.create_replicator_checkpoint_manager( config.local_checkpoint_directory, @@ -98,6 +102,7 @@ def create_checkpoint_manager(config, mesh, init_state_fn): config.enable_autocheckpoint, config.checkpoint_todelete_subdir, config.checkpoint_todelete_full_path, + config.checkpoint_storage_target_data_file_size_bytes, ) # Use Colocated Python checkpointing dispatchers optimization (Single Controller only). @@ -261,7 +266,9 @@ def setup_train_loop(config, recorder, devices=None): _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices) else: model = model_creation_utils.from_config(config, devices) - learning_rate_schedule, tx = create_training_optimizer(config, model) + learning_rate_schedule, tx = create_training_optimizer( + config, model, mesh=mesh + ) if config.pure_nnx: # For NNX, the train state is wrapped in the TrainStateNNX module. diff --git a/tests/unit/gdn_cp_test.py b/tests/unit/gdn_cp_test.py new file mode 100644 index 0000000000..a95cf5c851 --- /dev/null +++ b/tests/unit/gdn_cp_test.py @@ -0,0 +1,289 @@ +# 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 +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Checks the GatedDeltaNet cross-device state composition on a CPU mesh. + +Runs as a subprocess so the forced device count takes effect before JAX +initializes; the parent pytest process has already initialized JAX with the +default device count. + +The child checks three things, each on a 1-D context mesh and a 2-D fsdp x +context mesh, and the first also with a tuple axis name, which is what +`qwen3.py` passes when more than one context knob is set. + + - Forward. `gdn_cp.incoming_state` against a sequential single-device + composition of the same affine maps. + - Backward. `gdn_cp` defines no custom_vjp, so the backward pass is whatever + JAX's transpose rules for `lax.scan`, `lax.ppermute` and `lax.psum` compose + into. A `ppermute` transposes to the inverse permutation, so a cotangent + has to travel back down the rank chain for an early shard to see any + gradient from a late one. + - End to end. `jax_chunk_gated_delta_rule` itself, sharded against unsharded, + which catches a wrong chunk layout or a wrong replay. + +Every one of these fails silently: the loss still falls while the state +entering each shard is wrong. So each check has a negative control that +confirms the comparison can fail. +""" + +import os +import subprocess +import sys +from functools import partial + +import jax +import jax.numpy as jnp +from jax.sharding import Mesh, PartitionSpec as P +import numpy as np +import pytest + +from maxtext.kernels.attention import gdn_cp + + +@pytest.mark.cpu_only +def test_gdn_cp_prefix_scan_matches_sequential_reference_on_cpu_mesh(): + env = os.environ.copy() + env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + " --xla_force_host_platform_device_count=8" + env["JAX_PLATFORMS"] = "cpu" + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert "GDN_CP_CHECKS_PASSED" in result.stdout + + +def _chunk_inputs(num_chunks, lead, chunk_len, k_dim, v_dim): + """Random w, u, k and g in the layout compose_local scans over.""" + keys = jax.random.split(jax.random.PRNGKey(1), 4) + w = jax.random.normal(keys[0], (num_chunks,) + lead + (chunk_len, k_dim), jnp.float32) * 0.1 + u = jax.random.normal(keys[1], (num_chunks,) + lead + (chunk_len, v_dim), jnp.float32) * 0.1 + k = jax.random.normal(keys[2], (num_chunks,) + lead + (chunk_len, k_dim), jnp.float32) * 0.1 + # Log-decay gates are negative and increase along the chunk, as in Qwen3-Next. + ramp = jnp.arange(chunk_len, dtype=jnp.float32) / chunk_len + g = -jax.nn.softplus(jax.random.normal(keys[3], (num_chunks,) + lead + (chunk_len,), jnp.float32)) * (1.0 - ramp) + return w, u, k, g + + +def _direct_recurrence(w, u, k, g, h): + """Applies the inter-chunk step chunk by chunk, in its original form. + + This reference never builds A or B and never calls compose, so it checks the + affine reformulation the whole module rests on rather than restating it. + """ + for i in range(w.shape[0]): + w_c, u_c, k_c, g_c = w[i], u[i], k[i], g[i] + g_last = g_c[..., -1] + k_g_T = (k_c * jnp.exp(g_last[..., None] - g_c)[..., None]).swapaxes(-1, -2) + h = jnp.exp(g_last)[..., None, None] * h + jnp.matmul(k_g_T, u_c - jnp.matmul(w_c, h)) + return h + + +def test_compose_local_folds_the_chunks_into_one_affine_map(): + """A_loc @ h + B_loc equals running the recurrence chunk by chunk.""" + lead, k_dim, v_dim = (1, 2), 8, 8 + w, u, k, g = _chunk_inputs(num_chunks=6, lead=lead, chunk_len=4, k_dim=k_dim, v_dim=v_dim) + h_init = jax.random.normal(jax.random.PRNGKey(2), lead + (k_dim, v_dim), jnp.float32) * 0.1 + + affine_a, affine_b = jax.jit(gdn_cp.compose_local)(w, u, k, g) + folded = jnp.matmul(affine_a, h_init) + affine_b + np.testing.assert_allclose(jax.device_get(folded), jax.device_get(_direct_recurrence(w, u, k, g, h_init)), atol=1e-5) + + # Negative control: reversing the chunks changes the answer, so the test sees order. + reversed_chunks = _direct_recurrence(w[::-1], u[::-1], k[::-1], g[::-1], h_init) + assert float(jnp.max(jnp.abs(folded - reversed_chunks))) > 1e-4 + + +def test_incoming_state_is_the_identity_on_a_single_device_mesh(): + """With one shard there is nothing before it, so it enters with h_init.""" + mesh = Mesh(np.array(jax.devices()[:1]), ("context",)) + batch, heads, k_dim, v_dim = 1, 2, 8, 8 + keys = jax.random.split(jax.random.PRNGKey(3), 3) + affine_a = jax.random.normal(keys[0], (1, batch, heads, k_dim, k_dim), jnp.float32) * (0.3 / k_dim**0.5) + affine_b = jax.random.normal(keys[1], (1, batch, heads, k_dim, v_dim), jnp.float32) * 0.1 + h_init = jax.random.normal(keys[2], (batch, heads, k_dim, v_dim), jnp.float32) * 0.1 + + mapped = partial( + jax.shard_map, + mesh=mesh, + in_specs=(P("context", None, None, None, None), P("context", None, None, None, None), P(None, None, None, None)), + out_specs=(P("context", None, None, None), P(None, None, None, None)), + check_vma=False, + ) + incoming, final = jax.jit(mapped(lambda a, b, h: gdn_cp.incoming_state(a[0], b[0], h, "context")))( + affine_a, affine_b, h_init + ) + np.testing.assert_allclose(jax.device_get(incoming.reshape(h_init.shape)), jax.device_get(h_init), atol=1e-6) + expected_final = jnp.matmul(affine_a[0], h_init) + affine_b[0] + np.testing.assert_allclose(jax.device_get(final), jax.device_get(expected_final), atol=1e-5) + + +def _sequential_reference(affine_a, affine_b, h_init): + """Device i enters with every pair before it composed and applied to h_init.""" + incoming = [] + a_run = jnp.broadcast_to(jnp.eye(affine_a.shape[-1], dtype=affine_a.dtype), affine_a.shape[1:]) + b_run = jnp.zeros(affine_b.shape[1:], affine_b.dtype) + for i in range(affine_a.shape[0]): + incoming.append(jnp.matmul(a_run, h_init) + b_run) + a_run, b_run = gdn_cp.compose((a_run, b_run), (affine_a[i], affine_b[i])) + return jnp.stack(incoming), jnp.matmul(a_run, h_init) + b_run + + +def _run_composition_checks(mesh, cp_axis): + """Runs the prefix scan and the negative control on one mesh.""" + axis_names = (cp_axis,) if isinstance(cp_axis, str) else cp_axis + num_shards = int(np.prod([mesh.shape[name] for name in axis_names])) + batch, heads, k_dim, v_dim = 1, 2, 8, 8 + keys = jax.random.split(jax.random.PRNGKey(0), 3) + # A contractive A keeps the composition over every shard finite. + affine_a = jax.random.normal(keys[0], (num_shards, batch, heads, k_dim, k_dim), jnp.float32) * (0.3 / k_dim**0.5) + affine_b = jax.random.normal(keys[1], (num_shards, batch, heads, k_dim, v_dim), jnp.float32) * 0.1 + h_init = jax.random.normal(keys[2], (batch, heads, k_dim, v_dim), jnp.float32) * 0.1 + + def sharded(fn): + mapped = partial( + jax.shard_map, + mesh=mesh, + in_specs=(P(cp_axis, None, None, None, None), P(cp_axis, None, None, None, None), P(None, None, None, None)), + out_specs=(P(cp_axis, None, None, None), P(None, None, None, None)), + check_vma=False, + ) + return jax.jit(mapped(lambda a, b, h: fn(a[0], b[0], h, cp_axis)))(affine_a, affine_b, h_init) + + reference_incoming, reference_final = _sequential_reference(affine_a, affine_b, h_init) + # shard_map concatenates each shard's (batch, heads, k_dim, v_dim) along the + # batch axis, so at batch 1 the leading axis of the result is the shard index. + incoming, final = sharded(gdn_cp.incoming_state) + incoming = incoming.reshape(reference_incoming.shape) + np.testing.assert_allclose(jax.device_get(incoming), jax.device_get(reference_incoming), atol=1e-5) + np.testing.assert_allclose(jax.device_get(final), jax.device_get(reference_final), atol=1e-5) + + # Negative control: break the chain so every shard starts from h_init. + broken, _ = sharded(lambda a, b, h, axis: (h, h)) + assert float(jnp.max(jnp.abs(broken.reshape(reference_incoming.shape) - reference_incoming))) > 1e-4 + + +def _run_gradient_checks(mesh, cp_axis): + """The backward pass, which the module never writes by hand. + + `gdn_cp` defines no custom_vjp, so the backward pass is whatever JAX's own + transpose rules for `lax.scan`, `lax.ppermute` and `lax.psum` compose into. + That is worth a test rather than an assumption: a `ppermute` transposes to the + inverse permutation, so a cotangent has to travel back down the rank chain for + an early shard to see any gradient from a late one at all. Getting it wrong is + silent in the same way a wrong forward prefix is. + """ + axis_names = (cp_axis,) if isinstance(cp_axis, str) else cp_axis + num_shards = int(np.prod([mesh.shape[name] for name in axis_names])) + batch, heads, k_dim, v_dim = 1, 2, 8, 8 + keys = jax.random.split(jax.random.PRNGKey(4), 3) + affine_a = jax.random.normal(keys[0], (num_shards, batch, heads, k_dim, k_dim), jnp.float32) * (0.3 / k_dim**0.5) + affine_b = jax.random.normal(keys[1], (num_shards, batch, heads, k_dim, v_dim), jnp.float32) * 0.1 + h_init = jax.random.normal(keys[2], (batch, heads, k_dim, v_dim), jnp.float32) * 0.1 + + def cp_loss(a, b, h): + mapped = jax.shard_map( + lambda a_s, b_s, h_s: gdn_cp.incoming_state(a_s[0], b_s[0], h_s, cp_axis), + mesh=mesh, + in_specs=(P(cp_axis, None, None, None, None), P(cp_axis, None, None, None, None), P(None, None, None, None)), + out_specs=(P(cp_axis, None, None, None), P(None, None, None, None)), + check_vma=False, + ) + incoming, final = mapped(a, b, h) + return jnp.sum(incoming**2) + jnp.sum(final**2) + + def reference_loss(a, b, h): + incoming, final = _sequential_reference(a, b, h) + return jnp.sum(incoming**2) + jnp.sum(final**2) + + grad_cp = jax.jit(jax.grad(cp_loss, argnums=(0, 1, 2)))(affine_a, affine_b, h_init) + grad_reference = jax.jit(jax.grad(reference_loss, argnums=(0, 1, 2)))(affine_a, affine_b, h_init) + for actual, expected in zip(grad_cp, grad_reference): + np.testing.assert_allclose(jax.device_get(actual), jax.device_get(expected), atol=1e-5, rtol=1e-4) + + # Every shard sees gradient. The last shard's pair reaches the first shard's + # cotangent only through the reverse ppermute, so a zero here means the + # cross-device backward is not connected. + per_shard = jnp.linalg.norm(grad_cp[0].reshape(num_shards, -1), axis=1) + assert float(jnp.min(per_shard)) > 1e-6, per_shard + + +def _run_end_to_end_checks(mesh, cp_axis): + """`jax_chunk_gated_delta_rule` itself, sharded against unsharded. + + The composition tests above check gdn_cp in isolation. This one runs the real + GatedDeltaNet function over a sharded sequence and compares every output + token, and the final state, against the same function on one device. It is + the check that catches a wrong chunk layout or a wrong replay, neither of + which the isolated tests can see. + """ + # pylint: disable=import-outside-toplevel + from maxtext.models.qwen3 import jax_chunk_gated_delta_rule + + num_shards = int(np.prod([mesh.shape[name] for name in ((cp_axis,) if isinstance(cp_axis, str) else cp_axis)])) + batch, heads, k_dim, v_dim, chunk = 1, 2, 16, 16, 8 + seq = num_shards * chunk * 2 # two local chunks per shard + keys = jax.random.split(jax.random.PRNGKey(5), 6) + query = jax.random.normal(keys[0], (batch, seq, heads, k_dim), jnp.float32) * 0.1 + key_t = jax.random.normal(keys[1], (batch, seq, heads, k_dim), jnp.float32) * 0.1 + value = jax.random.normal(keys[2], (batch, seq, heads, v_dim), jnp.float32) * 0.1 + g = -jax.nn.softplus(jax.random.normal(keys[3], (batch, seq, heads), jnp.float32)) + beta = jax.nn.sigmoid(jax.random.normal(keys[4], (batch, seq, heads), jnp.float32)) + h_init = jax.random.normal(keys[5], (batch, heads, k_dim, v_dim), jnp.float32) * 0.1 + + def call(cp): + return jax_chunk_gated_delta_rule( + query, key_t, value, g, beta, chunk_size=chunk, initial_state=h_init, cp_axis=cp, compute_dtype=jnp.float32 + ) + + reference_out, reference_state = jax.jit(lambda: call(None))() + + qkv_spec = P(None, cp_axis, None, None) + g_spec = P(None, cp_axis, None) + state_spec = P(None, None, None, None) + + def sharded_call(q, k, v, gg, bb, hh): + return jax_chunk_gated_delta_rule( + q, k, v, gg, bb, chunk_size=chunk, initial_state=hh, cp_axis=cp_axis, compute_dtype=jnp.float32 + ) + + cp_out, cp_state = jax.jit( + jax.shard_map( + sharded_call, + mesh=mesh, + in_specs=(qkv_spec, qkv_spec, qkv_spec, g_spec, g_spec, state_spec), + out_specs=(qkv_spec, state_spec), + check_vma=False, + ) + )(query, key_t, value, g, beta, h_init) + + np.testing.assert_allclose(jax.device_get(cp_out), jax.device_get(reference_out), atol=2e-4, rtol=2e-3) + np.testing.assert_allclose(jax.device_get(cp_state), jax.device_get(reference_state), atol=2e-4, rtol=2e-3) + + # Negative control: shift the sequence by one shard. If the comparison were + # insensitive to which tokens land on which device, this would still pass. + rolled = jnp.roll(reference_out, shift=chunk, axis=1) + assert float(jnp.max(jnp.abs(rolled - reference_out))) > 1e-4 + + +if __name__ == "__main__": + _devices = np.array(jax.devices()) + assert len(_devices) == 8, jax.devices() + _run_composition_checks(Mesh(_devices, ("context",)), "context") + _run_composition_checks(Mesh(_devices.reshape(2, 4), ("fsdp", "context")), "context") + # qwen3.py builds cp_axis as a tuple, so both context knobs can be live at once. + _run_composition_checks( + Mesh(_devices.reshape(2, 4), ("context", "context_usp_ulysses")), + ("context", "context_usp_ulysses"), + ) + _run_gradient_checks(Mesh(_devices, ("context",)), "context") + _run_gradient_checks(Mesh(_devices.reshape(2, 4), ("fsdp", "context")), "context") + _run_end_to_end_checks(Mesh(_devices, ("context",)), "context") + print("GDN_CP_CHECKS_PASSED") diff --git a/tests/unit/mhc_test.py b/tests/unit/mhc_test.py index bcb498e144..d7faad78fa 100644 --- a/tests/unit/mhc_test.py +++ b/tests/unit/mhc_test.py @@ -103,6 +103,7 @@ def _setup_mhc( use_mhc_pallas_kernel=False, mhc_pallas_kernel_fwd_block_size=None, mhc_pallas_kernel_bwd_block_size=None, + mhc_pallas_kernel_bwd_feature_block_size=None, dim=16, sequence_length=7, per_device_batch_size=None, @@ -138,9 +139,17 @@ def _setup_mhc( "engram_layers": [], } if mhc_pallas_kernel_fwd_block_size is not None: - kwargs["mhc_pallas_kernel_fwd_block_size"] = mhc_pallas_kernel_fwd_block_size + kwargs["mhc_pallas_kernel_fwd_block_size"] = ( + mhc_pallas_kernel_fwd_block_size + ) if mhc_pallas_kernel_bwd_block_size is not None: - kwargs["mhc_pallas_kernel_bwd_block_size"] = mhc_pallas_kernel_bwd_block_size + kwargs["mhc_pallas_kernel_bwd_block_size"] = ( + mhc_pallas_kernel_bwd_block_size + ) + if mhc_pallas_kernel_bwd_feature_block_size is not None: + kwargs["mhc_pallas_kernel_bwd_feature_block_size"] = ( + mhc_pallas_kernel_bwd_feature_block_size + ) if dtype is not None: kwargs["dtype"] = dtype kwargs["weight_dtype"] = dtype @@ -179,7 +188,9 @@ def test_moe_layer_output_shape(self, rate): self._setup_mhc(rate) with nn_partitioning.axis_rules(self.config.logical_axis_rules): - module = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + module = mhc.ManifoldConstrainedHyperConnections( + self.config, self.dim, self.mesh, self.rngs + ) layer = moe.RoutedMoE( config=self.config, num_experts=self.config.num_experts, @@ -194,7 +205,9 @@ def test_moe_layer_output_shape(self, rate): ) b, s, k, d = self.x.shape - output, metadata = module(self.pre_norm, layer, x=self.x, mhc_type=HyperConnectionType.MLP_MOE) + output, metadata = module( + self.pre_norm, layer, x=self.x, mhc_type=HyperConnectionType.MLP_MOE + ) # metadata includes load_balance_loss & moe_bias_updates self.assertLen(metadata, 2) for key, value in metadata.items(): @@ -205,7 +218,9 @@ def test_moe_layer_output_shape(self, rate): def test_dense_layer_output_shape(self, rate): self._setup_mhc(rate) with nn_partitioning.axis_rules(self.config.logical_axis_rules): - module = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + module = mhc.ManifoldConstrainedHyperConnections( + self.config, self.dim, self.mesh, self.rngs + ) layer = linears.MlpBlock( config=self.config, mesh=self.mesh, @@ -220,7 +235,9 @@ def test_dense_layer_output_shape(self, rate): ) b, s, k, d = self.x.shape - output, metadata = module(self.pre_norm, layer, x=self.x, mhc_type=HyperConnectionType.MLP_DENSE) + output, metadata = module( + self.pre_norm, layer, x=self.x, mhc_type=HyperConnectionType.MLP_DENSE + ) self.assertDictEqual(metadata, {}) self.assertEqual(output.shape, (b, s, k, d)) @@ -233,7 +250,9 @@ def test_attention_layer_output_shape(self, rate): self.config.emb_dim, ) with nn_partitioning.axis_rules(self.config.logical_axis_rules): - module = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + module = mhc.ManifoldConstrainedHyperConnections( + self.config, self.dim, self.mesh, self.rngs + ) layer = attention_mla.MLA( config=self.config, num_query_heads=self.config.num_query_heads, @@ -280,7 +299,9 @@ def test_mhc_lite_doubly_stochastic(self): """Verify that mHC-lite output is doubly stochastic (rows/cols sum to 1).""" self._setup_mhc(4, enable_mhc_lite=True) with nn_partitioning.axis_rules(self.config.logical_axis_rules): - module = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + module = mhc.ManifoldConstrainedHyperConnections( + self.config, self.dim, self.mesh, self.rngs + ) b, s, k, d = self.x.shape @@ -290,7 +311,9 @@ def test_mhc_lite_doubly_stochastic(self): # Output from mHC-lite mapping (using post_matmul API) res_alpha = jnp.asarray(module.res_alpha[...], module.dtype) - h_res = jnp.einsum("bsm,mn -> bsn", norm_x, res_alpha, precision=module.matmul_precision) + h_res = jnp.einsum( + "bsm,mn -> bsn", norm_x, res_alpha, precision=module.matmul_precision + ) res_mapping_out = module.res_mapping(h_res) row_sums = jnp.sum(res_mapping_out, axis=-1) @@ -319,7 +342,9 @@ def test_weight_concatenation_equivalence(self): """Verify that fused projection matches sequential projections.""" self._setup_mhc(4) with nn_partitioning.axis_rules(self.config.logical_axis_rules): - module = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + module = mhc.ManifoldConstrainedHyperConnections( + self.config, self.dim, self.mesh, self.rngs + ) b, s, k, d = self.x.shape x_flat = jnp.reshape(self.x, (b, s, k * d)) @@ -330,13 +355,26 @@ def test_weight_concatenation_equivalence(self): post_alpha = jnp.asarray(module.post_alpha[...], module.dtype) res_alpha = jnp.asarray(module.res_alpha[...], module.dtype) - h_pre_seq = jnp.einsum("bsm,mk -> bsk", norm_x, pre_alpha, precision=module.matmul_precision) - h_post_seq = jnp.einsum("bsm,mk -> bsk", norm_x, post_alpha, precision=module.matmul_precision) - h_res_seq = jnp.einsum("bsm,mn -> bsn", norm_x, res_alpha, precision=module.matmul_precision) + h_pre_seq = jnp.einsum( + "bsm,mk -> bsk", norm_x, pre_alpha, precision=module.matmul_precision + ) + h_post_seq = jnp.einsum( + "bsm,mk -> bsk", norm_x, post_alpha, precision=module.matmul_precision + ) + h_res_seq = jnp.einsum( + "bsm,mn -> bsn", norm_x, res_alpha, precision=module.matmul_precision + ) # Fused Projection (New way) - alpha_concat = jnp.concatenate([pre_alpha, post_alpha, res_alpha], axis=-1) - h_concat = jnp.einsum("bsm,mn -> bsn", norm_x, alpha_concat, precision=module.matmul_precision) + alpha_concat = jnp.concatenate( + [pre_alpha, post_alpha, res_alpha], axis=-1 + ) + h_concat = jnp.einsum( + "bsm,mn -> bsn", + norm_x, + alpha_concat, + precision=module.matmul_precision, + ) h_pre_fused = h_concat[..., : module.k] h_post_fused = h_concat[..., module.k : 2 * module.k] @@ -380,7 +418,9 @@ def test_feature_flag_gates_lite(self): self.rngs = nnx.Rngs(params=jax.random.key(0), dropout=jax.random.key(42)) with nn_partitioning.axis_rules(self.config.logical_axis_rules): - module = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + module = mhc.ManifoldConstrainedHyperConnections( + self.config, self.dim, self.mesh, self.rngs + ) # Shape of res_alpha should be (4*16, 4*4) = (64, 16) instead of (64, 24) self.assertEqual(module.res_alpha.shape, (64, 16)) @@ -405,7 +445,9 @@ def test_use_mhc_pallas_kernel_dispatch(self, use_mhc_pallas_kernel): dtype="bfloat16", ) with nn_partitioning.axis_rules(self.config.logical_axis_rules): - module = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + module = mhc.ManifoldConstrainedHyperConnections( + self.config, self.dim, self.mesh, self.rngs + ) layer = linears.MlpBlock( config=self.config, mesh=self.mesh, @@ -435,8 +477,12 @@ def fake_post(*args, **kwargs): return real_post(*args, **kwargs) with ( - mock.patch.object(mhc.mhc_kernel, "pre", side_effect=fake_pre) as mock_pre, - mock.patch.object(mhc.mhc_kernel, "post", side_effect=fake_post) as mock_post, + mock.patch.object( + mhc.mhc_kernel, "pre", side_effect=fake_pre + ) as mock_pre, + mock.patch.object( + mhc.mhc_kernel, "post", side_effect=fake_post + ) as mock_post, ): output, _ = module( self.pre_norm, @@ -453,12 +499,14 @@ def fake_post(*args, **kwargs): self.assertIsNotNone(config_pre) self.assertEqual(config_pre.block_size, 256) self.assertEqual(config_pre.bwd_block_size, 128) + self.assertEqual(config_pre.bwd_feature_block_size, 1024) _, kwargs_post = mock_post.call_args config_post = kwargs_post.get("config") self.assertIsNotNone(config_post) self.assertEqual(config_post.block_size, 256) self.assertEqual(config_post.bwd_block_size, 128) + self.assertEqual(config_post.bwd_feature_block_size, 1024) else: mock_pre.assert_not_called() mock_post.assert_not_called() @@ -473,13 +521,16 @@ def test_use_mhc_pallas_kernel_custom_block_size(self): use_mhc_pallas_kernel=True, mhc_pallas_kernel_fwd_block_size=128, mhc_pallas_kernel_bwd_block_size=64, + mhc_pallas_kernel_bwd_feature_block_size=512, dim=128, sequence_length=128, per_device_batch_size=1, dtype="bfloat16", ) with nn_partitioning.axis_rules(self.config.logical_axis_rules): - module = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + module = mhc.ManifoldConstrainedHyperConnections( + self.config, self.dim, self.mesh, self.rngs + ) layer = linears.MlpBlock( config=self.config, mesh=self.mesh, @@ -509,8 +560,12 @@ def fake_post(*args, **kwargs): return real_post(*args, **kwargs) with ( - mock.patch.object(mhc.mhc_kernel, "pre", side_effect=fake_pre) as mock_pre, - mock.patch.object(mhc.mhc_kernel, "post", side_effect=fake_post) as mock_post, + mock.patch.object( + mhc.mhc_kernel, "pre", side_effect=fake_pre + ) as mock_pre, + mock.patch.object( + mhc.mhc_kernel, "post", side_effect=fake_post + ) as mock_post, ): output, _ = module( self.pre_norm, @@ -526,12 +581,14 @@ def fake_post(*args, **kwargs): self.assertIsNotNone(config_pre) self.assertEqual(config_pre.block_size, 128) self.assertEqual(config_pre.bwd_block_size, 64) + self.assertEqual(config_pre.bwd_feature_block_size, 512) _, kwargs_post = mock_post.call_args config_post = kwargs_post.get("config") self.assertIsNotNone(config_post) self.assertEqual(config_post.block_size, 128) self.assertEqual(config_post.bwd_block_size, 64) + self.assertEqual(config_post.bwd_feature_block_size, 512) self.assertEqual(output.shape, self.x.shape) @@ -555,7 +612,9 @@ def test_use_mhc_pallas_kernel_requires_enable_mhc_lite(self): mock_config.matmul_precision = "default" mock_config.mhc_expansion_rate = 4 with self.assertRaises(ValueError): - mhc.ManifoldConstrainedHyperConnections(mock_config, 16, self.mesh, self.rngs) + mhc.ManifoldConstrainedHyperConnections( + mock_config, 16, self.mesh, self.rngs + ) def test_layer_vjp_parity_kernel_vs_baseline(self): """Verify that ManifoldConstrainedHyperConnections with kernel matches baseline under VJP.""" @@ -569,13 +628,17 @@ def test_layer_vjp_parity_kernel_vs_baseline(self): dtype="bfloat16", ) with nn_partitioning.axis_rules(self.config.logical_axis_rules): - module_baseline = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + module_baseline = mhc.ManifoldConstrainedHyperConnections( + self.config, self.dim, self.mesh, self.rngs + ) def layer_fn(inputs): return inputs * 2.0 def forward_baseline(x): - out, _ = module_baseline(self.pre_norm, layer_fn, x=x, mhc_type=HyperConnectionType.MLP_DENSE) + out, _ = module_baseline( + self.pre_norm, layer_fn, x=x, mhc_type=HyperConnectionType.MLP_DENSE + ) return out out_base, vjp_base = jax.vjp(forward_baseline, self.x) @@ -588,12 +651,15 @@ def forward_baseline(x): use_mhc_pallas_kernel=True, mhc_pallas_kernel_fwd_block_size=128, mhc_pallas_kernel_bwd_block_size=64, + mhc_pallas_kernel_bwd_feature_block_size=256, dim=128, sequence_length=128, per_device_batch_size=1, dtype="bfloat16", ) - module_kernel = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + module_kernel = mhc.ManifoldConstrainedHyperConnections( + self.config, self.dim, self.mesh, self.rngs + ) real_pre = mhc.mhc_kernel.pre real_post = mhc.mhc_kernel.post @@ -616,7 +682,12 @@ def fake_post(*args, **kwargs): ): def forward_kernel(x): - out, _ = module_kernel(self.pre_norm, layer_fn, x=x, mhc_type=HyperConnectionType.MLP_DENSE) + out, _ = module_kernel( + self.pre_norm, + layer_fn, + x=x, + mhc_type=HyperConnectionType.MLP_DENSE, + ) return out out_kern, vjp_kern = jax.vjp(forward_kernel, self.x) @@ -641,20 +712,41 @@ def _make_kernel_inputs(batch=2, sequence=64, streams=4, embedding=256, seed=0): """Generates synthetic inputs and parameters for testing mHC kernels.""" key = jax.random.PRNGKey(seed) keys = jax.random.split(key, 12) - x = jax.random.normal(keys[0], (batch, sequence, streams, embedding), dtype=jnp.bfloat16) - norm_scale = jax.random.normal(keys[1], (streams * embedding,), dtype=jnp.bfloat16) - pre_alpha = jax.random.normal(keys[2], (streams * embedding, streams), dtype=jnp.bfloat16) * 0.1 + x = jax.random.normal( + keys[0], (batch, sequence, streams, embedding), dtype=jnp.bfloat16 + ) + norm_scale = jax.random.normal( + keys[1], (streams * embedding,), dtype=jnp.bfloat16 + ) + pre_alpha = ( + jax.random.normal( + keys[2], (streams * embedding, streams), dtype=jnp.bfloat16 + ) + * 0.1 + ) pre_bias = jax.random.normal(keys[3], (streams,), dtype=jnp.bfloat16) * 0.1 pre_scale = jnp.array([1.0], dtype=jnp.bfloat16) - post_alpha = jax.random.normal(keys[5], (streams * embedding, streams), dtype=jnp.bfloat16) * 0.1 + post_alpha = ( + jax.random.normal( + keys[5], (streams * embedding, streams), dtype=jnp.bfloat16 + ) + * 0.1 + ) post_bias = jax.random.normal(keys[6], (streams,), dtype=jnp.bfloat16) * 0.1 post_scale = jnp.array([1.0], dtype=jnp.bfloat16) num_perms = math.factorial(streams) - res_alpha = jax.random.normal(keys[8], (streams * embedding, num_perms), dtype=jnp.bfloat16) * 0.1 + res_alpha = ( + jax.random.normal( + keys[8], (streams * embedding, num_perms), dtype=jnp.bfloat16 + ) + * 0.1 + ) res_bias = jax.random.normal(keys[9], (num_perms,), dtype=jnp.bfloat16) * 0.1 res_scale = jnp.array([1.0], dtype=jnp.bfloat16) permutations = _get_permutation_matrices(streams) - cotangent = jax.random.normal(keys[11], (batch, sequence, streams, embedding), dtype=jnp.bfloat16) + cotangent = jax.random.normal( + keys[11], (batch, sequence, streams, embedding), dtype=jnp.bfloat16 + ) weights = mhc_kernel_common.MhcWeights( norm_scale=norm_scale, pre_alpha=pre_alpha, @@ -670,7 +762,9 @@ def _make_kernel_inputs(batch=2, sequence=64, streams=4, embedding=256, seed=0): return x, weights, permutations, cotangent -def _run_pipeline_reference(x, weights: mhc_kernel_common.MhcWeights, permutations): +def _run_pipeline_reference( + x, weights: mhc_kernel_common.MhcWeights, permutations +): """Runs the native JAX/XLA reference implementation of the mHC pipeline.""" batch, sequence, streams, embedding = x.shape tokens = batch * sequence @@ -680,7 +774,9 @@ def _run_pipeline_reference(x, weights: mhc_kernel_common.MhcWeights, permutatio flattened_f32 = x_flat.reshape(tokens, flattened_size).astype(jnp.float32) normalized = ( flattened_f32 - * jax.lax.rsqrt(jnp.mean(flattened_f32 * flattened_f32, axis=-1, keepdims=True) + 1e-5) + * jax.lax.rsqrt( + jnp.mean(flattened_f32 * flattened_f32, axis=-1, keepdims=True) + 1e-5 + ) * weights.norm_scale.astype(jnp.float32) ).astype(x.dtype) @@ -722,7 +818,9 @@ def _run_pipeline_reference(x, weights: mhc_kernel_common.MhcWeights, permutatio ) residual = jnp.dot( weights_res, - permutations.reshape(permutation_count, streams * streams).astype(jnp.float32), + permutations.reshape(permutation_count, streams * streams).astype( + jnp.float32 + ), ).reshape(tokens, streams, streams) residual_mix = jnp.einsum( @@ -731,7 +829,10 @@ def _run_pipeline_reference(x, weights: mhc_kernel_common.MhcWeights, permutatio x_flat, preferred_element_type=jnp.float32, ) - post_mix = h_post.astype(jnp.float32)[:, :, None] * layer_input.astype(jnp.float32)[:, None, :] + post_mix = ( + h_post.astype(jnp.float32)[:, :, None] + * layer_input.astype(jnp.float32)[:, None, :] + ) return (residual_mix + post_mix).astype(x.dtype).reshape(x.shape) @@ -762,7 +863,9 @@ class TestMhcKernelsFwd(parameterized.TestCase): """Unit tests for MaxText mHC-lite Pallas forward kernel.""" def test_doubly_stochastic(self): - x, weights, permutations, _ = _make_kernel_inputs(batch=1, sequence=128, streams=4, embedding=128) + x, weights, permutations, _ = _make_kernel_inputs( + batch=1, sequence=128, streams=4, embedding=128 + ) config = mhc_kernel.MhcKernelConfig(interpret=True) _, context = mhc_kernel.pre( x, @@ -772,15 +875,21 @@ def test_doubly_stochastic(self): ) row_sums = jnp.sum(context.residual, axis=-1) col_sums = jnp.sum(context.residual, axis=-2) - np.testing.assert_allclose(row_sums, np.ones_like(row_sums), rtol=1e-3, atol=1e-3) - np.testing.assert_allclose(col_sums, np.ones_like(col_sums), rtol=1e-3, atol=1e-3) + np.testing.assert_allclose( + row_sums, np.ones_like(row_sums), rtol=1e-3, atol=1e-3 + ) + np.testing.assert_allclose( + col_sums, np.ones_like(col_sums), rtol=1e-3, atol=1e-3 + ) @parameterized.named_parameters( ("mosaic", "mosaic"), ("auto", None), ) def test_forward_parity(self, implementation): - x, weights, permutations, _ = _make_kernel_inputs(batch=2, sequence=64, streams=4, embedding=256) + x, weights, permutations, _ = _make_kernel_inputs( + batch=2, sequence=64, streams=4, embedding=256 + ) expected = _run_pipeline_reference(x, weights, permutations) actual = _run_pipeline_api( x, @@ -792,7 +901,9 @@ def test_forward_parity(self, implementation): np.testing.assert_allclose(actual, expected, rtol=5e-2, atol=5e-2) def test_unsupported_shape_raises_error(self): - x, weights, permutations, _ = _make_kernel_inputs(batch=1, sequence=16, streams=2, embedding=128) + x, weights, permutations, _ = _make_kernel_inputs( + batch=1, sequence=16, streams=2, embedding=128 + ) config = mhc_kernel.MhcKernelConfig(interpret=True) with self.assertRaises(mhc_kernel_common.UnsupportedInputError): mhc_kernel.pre( @@ -807,7 +918,9 @@ class TestMhcKernelsBwd(parameterized.TestCase): """Unit tests for MaxText mHC-lite Pallas backward kernel.""" def test_forward_and_backward_vjp_parity(self): - x, weights, permutations, cotangent = _make_kernel_inputs(batch=2, sequence=64, streams=4, embedding=256) + x, weights, permutations, cotangent = _make_kernel_inputs( + batch=2, sequence=64, streams=4, embedding=256 + ) expected_out, expected_vjp_fn = jax.vjp( lambda x_, w_: _run_pipeline_reference(x_, w_, permutations), x, @@ -816,7 +929,9 @@ def test_forward_and_backward_vjp_parity(self): expected_dx, expected_dw = expected_vjp_fn(cotangent) actual_out, actual_vjp_fn = jax.vjp( - lambda x_, w_: _run_pipeline_api(x_, w_, permutations, implementation=None, interpret=True), + lambda x_, w_: _run_pipeline_api( + x_, w_, permutations, implementation=None, interpret=True + ), x, weights, ) @@ -825,11 +940,17 @@ def test_forward_and_backward_vjp_parity(self): np.testing.assert_allclose(actual_out, expected_out, rtol=5e-2, atol=5e-2) actual_grads = (actual_dx,) + tuple(jax.tree_util.tree_leaves(actual_dw)) - expected_grads = (expected_dx,) + tuple(jax.tree_util.tree_leaves(expected_dw)) + expected_grads = (expected_dx,) + tuple( + jax.tree_util.tree_leaves(expected_dw) + ) self.assertEqual(len(actual_grads), len(expected_grads)) - for i, (actual_g, expected_g) in enumerate(zip(actual_grads, expected_grads)): + for i, (actual_g, expected_g) in enumerate( + zip(actual_grads, expected_grads) + ): tol = 0.05 if actual_g.size == 1 else 0.02 - scale = max(float(np.max(np.abs(np.asarray(expected_g, np.float32)))), 1e-7) + scale = max( + float(np.max(np.abs(np.asarray(expected_g, np.float32)))), 1e-7 + ) np.testing.assert_allclose( np.asarray(actual_g, np.float32), np.asarray(expected_g, np.float32), @@ -839,7 +960,9 @@ def test_forward_and_backward_vjp_parity(self): ) def test_forward_and_backward_vjp_parity_feature_tiled(self): - x, weights, permutations, cotangent = _make_kernel_inputs(batch=2, sequence=64, streams=4, embedding=256) + x, weights, permutations, cotangent = _make_kernel_inputs( + batch=2, sequence=64, streams=4, embedding=256 + ) config = mhc_kernel.MhcKernelConfig(bwd_feature_block_size=128) expected_out, expected_vjp_fn = jax.vjp( lambda x_, w_: _run_pipeline_reference(x_, w_, permutations), @@ -849,7 +972,9 @@ def test_forward_and_backward_vjp_parity_feature_tiled(self): expected_dx, expected_dw = expected_vjp_fn(cotangent) actual_out, actual_vjp_fn = jax.vjp( - lambda x_, w_: _run_pipeline_api(x_, w_, permutations, config=config, interpret=True), + lambda x_, w_: _run_pipeline_api( + x_, w_, permutations, config=config, interpret=True + ), x, weights, ) @@ -858,11 +983,17 @@ def test_forward_and_backward_vjp_parity_feature_tiled(self): np.testing.assert_allclose(actual_out, expected_out, rtol=5e-2, atol=5e-2) actual_grads = (actual_dx,) + tuple(jax.tree_util.tree_leaves(actual_dw)) - expected_grads = (expected_dx,) + tuple(jax.tree_util.tree_leaves(expected_dw)) + expected_grads = (expected_dx,) + tuple( + jax.tree_util.tree_leaves(expected_dw) + ) self.assertEqual(len(actual_grads), len(expected_grads)) - for i, (actual_g, expected_g) in enumerate(zip(actual_grads, expected_grads)): + for i, (actual_g, expected_g) in enumerate( + zip(actual_grads, expected_grads) + ): tol = 0.05 if actual_g.size == 1 else 0.02 - scale = max(float(np.max(np.abs(np.asarray(expected_g, np.float32)))), 1e-7) + scale = max( + float(np.max(np.abs(np.asarray(expected_g, np.float32)))), 1e-7 + ) np.testing.assert_allclose( np.asarray(actual_g, np.float32), np.asarray(expected_g, np.float32), @@ -876,7 +1007,9 @@ class TestMhcCostEstimates(unittest.TestCase): """Unit tests for analytical CostEstimate computations on MhcDims.""" def setUp(self): - self.dims = mhc_kernel_common.MhcDims(tokens=256, streams=4, embedding=512, num_permutations=24) + self.dims = mhc_kernel_common.MhcDims( + tokens=256, streams=4, embedding=512, num_permutations=24 + ) def test_dims_properties(self): self.assertEqual(self.dims.flattened_size, 4 * 512) @@ -897,10 +1030,20 @@ def test_cost_estimates_non_zero_and_valid(self): for name, cost in costs: self.assertIsInstance(cost.flops, int, msg=f"{name} flops must be int") self.assertGreater(cost.flops, 0, msg=f"{name} flops must be positive") - self.assertIsInstance(cost.bytes_accessed, int, msg=f"{name} bytes_accessed must be int") - self.assertGreater(cost.bytes_accessed, 0, msg=f"{name} bytes_accessed must be positive") - self.assertIsInstance(cost.transcendentals, int, msg=f"{name} transcendentals must be int") - self.assertGreaterEqual(cost.transcendentals, 0, msg=f"{name} transcendentals must be non-negative") + self.assertIsInstance( + cost.bytes_accessed, int, msg=f"{name} bytes_accessed must be int" + ) + self.assertGreater( + cost.bytes_accessed, 0, msg=f"{name} bytes_accessed must be positive" + ) + self.assertIsInstance( + cost.transcendentals, int, msg=f"{name} transcendentals must be int" + ) + self.assertGreaterEqual( + cost.transcendentals, + 0, + msg=f"{name} transcendentals must be non-negative", + ) # coeff kernels have transcendentals (sigmoid + softmax) self.assertGreater(self.dims.coeff_fwd_cost().transcendentals, 0) diff --git a/tests/unit/muon_utils_test.py b/tests/unit/muon_utils_test.py index 0b78056a1f..de066fd46e 100644 --- a/tests/unit/muon_utils_test.py +++ b/tests/unit/muon_utils_test.py @@ -16,18 +16,18 @@ # pylint: disable=protected-access -import io import contextlib +import io import unittest from unittest import mock -import jax -import jax.numpy as jnp from flax import linen as nn from flax import nnx -from optax.contrib._muon import MuonDimensionNumbers as mdn - +import jax +import jax.numpy as jnp +from maxtext.optimizers.muon import MuonDimensionNumbers as mdn from maxtext.utils import muon_utils +import numpy as np class TestIsPathContainAny(unittest.TestCase): @@ -81,11 +81,6 @@ def test_moe_prefused_wi_uses_last_two_axes(self): self.assertEqual(muon_utils.transform_logic(("decoder", "MoeBlock_0", "wi")), mdn((-2,), (-1,))) self.assertEqual(muon_utils.transform_logic(("decoder", "routed_experts", "wi")), mdn((-2,), (-1,))) - 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,))) - def test_qwen3_moe_block(self): self.assertEqual(muon_utils.transform_logic(("decoder", "moe_block", "wi_0")), mdn((-2,), (-1,))) self.assertEqual(muon_utils.transform_logic(("decoder", "moe_block", "wo")), mdn((-2,), (-1,))) @@ -94,37 +89,89 @@ def test_gpt_oss_mlp_moe(self): self.assertEqual(muon_utils.transform_logic(("decoder", "GptOssMlp", "wi_0")), mdn((-2,), (-1,))) self.assertEqual(muon_utils.transform_logic(("decoder", "GptOssMlp", "wo")), mdn((-2,), (-1,))) - def test_moe_gate_falls_through_to_standard(self): - # 'gate' is inside MoeBlock_0 but not one of (wi, wi_0, wi_1, wo) β†’ standard. - self.assertEqual(muon_utils.transform_logic(("decoder", "MoeBlock_0", "gate", "kernel")), mdn((0,), (-1,))) - self.assertEqual(muon_utils.transform_logic(("decoder", "routed_experts", "gate", "kernel")), mdn((0,), (-1,))) + def test_moe_gate_uses_standard_axes(self): + # 'gate' is a 2D router matrix [in_features, (num_layers), num_experts] -> standard (0,), (-1,). + self.assertEqual( + muon_utils.transform_logic(("decoder", "MoeBlock_0", "gate", "kernel")), + mdn((0,), (-1,)), + ) + + def test_gate_up_proj_uses_last_two_axes(self): + self.assertEqual( + muon_utils.transform_logic(("decoder", "MoeBlock_0", "gate_up_proj")), + mdn((-2,), (-1,)), + ) + self.assertEqual( + muon_utils.transform_logic(("decoder", "mlp", "gate_proj")), + mdn((0,), (-1,)), + ) + + def test_1d_tensor_is_excluded(self): + self.assertIsNone(muon_utils.transform_logic(("decoder", "mlp", "custom_param"), shape=(512,))) + + def test_2d_attention_projections_use_2d_axes(self): + self.assertEqual( + muon_utils.transform_logic(("decoder", "self_attention", "query"), shape=(512, 512)), + mdn((0,), (-1,)), + ) + self.assertEqual( + muon_utils.transform_logic(("decoder", "self_attention", "out"), shape=(512, 512)), + mdn((0,), (-1,)), + ) # --- 2.2 Self-attention --- def test_self_attention_out_projection(self): - self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "out")), mdn((0, -2), (-1,))) + self.assertEqual( + muon_utils.transform_logic(("decoder", "self_attention", "out")), + mdn((0, -2), (-1,)), + ) def test_self_attention_query_projection(self): - self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "query")), mdn((0,), (-2, -1))) + self.assertEqual( + muon_utils.transform_logic(("decoder", "self_attention", "query")), + mdn((0,), (-2, -1)), + ) def test_self_attention_key_projection(self): - self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "key")), mdn((0,), (-2, -1))) + self.assertEqual( + muon_utils.transform_logic(("decoder", "self_attention", "key")), + mdn((0,), (-2, -1)), + ) def test_self_attention_value_projection(self): - self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "value")), mdn((0,), (-2, -1))) + self.assertEqual( + muon_utils.transform_logic(("decoder", "self_attention", "value")), + mdn((0,), (-2, -1)), + ) def test_self_attention_wq_b_and_wkv_b(self): - self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "wq_b")), mdn((0,), (-2, -1))) - self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "wkv_b")), mdn((0,), (-2, -1))) + self.assertEqual( + muon_utils.transform_logic(("decoder", "self_attention", "wq_b")), + mdn((0,), (-2, -1)), + ) + self.assertEqual( + muon_utils.transform_logic(("decoder", "self_attention", "wkv_b")), + mdn((0,), (-2, -1)), + ) - def test_self_attention_mla_wq_a_is_excluded_from_special(self): + def test_self_attention_mla_wq_a_is_standard(self): # wq_a / wkv_a are MLA down-projections; they fall through the self_attention branch - # without matching anything, so the function returns the default standard mdn((0,), (-1,)). - self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "wq_a")), mdn((0,), (-1,))) - self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "wkv_a")), mdn((0,), (-1,))) + # to standard mdn((0,), (-1,)). + self.assertEqual( + muon_utils.transform_logic(("decoder", "self_attention", "wq_a")), + mdn((0,), (-1,)), + ) + self.assertEqual( + muon_utils.transform_logic(("decoder", "self_attention", "wkv_a")), + mdn((0,), (-1,)), + ) # --- 3. Standard --- def test_standard_weight(self): - self.assertEqual(muon_utils.transform_logic(("decoder", "mlp", "kernel")), mdn((0,), (-1,))) + self.assertEqual( + muon_utils.transform_logic(("decoder", "mlp", "kernel")), + mdn((0,), (-1,)), + ) # --- 4. DeepSeek V4 Specific --- def test_deepseek_v4_exclusions(self): @@ -132,6 +179,7 @@ def test_deepseek_v4_exclusions(self): self.assertIsNone(muon_utils.transform_logic(("decoder", "hc_head", "hc_scale"))) self.assertIsNone(muon_utils.transform_logic(("decoder", "hc_head", "hc_base"))) self.assertIsNone(muon_utils.transform_logic(("decoder", "layers", "layers_0", "mhc_attention", "res_beta"))) + self.assertIsNone(muon_utils.transform_logic(("decoder", "layers", "layers_0", "mhc_attention", "res_alpha_scale"))) self.assertIsNone(muon_utils.transform_logic(("decoder", "layers", "layers_0", "self_attention", "sinks"))) self.assertIsNone( muon_utils.transform_logic(("decoder", "layers", "layers_0", "mlp", "MoeBlock_0", "gate", "tid2eid")) @@ -140,9 +188,24 @@ def test_deepseek_v4_exclusions(self): muon_utils.transform_logic(("decoder", "layers", "layers_0", "self_attention", "rotary_embedding", "inv_freq")) ) + def test_deepseek_v4_mhc_alpha_projections(self): + # Alpha projection matrices in MHC are 2D dense linear projections + self.assertEqual( + muon_utils.transform_logic(("decoder", "layers", "layers_0", "mhc_attention", "pre_alpha")), + mdn((0,), (-1,)), + ) + self.assertEqual( + muon_utils.transform_logic(("decoder", "layers", "layers_0", "mhc_attention", "post_alpha")), + mdn((0,), (-1,)), + ) + self.assertEqual( + muon_utils.transform_logic(("decoder", "layers", "layers_0", "mhc_attention", "res_alpha")), + mdn((0,), (-1,)), + ) + def test_deepseek_v4_self_attention_grouped_projection(self): - # o_a_proj projects with reduction on in_features_per_group (-2) - # and output on out_features_per_group (-1) + # o_a_proj is a grouped linear layer [o_groups, (layers), in_features_per_group, out_features_per_group] + # with reduction on in_features_per_group (-2) and output on out_features_per_group (-1) self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "o_a_proj")), mdn((-2,), (-1,))) def test_deepseek_v4_position_bias_is_standard_weight(self): @@ -157,6 +220,41 @@ def test_deepseek_v4_position_bias_is_standard_weight(self): mdn((0,), (-1,)), ) + # --- 5. Qwen3-Next Specific --- + 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,)), + ) + + def test_qwen3_next_gdn_projections(self): + self.assertEqual( + muon_utils.transform_logic(("decoder", "gdn", "in_proj_qkvz")), + mdn((0,), (-1,)), + ) + self.assertEqual( + muon_utils.transform_logic(("decoder", "gdn", "in_proj_ba")), + mdn((0,), (-1,)), + ) + self.assertEqual( + muon_utils.transform_logic(("decoder", "gdn", "out_proj")), + mdn((0,), (-1,)), + ) + + 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", "shared_expert_gate"))) + class TestGetTransformTree(unittest.TestCase): """Tests for get_transform_tree: recursive dict walk that applies transform_logic.""" @@ -176,16 +274,22 @@ def test_non_dict_leaf_at_root_returns_transform(self): self.assertEqual(muon_utils.get_transform_tree(0), mdn((0,), (-1,))) +class _AttentionSubModule(nnx.Module): + + def __init__(self): + self.out = nnx.Param(jnp.ones((2, 4, 8))) + + class _MoeLikeNNXModel(nnx.Module): """Small NNX model whose param paths exercise the NNX branch of get_muon_weight_dimension_numbers.""" def __init__(self, rngs): # Names are chosen so transform_logic matches each of the three meaningful branches: # - w_standard: default mdn - # - self_attention_out: attention-out mdn + # - self_attention.out: attention-out mdn # - scale: excluded (None) self.w_standard = nnx.Param(jnp.ones((4, 8))) - self.self_attention_out = nnx.Param(jnp.ones((4, 8))) + self.self_attention = _AttentionSubModule() self.scale = nnx.Param(jnp.ones((8,))) @@ -201,7 +305,8 @@ def test_nnx_model_dispatches_to_tree_map_with_path(self): # Result is an nnx.State whose top-level keys mirror the model attributes. self.assertIn("w_standard", result) - self.assertIn("self_attention_out", result) + self.assertIn("self_attention", result) + self.assertIn("out", result["self_attention"]) self.assertIn("scale", result) # NNX Variables are walked by jax.tree_util.tree_map_with_path, so the returned @@ -210,6 +315,30 @@ def test_nnx_model_dispatches_to_tree_map_with_path(self): self.assertIsNone(result["scale"]) # 'w_standard' does not trigger any special rule β†’ standard mdn. self.assertEqual(result["w_standard"], mdn((0,), (-1,))) + self.assertEqual(result["self_attention"]["out"], mdn((0, -2), (-1,))) + + def test_nnx_model_with_logical_axis_rules(self): + """Verifies that config.logical_axis_rules is active within get_muon_weight_dimension_numbers.""" + config = mock.MagicMock() + config.logical_axis_rules = (("embed", "fsdp"), ("mlp", "tensor")) + result = muon_utils.get_muon_weight_dimension_numbers( + self.model, config=config + ) + self.assertEqual(result["w_standard"], mdn((0,), (-1,))) + self.assertEqual(result["self_attention"]["out"], mdn((0, -2), (-1,))) + self.assertIsNone(result["scale"]) + + def test_nnx_model_with_mesh_populates_named_sharding(self): + """Verifies that NamedSharding is properly attached to MuonDimensionNumbers when mesh is present.""" + devices = np.array(jax.devices()[:1]).reshape((1, 1)) + mesh = jax.sharding.Mesh(devices, ("data", "model")) + result = muon_utils.get_muon_weight_dimension_numbers( + self.model, config=None, mesh=mesh + ) + self.assertIsNotNone(result["w_standard"].sharding) + self.assertIsInstance( + result["w_standard"].sharding, jax.sharding.NamedSharding + ) def test_nnx_verbose_path_executes_print_debug(self): """verbose=True should also execute _print_structure_debug without raising.""" @@ -225,7 +354,6 @@ class TestGetMuonWeightDimensionNumbersLinen(unittest.TestCase): def test_linen_branch_uses_get_abstract_param(self): """Linen models dispatch to maxtext_utils.get_abstract_param + get_transform_tree.""" - # Build a Linen nn.Module so isinstance(model, nnx.Module) is False. class LinenStub(nn.Module): @@ -235,8 +363,6 @@ def __call__(self, x): model = LinenStub() - # Mock the heavy get_abstract_param call with a pre-shaped dict that exercises - # both a standard weight path and an excluded path. fake_abstract_param = { "params": { "self_attention": {"out": object()}, @@ -261,7 +387,10 @@ def test_handles_logically_partitioned_leaf(self): buf = io.StringIO() with contextlib.redirect_stdout(buf): - muon_utils._print_structure_debug(tree, muon_weight_dimension_numbers={"params": {"kernel": mdn((0,), (-1,))}}) + muon_utils._print_structure_debug( + tree, + muon_weight_dimension_numbers={"params": {"kernel": mdn((0,), (-1,))}}, + ) out = buf.getvalue() self.assertIn("(4, 8)", out) self.assertIn("embed", out) @@ -277,5 +406,88 @@ def test_handles_shape_dtype_struct_leaf(self): self.assertIn("(16, 32)", out) +class TestLayerScanInvariance(unittest.TestCase): + """Verifies Muon dimension mapping is invariant to layer scanning.""" + + def test_transform_tree_identical_with_and_without_layer_scan(self): + """Parameter paths return identical dimension specs with/without scan.""" + unscanned_tree = { + "decoder": { + "self_attention": { + "query": jax.ShapeDtypeStruct((512, 8, 64), jnp.float32), + "out": jax.ShapeDtypeStruct((8, 64, 512), jnp.float32), + }, + "mlp": { + "wi_0": jax.ShapeDtypeStruct((512, 2048), jnp.float32), + "wo": jax.ShapeDtypeStruct((2048, 512), jnp.float32), + }, + "norm": { + "scale": jax.ShapeDtypeStruct((512,), jnp.float32), + }, + } + } + # Place layer dimension (num_layers=12) on axis 1 (param_scan_axis=1) for + # all layers to simulate layer scanning. + scanned_tree = { + "decoder": { + "self_attention": { + "query": jax.ShapeDtypeStruct((512, 12, 8, 64), jnp.float32), + "out": jax.ShapeDtypeStruct((8, 12, 64, 512), jnp.float32), + }, + "mlp": { + "wi_0": jax.ShapeDtypeStruct((512, 12, 2048), jnp.float32), + "wo": jax.ShapeDtypeStruct((2048, 12, 512), jnp.float32), + }, + "norm": { + "scale": jax.ShapeDtypeStruct((512, 12), jnp.float32), + }, + } + } + + unscanned_mdn = muon_utils.get_transform_tree(unscanned_tree) + scanned_mdn = muon_utils.get_transform_tree(scanned_tree) + + self.assertEqual(unscanned_mdn, scanned_mdn) + + def test_relative_dimensions_resolve_consistent_matrix_features(self): + """Dimension numbers index same feature axes regardless of scan axis.""" + # Standard MLP: 2D unscanned vs 3D scanned (param_scan_axis = 1, num_layers = 12) + mlp_unscanned_shape = (512, 2048) + mlp_scanned_shape = (512, 12, 2048) + mlp_mdn = muon_utils.transform_logic(("decoder", "mlp", "kernel")) + + # Reduction axis (0) and Output axis (-1) point to identical feature sizes + self.assertEqual(mlp_unscanned_shape[mlp_mdn.reduction_axis[0]], 512) + self.assertEqual(mlp_scanned_shape[mlp_mdn.reduction_axis[0]], 512) + self.assertEqual(mlp_unscanned_shape[mlp_mdn.output_axis[0]], 2048) + self.assertEqual(mlp_scanned_shape[mlp_mdn.output_axis[0]], 2048) + + # Attention QKV: 3D unscanned vs 4D scanned (param_scan_axis = 1, num_layers = 12) + qkv_unscanned_shape = (512, 8, 64) + qkv_scanned_shape = (512, 12, 8, 64) + qkv_mdn = muon_utils.transform_logic(("decoder", "self_attention", "query")) + + self.assertEqual(qkv_unscanned_shape[qkv_mdn.reduction_axis[0]], 512) + self.assertEqual(qkv_scanned_shape[qkv_mdn.reduction_axis[0]], 512) + self.assertEqual(tuple(qkv_unscanned_shape[ax] for ax in qkv_mdn.output_axis), (8, 64)) + self.assertEqual(tuple(qkv_scanned_shape[ax] for ax in qkv_mdn.output_axis), (8, 64)) + + # Attention Out: 3D unscanned vs 4D scanned (param_scan_axis = 1, num_layers = 12) + out_unscanned_shape = (8, 64, 512) + out_scanned_shape = (8, 12, 64, 512) + out_mdn = muon_utils.transform_logic(("decoder", "self_attention", "out")) + + self.assertEqual( + tuple(out_unscanned_shape[ax] for ax in out_mdn.reduction_axis), + (8, 64), + ) + self.assertEqual( + tuple(out_scanned_shape[ax] for ax in out_mdn.reduction_axis), + (8, 64), + ) + self.assertEqual(out_unscanned_shape[out_mdn.output_axis[0]], 512) + self.assertEqual(out_scanned_shape[out_mdn.output_axis[0]], 512) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index 25718ba94f..04adf447d2 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -28,22 +28,18 @@ from unittest import mock from unittest.mock import MagicMock, patch -import pytest - -import jax -import jax.numpy as jnp -import numpy as np from flax import linen as nn from flax import nnx +import jax +import jax.numpy as jnp from jax.sharding import Mesh - from maxtext.common.common_types import ( + AttentionType, DECODING_ACTIVE_SEQUENCE_INDICATOR, + DecoderBlockType, MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, - AttentionType, - DecoderBlockType, MultimodalInput, ) from maxtext.configs import pyconfig @@ -57,6 +53,8 @@ from maxtext.models.llama2 import LlamaDecoderLayer from maxtext.utils import maxtext_utils, maxtext_utils_nnx from tests.utils.test_helpers import get_test_config_path +import numpy as np +import pytest # --------------------------------------------------------------------------- # Shared minimal config overrides used across most tests @@ -625,7 +623,9 @@ def fake_apply_embedding( _deterministic, _model_mode, multimodal_input=None, + decoder_input_embeddings=None, ): + del decoder_input_embeddings captured["multimodal_input"] = multimodal_input batch = self.cfg.global_batch_size_to_train_on seq_len = self.cfg.max_target_length @@ -655,6 +655,26 @@ def fake_apply_embedding( self.assertTrue(jnp.array_equal(forwarded.audio_masks, sentinel_aud_mask)) self.assertTrue(jnp.array_equal(forwarded.bidirectional_mask, sentinel_bidir)) + def test_precomputed_embeddings_bypass_initial_multimodal_merge(self): + """Complete input embeddings must not be merged with vision embeddings again.""" + ids, _, positions = self._make_token_inputs() + embeddings = jnp.ones( + (self.cfg.global_batch_size_to_train_on, self.cfg.max_target_length, self.cfg.emb_dim), + dtype=self.cfg.dtype, + ) + + result = self.decoder._apply_embedding( # pylint: disable=protected-access + lambda *_args, **_kwargs: self.fail("token embedding should be bypassed"), + ids, + positions, + True, + MODEL_MODE_TRAIN, + multimodal_input=object(), + decoder_input_embeddings=embeddings, + ) + + self.assertTrue(jnp.array_equal(result, embeddings)) + def test_different_random_seeds_produce_different_logits(self): """Two randomly-initialised decoders should not produce identical logits.""" cfg = self.cfg @@ -717,6 +737,107 @@ def test_scan_layers(self): self.assertEqual(logits.shape, (batch, seq_len, cfg.vocab_size)) +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, + full_attention_layer_offset=3, + ) + + 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): """Small stand-in that exposes cache ordering and mutable-state updates.""" @@ -1077,6 +1198,56 @@ def make_block(): self.assertEqual(len(updated_kvs), num_layers) np.testing.assert_allclose(y_external, y_scanned, rtol=1e-5, atol=1e-5) + def test_qwen3_next_scanned_layers(self): + """Test NNXDecoder with qwen3_next block, dense prefix, and scan_layers=True.""" + cfg = _make_config( + decoder_block="qwen3_next", + scan_layers=True, + first_num_dense_layers=1, + inhomogeneous_layer_cycle_interval=3, + full_attention_layer_offset=0, + num_decoder_layers=4, + base_emb_dim=128, + base_num_query_heads=4, + base_num_kv_heads=4, + base_mlp_dim=256, + base_moe_mlp_dim=128, + shared_experts=1, + num_experts=4, + num_experts_per_tok=2, + gdn_num_key_heads=4, + gdn_num_value_heads=4, + gdn_key_head_dim=32, + gdn_value_head_dim=32, + vocab_size=256, + max_target_length=16, + ) + decoder = NNXDecoder( + config=cfg, + mesh=self.mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=self.rngs, + ) + shared_embedding = self._make_shared_embedding(cfg) + ids, segment_ids, positions = self._make_token_inputs(cfg) + + logits, _, _ = decoder( + shared_embedding, + ids, + positions, + decoder_segment_ids=segment_ids, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + self.assertEqual( + logits.shape, + ( + cfg.global_batch_size_to_train_on, + cfg.max_target_length, + cfg.vocab_size, + ), + ) + @pytest.mark.tpu_only class TestGemma4SmallNNXDecoder(unittest.TestCase): @@ -1317,7 +1488,5 @@ def __call__(self, x, **kwargs): ) finally: maxtext_utils_nnx.nnx_add_and_sync_scan_axis = original_add_scan_axis - - if __name__ == "__main__": unittest.main() diff --git a/tests/unit/optimizers_test.py b/tests/unit/optimizers_test.py index 0057dab45c..ba3609b474 100644 --- a/tests/unit/optimizers_test.py +++ b/tests/unit/optimizers_test.py @@ -13,24 +13,23 @@ # limitations under the License. """Unit tests for all optimizers.""" + import re +from typing import NamedTuple import unittest -from unittest.mock import patch, MagicMock -import jax -import optax -import jax.numpy as jnp +from unittest.mock import MagicMock, patch -import pytest -from absl.testing import parameterized +from absl.testing import absltest, parameterized from flax import nnx -from optax.contrib import MuonDimensionNumbers as mdn - +import jax +import jax.numpy as jnp from maxtext.configs import pyconfig from maxtext.optimizers import optimizers +from maxtext.optimizers.muon import MuonDimensionNumbers as mdn from maxtext.utils import maxtext_utils, muon_utils from tests.utils.test_helpers import get_test_config_path -from typing import NamedTuple - +import optax +import pytest # deepseek2, specific: q_lora_rank=0 # applicable: deepseek2-16, but not deepseek2-236b (q_lora_rank=1536) @@ -156,7 +155,9 @@ "decoder": { "decoder_norm": {"scale": None}, "layers": {f"layers_{i}": _GEMMA3_LAYER for i in range(6)}, - "layers_remainder": {f"layers_{i}": _GEMMA3_LAYER for i in range(4)}, + "layers_remainder": { + f"layers_{i}": _GEMMA3_LAYER for i in range(4) + }, }, "token_embedder": {"embedding": None}, } @@ -170,15 +171,29 @@ "decoder_norm": {"scale": None}, "layers": { "mlp": { - "wi_0": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, - "wi_1": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, - "wo": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, + "wi_0": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-1,)) + }, + "wi_1": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-1,)) + }, + "wo": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-1,)) + }, }, "self_attention": { - "query": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))}, - "key": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))}, - "value": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))}, - "out": {"kernel": mdn(reduction_axis=(0, -2), output_axis=(-1,))}, + "query": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)) + }, + "key": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)) + }, + "value": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)) + }, + "out": { + "kernel": mdn(reduction_axis=(0, -2), output_axis=(-1,)) + }, }, "post_self_attention_layer_norm": {"scale": None}, "pre_self_attention_layer_norm": {"scale": None}, @@ -198,15 +213,29 @@ "decoder_norm": {"scale": None}, "layers": { "mlp": { - "wi_0": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, - "wi_1": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, - "wo": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, + "wi_0": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-1,)) + }, + "wi_1": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-1,)) + }, + "wo": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-1,)) + }, }, "self_attention": { - "query": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))}, - "key": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))}, - "value": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))}, - "out": {"kernel": mdn(reduction_axis=(0, -2), output_axis=(-1,))}, + "query": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)) + }, + "key": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)) + }, + "value": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)) + }, + "out": { + "kernel": mdn(reduction_axis=(0, -2), output_axis=(-1,)) + }, "key_norm": {"scale": None}, "query_norm": {"scale": None}, }, @@ -218,7 +247,6 @@ } } - # qwen3 MoE (e.g. qwen3-30b-a3b) QWEN3_MOE_DIMENSION_NUMBER = { "params": { @@ -226,7 +254,7 @@ "decoder_norm": {"scale": None}, "layers": { "moe_block": { - "gate": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, + "gate": {"kernel": mdn((0,), (-1,))}, "wi_0": mdn(reduction_axis=(-2,), output_axis=(-1,)), "wi_1": mdn(reduction_axis=(-2,), output_axis=(-1,)), "wo": mdn(reduction_axis=(-2,), output_axis=(-1,)), @@ -234,10 +262,18 @@ "post_self_attention_layer_norm": {"scale": None}, "pre_self_attention_layer_norm": {"scale": None}, "self_attention": { - "query": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))}, - "key": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))}, - "value": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))}, - "out": {"kernel": mdn(reduction_axis=(0, -2), output_axis=(-1,))}, + "query": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)) + }, + "key": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)) + }, + "value": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)) + }, + "out": { + "kernel": mdn(reduction_axis=(0, -2), output_axis=(-1,)) + }, "key_norm": {"scale": None}, "query_norm": {"scale": None}, }, @@ -256,9 +292,11 @@ "decoder_norm": {"scale": None}, "layers": { "latent_norm": {"scale": None}, - "layer_up_projection": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, + "layer_up_projection": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-1,)) + }, "moe_block": { - "gate": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, + "gate": {"kernel": mdn((0,), (-1,))}, "wi_0": mdn(reduction_axis=(-2,), output_axis=(-1,)), "wi_1": mdn(reduction_axis=(-2,), output_axis=(-1,)), "wo": mdn(reduction_axis=(-2,), output_axis=(-1,)), @@ -266,10 +304,18 @@ "post_self_attention_layer_norm": {"scale": None}, "pre_self_attention_layer_norm": {"scale": None}, "self_attention": { - "query": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))}, - "key": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))}, - "value": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))}, - "out": {"kernel": mdn(reduction_axis=(0, -2), output_axis=(-1,))}, + "query": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)) + }, + "key": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)) + }, + "value": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)) + }, + "out": { + "kernel": mdn(reduction_axis=(0, -2), output_axis=(-1,)) + }, "key_norm": {"scale": None}, "query_norm": {"scale": None}, }, @@ -284,7 +330,7 @@ # qwen3-next (e.g. qwen3-next-80b-a3b) _QWEN3_NEXT_MLP = { "routed_experts": { - "gate": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, + "gate": {"kernel": mdn((0,), (-1,))}, "wi_0": mdn(reduction_axis=(-2,), output_axis=(-1,)), "wi_1": mdn(reduction_axis=(-2,), output_axis=(-1,)), "wo": mdn(reduction_axis=(-2,), output_axis=(-1,)), @@ -351,15 +397,27 @@ # gpt-oss (e.g. gpt-oss-20b) _GPT_OSS_ATTENTION = { - "query": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)), "bias": None}, - "key": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)), "bias": None}, - "value": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)), "bias": None}, - "out": {"kernel": mdn(reduction_axis=(0, -2), output_axis=(-1,)), "bias": None}, + "query": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)), + "bias": None, + }, + "key": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)), + "bias": None, + }, + "value": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1)), + "bias": None, + }, + "out": { + "kernel": mdn(reduction_axis=(0, -2), output_axis=(-1,)), + "bias": None, + }, "sinks": None, } _GPT_OSS_MLP = { - "gate": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,)), "bias": None}, + "gate": {"kernel": mdn((0,), (-1,)), "bias": None}, "wi_0": mdn(reduction_axis=(-2,), output_axis=(-1,)), "wi_0_bias": None, "wi_1": mdn(reduction_axis=(-2,), output_axis=(-1,)), @@ -391,20 +449,7 @@ # deepseek4 building blocks -_DEEPSEEK4_MHC_ATTENTION = { - "mhc_norm": {"scale": None}, - "post_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)), - "post_alpha_scale": None, - "post_beta": None, - "pre_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)), - "pre_alpha_scale": None, - "pre_beta": None, - "res_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)), - "res_alpha_scale": None, - "res_beta": None, -} - -_DEEPSEEK4_MHC_MLP = { +_DEEPSEEK4_MHC = { "mhc_norm": {"scale": None}, "post_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)), "post_alpha_scale": None, @@ -419,7 +464,7 @@ _DEEPSEEK4_MLP = { "MoeBlock_0": { - "gate": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, + "gate": {"kernel": mdn((0,), (-1,))}, "wi_0": mdn(reduction_axis=(-2,), output_axis=(-1,)), "wi_1": mdn(reduction_axis=(-2,), output_axis=(-1,)), "wo": mdn(reduction_axis=(-2,), output_axis=(-1,)), @@ -448,12 +493,16 @@ "csa_compressor": { "gate_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, "indexer": { - "gate_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, + "gate_proj": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-1,)) + }, "kv_norm": {"scale": None}, "kv_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, "position_bias": mdn(reduction_axis=(0,), output_axis=(-1,)), "q_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, - "weights_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, + "weights_proj": { + "kernel": mdn(reduction_axis=(0,), output_axis=(-1,)) + }, }, "kv_norm": {"scale": None}, "kv_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))}, @@ -487,8 +536,8 @@ } _DEEPSEEK4_LAYER_BASIC = { - "mhc_attention": _DEEPSEEK4_MHC_ATTENTION, - "mhc_mlp": _DEEPSEEK4_MHC_MLP, + "mhc_attention": _DEEPSEEK4_MHC, + "mhc_mlp": _DEEPSEEK4_MHC, "mlp": _DEEPSEEK4_MLP, "post_self_attention_layer_norm": {"scale": None}, "pre_self_attention_layer_norm": {"scale": None}, @@ -496,8 +545,8 @@ } _DEEPSEEK4_LAYER_CSA_PREFIX = { - "mhc_attention": _DEEPSEEK4_MHC_ATTENTION, - "mhc_mlp": _DEEPSEEK4_MHC_MLP, + "mhc_attention": _DEEPSEEK4_MHC, + "mhc_mlp": _DEEPSEEK4_MHC, "mlp": _DEEPSEEK4_MLP, "post_self_attention_layer_norm": {"scale": None}, "pre_self_attention_layer_norm": {"scale": None}, @@ -505,8 +554,8 @@ } _DEEPSEEK4_LAYER_CSA_SCANNED = { - "mhc_attention": _DEEPSEEK4_MHC_ATTENTION, - "mhc_mlp": _DEEPSEEK4_MHC_MLP, + "mhc_attention": _DEEPSEEK4_MHC, + "mhc_mlp": _DEEPSEEK4_MHC, "mlp": _DEEPSEEK4_MLP_SCANNED, "post_self_attention_layer_norm": {"scale": None}, "pre_self_attention_layer_norm": {"scale": None}, @@ -514,8 +563,8 @@ } _DEEPSEEK4_LAYER_HCA_SCANNED = { - "mhc_attention": _DEEPSEEK4_MHC_ATTENTION, - "mhc_mlp": _DEEPSEEK4_MHC_MLP, + "mhc_attention": _DEEPSEEK4_MHC, + "mhc_mlp": _DEEPSEEK4_MHC, "mlp": _DEEPSEEK4_MLP_SCANNED, "post_self_attention_layer_norm": {"scale": None}, "pre_self_attention_layer_norm": {"scale": None}, @@ -564,18 +613,24 @@ class MuonDimensionTest(parameterized.TestCase): ("gemma3-4b", "gemma3-4b", GEMMA3_DIMENSION_NUMBER), ("qwen3-0.6b", "qwen3-0.6b", QWEN3_DIMENSION_NUMBER), ("qwen3-30b-a3b", "qwen3-30b-a3b", QWEN3_MOE_DIMENSION_NUMBER), - ("qwen3-custom-30b-a3b", "qwen3-custom-30b-a3b", QWEN3_CUSTOM_MOE_DIMENSION_NUMBER), + ( + "qwen3-custom-30b-a3b", + "qwen3-custom-30b-a3b", + QWEN3_CUSTOM_MOE_DIMENSION_NUMBER, + ), ("qwen3-next-80b-a3b", "qwen3-next-80b-a3b", QWEN3_NEXT_DIMENSION_NUMBER), ("gpt-oss-20b", "gpt-oss-20b", GPT_OSS_DIMENSION_NUMBER), ) @pytest.mark.tpu_only def test_model_integration(self, model_name, expected_output): - """ - Initializes the specified MaxText model and asserts that the generated + """Initializes the specified MaxText model and asserts that the generated + Muon dimension numbers match the hardcoded reference. """ is_pure_nnx = model_name in {"deepseek4-284b"} - actual_output = muon_utils.get_model_mdn(model_name, scan_layers=True, pure_nnx=is_pure_nnx) + actual_output = muon_utils.get_model_mdn( + model_name, scan_layers=True, pure_nnx=is_pure_nnx + ) if "params" in expected_output and "params" in actual_output: self.assertEqual(actual_output["params"], expected_output["params"]) else: @@ -596,12 +651,21 @@ def test_get_adamw_mask_with_empty_mask(self): def test_get_adamw_mask_with_valid_mask(self): """Directly test the get_adamw_mask function with valid mask""" # Case 2: Mask in config - argv = ["", get_test_config_path(), "run_name=test", "adamw_mask=['bias', '.*norm', '.*ln.*']"] + argv = [ + "", + get_test_config_path(), + "run_name=test", + "adamw_mask=['bias', '.*norm', '.*ln.*']", + ] config = pyconfig.initialize(argv) mask_fn = optimizers.get_adamw_mask(config) self.assertTrue(callable(mask_fn)) - params = {"layer1": {"kernel": 1, "bias": 2}, "layer2": {"layer_norm": {"scale": 3}}, "layer3": {"ln": {"scale": 4}}} + params = { + "layer1": {"kernel": 1, "bias": 2}, + "layer2": {"layer_norm": {"scale": 3}}, + "layer3": {"ln": {"scale": 4}}, + } mask = mask_fn(params) self.assertTrue(mask["layer1"]["kernel"]) self.assertFalse(mask["layer1"]["bias"]) @@ -619,7 +683,10 @@ def test_get_adamw_mask_with_invalid_mask(self): optimizers.get_adamw_mask(config) def test_get_adamw_mask_with_getattrkey(self): - """Test that get_adamw_mask correctly handles GetAttrKey (e.g. from NamedTuples)""" + """Test that get_adamw_mask correctly handles GetAttrKey (e.g. + + from NamedTuples) + """ class MyParams(NamedTuple): kernel: jax.Array @@ -636,10 +703,10 @@ class MyParams(NamedTuple): self.assertFalse(mask.bias) @parameterized.named_parameters( - ("adamw", "adamw", "maxtext.optimizers.optimizers.optax.adamw"), - ("adam_pax", "adam_pax", "maxtext.optimizers.optimizers.adam_pax"), + ("adamw", "adamw", optimizers.optax, "adamw"), + ("adam_pax", "adam_pax", optimizers, "adam_pax"), ) - def test_optimizer_with_mask(self, opt_type, mock_path): + def test_optimizer_with_mask(self, opt_type, mock_target, mock_attr): """Test that optimizer receives the mask function from config and it works as expected""" # Create a config with a mask list including regex argv = [ @@ -652,7 +719,7 @@ def test_optimizer_with_mask(self, opt_type, mock_path): config = pyconfig.initialize(argv) learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) - with patch(mock_path) as mock_opt: + with patch.object(mock_target, mock_attr) as mock_opt: # Call get_optimizer optimizers.get_optimizer(config, learning_rate_schedule) @@ -665,7 +732,11 @@ def test_optimizer_with_mask(self, opt_type, mock_path): self.assertIsNotNone(mask_fn) # Test the behavior of mask_fn - params = {"layer1": {"kernel": 1, "bias": 2}, "layer2": {"layer_norm": {"scale": 3}}, "layer3": [4, 5]} + params = { + "layer1": {"kernel": 1, "bias": 2}, + "layer2": {"layer_norm": {"scale": 3}}, + "layer3": [4, 5], + } mask = mask_fn(params) @@ -680,16 +751,16 @@ def test_optimizer_with_mask(self, opt_type, mock_path): self.assertTrue(mask["layer3"][1]) @parameterized.named_parameters( - ("adamw", "adamw", "maxtext.optimizers.optimizers.optax.adamw"), - ("adam_pax", "adam_pax", "maxtext.optimizers.optimizers.adam_pax"), + ("adamw", "adamw", optimizers.optax, "adamw"), + ("adam_pax", "adam_pax", optimizers, "adam_pax"), ) - def test_optimizer_without_mask(self, opt_type, mock_path): + def test_optimizer_without_mask(self, opt_type, mock_target, mock_attr): """Test that optimizer receives None for mask when config is empty""" argv = ["", get_test_config_path(), "run_name=test", f"opt_type={opt_type}"] config = pyconfig.initialize(argv) learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config) - with patch(mock_path) as mock_opt: + with patch.object(mock_target, mock_attr) as mock_opt: # Call get_optimizer optimizers.get_optimizer(config, learning_rate_schedule) @@ -702,7 +773,8 @@ def test_optimizer_without_mask(self, opt_type, mock_path): class AdamPaxScalarLearningRateTest(parameterized.TestCase): """Cover both branches of adam_pax's callable-vs-scalar learning_rate_fn guard. - adam_pax accepts either a callable schedule (the usual case) or a pre-evaluated + adam_pax accepts either a callable schedule (the usual case) or a + pre-evaluated scalar (when wrapped by optax.inject_hyperparams). Both must produce identical parameter updates for the same effective learning rate. """ @@ -760,13 +832,18 @@ def learning_rate_schedule(step): # We can test the optimizer by creating some dummy params and gradients # and checking if the updates are zeroed out for non-trainable parameters. params = { - "layer1": {"kernel": jax.numpy.ones((2, 2)), "indexer": jax.numpy.ones((2, 2))}, + "layer1": { + "kernel": jax.numpy.ones((2, 2)), + "indexer": jax.numpy.ones((2, 2)), + }, "layer2": {"layer_norm": {"scale": jax.numpy.ones((2, 2))}}, "layer3": {"ln": {"scale": jax.numpy.ones((2, 2))}}, } # Give some non-zero gradients - grads = jax.tree_util.tree_map(lambda x: jax.numpy.ones_like(x) * 0.5, params) + grads = jax.tree_util.tree_map( + lambda x: jax.numpy.ones_like(x) * 0.5, params + ) # Initialize optimizer state opt_state = opt.init(params) @@ -781,11 +858,18 @@ def learning_rate_schedule(step): # 'layer1/indexer' matches, so it should be trained (update != 0) self.assertFalse(jax.numpy.all(updates["layer1"]["indexer"] == 0)) # 'layer2/layer_norm/scale' matches, so it should be trained (update != 0) - self.assertFalse(jax.numpy.all(updates["layer2"]["layer_norm"]["scale"] == 0)) + self.assertFalse( + jax.numpy.all(updates["layer2"]["layer_norm"]["scale"] == 0) + ) def test_get_optimizer_without_trainable_mask(self): """Test get_optimizer when trainable_parameters_mask is empty.""" - argv = ["", get_test_config_path(), "run_name=test", "trainable_parameters_mask=[]"] + argv = [ + "", + get_test_config_path(), + "run_name=test", + "trainable_parameters_mask=[]", + ] config = pyconfig.initialize(argv) # Use a constant learning rate > 0 to ensure non-zero updates @@ -809,7 +893,9 @@ class SkipStepOnSpikesTest(parameterized.TestCase): def _run_spike_test(self, spike_kwargs): inner_opt = optax.sgd(0.1) - opt = optimizers.skip_step_on_spikes(inner_opt, interval=4, scaling_factor=1.0) + opt = optimizers.skip_step_on_spikes( + inner_opt, interval=4, scaling_factor=1.0 + ) params = {"x": jnp.array([1.0])} opt_state = opt.init(params) @@ -818,18 +904,24 @@ def _run_spike_test(self, spike_kwargs): base_kwargs = {k: jnp.array(1.0) for k in spike_kwargs.keys()} # Step 0: count = 0 < 2, will not skip (count should be >= interval / 2) - updates, opt_state = opt.update({"x": jnp.array([1.0])}, opt_state, params, **base_kwargs) + updates, opt_state = opt.update( + {"x": jnp.array([1.0])}, opt_state, params, **base_kwargs + ) self.assertFalse(jnp.all(updates["x"] == 0.0)) self.assertFalse(opt_state["is_skipped"]) # Step 1: count = 1 < 2, will not skip. mean=1.0, std=0.0 (count should be >= interval / 2) - updates, opt_state = opt.update({"x": jnp.array([1.0])}, opt_state, params, **base_kwargs) + updates, opt_state = opt.update( + {"x": jnp.array([1.0])}, opt_state, params, **base_kwargs + ) self.assertFalse(jnp.all(updates["x"] == 0.0)) self.assertFalse(opt_state["is_skipped"]) # Step 2: count = 2. Spike! spike_kwargs_jnp = {k: jnp.array(v) for k, v in spike_kwargs.items()} - updates, opt_state = opt.update({"x": jnp.array([1.0])}, opt_state, params, **spike_kwargs_jnp) + updates, opt_state = opt.update( + {"x": jnp.array([1.0])}, opt_state, params, **spike_kwargs_jnp + ) self.assertTrue(jnp.all(updates["x"] == 0.0)) self.assertTrue(opt_state["is_skipped"]) @@ -844,7 +936,9 @@ def test_skip_step_on_both_spike(self): def test_no_skip_without_kwargs(self): inner_opt = optax.sgd(0.1) - opt = optimizers.skip_step_on_spikes(inner_opt, interval=4, scaling_factor=1.0) + opt = optimizers.skip_step_on_spikes( + inner_opt, interval=4, scaling_factor=1.0 + ) params = {"x": jnp.array([1.0])} opt_state = opt.init(params) @@ -862,14 +956,20 @@ class TestMuonLogic(unittest.TestCase): def test_is_path_contain_any(self): # pylint: disable=protected-access - self.assertTrue(muon_utils._is_path_contain_any(("a", "b"), ("x", "a", "z"))) - self.assertFalse(muon_utils._is_path_contain_any(("a", "b"), ("x", "y", "z"))) + self.assertTrue( + muon_utils._is_path_contain_any(("a", "b"), ("x", "a", "z")) + ) + self.assertFalse( + muon_utils._is_path_contain_any(("a", "b"), ("x", "y", "z")) + ) def test_transform_logic_exclusions(self): self.assertIsNone(muon_utils.transform_logic(("layer_0", "bias"))) self.assertIsNone(muon_utils.transform_logic(("layer_0", "scale"))) self.assertIsNone(muon_utils.transform_logic(("embedding", "kernel"))) - self.assertIsNone(muon_utils.transform_logic(("layer_0", "attention", "A_log"))) + self.assertIsNone( + muon_utils.transform_logic(("layer_0", "attention", "A_log")) + ) def test_transform_logic_moe(self): path = ("layers_0", "MoeBlock_0", "wi_0") @@ -885,19 +985,44 @@ def test_transform_logic_attention(self): self.assertEqual(muon_utils.transform_logic(path_q), mdn((0,), (-2, -1))) path_gpt_out = ("layers_0", "GptOssAttention", "out", "kernel") - self.assertEqual(muon_utils.transform_logic(path_gpt_out), mdn((0, -2), (-1,))) + self.assertEqual( + muon_utils.transform_logic(path_gpt_out), mdn((0, -2), (-1,)) + ) path_gpt_q = ("layers_0", "GptOssAttention", "query", "kernel") - self.assertEqual(muon_utils.transform_logic(path_gpt_q), mdn((0,), (-2, -1))) + self.assertEqual( + muon_utils.transform_logic(path_gpt_q), mdn((0,), (-2, -1)) + ) - path_qwen3_next_q = ("layers_0", "attention", "attention", "query", "kernel") - self.assertEqual(muon_utils.transform_logic(path_qwen3_next_q), mdn((0,), (-2, -1))) + path_qwen3_next_q = ( + "layers_0", + "attention", + "attention", + "query", + "kernel", + ) + self.assertEqual( + muon_utils.transform_logic(path_qwen3_next_q), mdn((0,), (-2, -1)) + ) - path_qwen3_next_out = ("layers_0", "attention", "attention", "out", "kernel") - self.assertEqual(muon_utils.transform_logic(path_qwen3_next_out), mdn((0,), (-1,))) + path_qwen3_next_out = ( + "layers_0", + "attention", + "attention", + "out", + "kernel", + ) + self.assertEqual( + muon_utils.transform_logic(path_qwen3_next_out), mdn((0,), (-1,)) + ) def test_get_transform_tree(self): - fake_tree = {"params": {"layer_0": {"kernel": "leaf", "bias": "leaf"}, "MoeBlock_0": {"wi_0": "leaf"}}} + fake_tree = { + "params": { + "layer_0": {"kernel": "leaf", "bias": "leaf"}, + "MoeBlock_0": {"wi_0": "leaf"}, + } + } result = muon_utils.get_transform_tree(fake_tree) self.assertEqual(result["params"]["layer_0"]["kernel"], mdn((0,), (-1,))) self.assertIsNone(result["params"]["layer_0"]["bias"]) @@ -956,8 +1081,8 @@ class DeepSeekAttention(nnx.Module): def __init__(self, rngs: nnx.Rngs): self.self_attention = nnx.Module() - self.self_attention.query = nnx.Linear(8, 8, rngs=rngs) - self.self_attention.out = nnx.Linear(8, 8, rngs=rngs) + self.self_attention.query = nnx.Param(jnp.ones((8, 4, 2))) + self.self_attention.out = nnx.Param(jnp.ones((4, 2, 8))) # Use eval_shape to create an abstract version of the model. model = nnx.eval_shape(lambda: DeepSeekAttention(nnx.Rngs(0))) @@ -965,9 +1090,9 @@ def __init__(self, rngs: nnx.Rngs): result = muon_utils.get_muon_weight_dimension_numbers(model, config) # Check attention query: [0] -> [-2, -1] - self.assertEqual(result.self_attention.query.kernel, mdn((0,), (-2, -1))) + self.assertEqual(result.self_attention.query, mdn((0,), (-2, -1))) # Check attention out: [0, -2] -> [-1] - self.assertEqual(result.self_attention.out.kernel, mdn((0, -2), (-1,))) + self.assertEqual(result.self_attention.out, mdn((0, -2), (-1,))) def test_muon_newton_schulz_config(self): """Verifies that muon optimizer configures Newton-Schulz parameters correctly based on model.""" @@ -986,8 +1111,10 @@ def test_muon_newton_schulz_config(self): config_ds4 = pyconfig.initialize(argv_ds4) with ( - patch("maxtext.optimizers.optimizers.get_muon_weight_dimension_numbers") as mock_get_mdn, - patch("maxtext.optimizers.optimizers.muon") as mock_muon, + patch.object( + optimizers, "get_muon_weight_dimension_numbers" + ) as mock_get_mdn, + patch.object(optimizers, "muon") as mock_muon, ): mock_get_mdn.return_value = {} optimizers.get_optimizer(config_ds4, learning_rate_schedule, model=model) @@ -998,15 +1125,25 @@ def test_muon_newton_schulz_config(self): self.assertEqual(kwargs["ns_coeffs"][-1], (2.0, -1.5, 0.5)) # Case 2: Standard Model (Llama2) (Defaults to 5-step schedule) - argv_llama = ["", get_test_config_path(), "run_name=test", "opt_type=muon", "model_name=llama2-7b"] + argv_llama = [ + "", + get_test_config_path(), + "run_name=test", + "opt_type=muon", + "model_name=llama2-7b", + ] config_llama = pyconfig.initialize(argv_llama) with ( - patch("maxtext.optimizers.optimizers.get_muon_weight_dimension_numbers") as mock_get_mdn, - patch("maxtext.optimizers.optimizers.muon") as mock_muon, + patch.object( + optimizers, "get_muon_weight_dimension_numbers" + ) as mock_get_mdn, + patch.object(optimizers, "muon") as mock_muon, ): mock_get_mdn.return_value = {} - optimizers.get_optimizer(config_llama, learning_rate_schedule, model=model) + optimizers.get_optimizer( + config_llama, learning_rate_schedule, model=model + ) mock_muon.assert_called_once() _, kwargs = mock_muon.call_args self.assertEqual(kwargs["ns_steps"], 5) @@ -1014,4 +1151,4 @@ def test_muon_newton_schulz_config(self): if __name__ == "__main__": - unittest.main() + absltest.main() diff --git a/tests/unit/qwen3_next_vs_reference_test.py b/tests/unit/qwen3_next_vs_reference_test.py index e9efad376f..0f6ecf70b4 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,50 @@ 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...") diff --git a/tests/unit/sharding_nnx_test.py b/tests/unit/sharding_nnx_test.py index 5e0db5e592..6c0ff8ca8c 100644 --- a/tests/unit/sharding_nnx_test.py +++ b/tests/unit/sharding_nnx_test.py @@ -203,7 +203,6 @@ def test_zero1_partitioned_optimizer_filters_masked_nodes(self): for leaf in leaves: self.assertIsInstance(leaf.get_value(), NamedSharding) - class TestNnxConstructNamedSharding(unittest.TestCase): """Unit tests for nnx_construct_named_sharding covering every branch. From e4f3f0a0110b53e13c6a2d09eaa2dcba7ce06339 Mon Sep 17 00:00:00 2001 From: Muskan Sharma Date: Fri, 4 Sep 2026 19:15:45 +0000 Subject: [PATCH 2/5] Add fallback import for must_fuse_call in jax_flash_attention.py --- src/maxtext/kernels/attention/jax_flash_attention.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/maxtext/kernels/attention/jax_flash_attention.py b/src/maxtext/kernels/attention/jax_flash_attention.py index 0c4e100ea5..d5d1583910 100644 --- a/src/maxtext/kernels/attention/jax_flash_attention.py +++ b/src/maxtext/kernels/attention/jax_flash_attention.py @@ -18,7 +18,11 @@ import jax from jax.experimental import layout from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask as mask_lib -from jax.experimental.xla_metadata import must_fuse_call +try: + from jax.experimental.xla_metadata import must_fuse_call +except (ImportError, AttributeError): + def must_fuse_call(name=""): + return lambda f: f import jax.numpy as jnp from maxtext.kernels.attention import splash_attention_kernel From eae4384629375cc742dedc33ee6e982df133784b Mon Sep 17 00:00:00 2001 From: Muskan Sharma Date: Fri, 4 Sep 2026 19:25:53 +0000 Subject: [PATCH 3/5] Fix create_orbax_checkpoint_manager argument count mismatch --- src/maxtext/common/checkpointing.py | 3 +++ src/maxtext/utils/train_utils.py | 1 - 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index 7191190f2b..ba4c38f36e 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -379,6 +379,9 @@ def create_orbax_checkpoint_manager( enable_autocheckpoint: bool = False, todelete_subdir: str | None = None, todelete_full_path: str | None = None, + checkpoint_storage_target_data_file_size_bytes: int | None = None, + *args, + **kwargs, ): """Returns specified Orbax (async or not) CheckpointManager or None if checkpointing is disabled.""" if not enable_checkpointing: diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index 768a6ce824..cf66072f54 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -102,7 +102,6 @@ def create_checkpoint_manager(config, mesh, init_state_fn): config.enable_autocheckpoint, config.checkpoint_todelete_subdir, config.checkpoint_todelete_full_path, - config.checkpoint_storage_target_data_file_size_bytes, ) # Use Colocated Python checkpointing dispatchers optimization (Single Controller only). From 74b3a4a2f44d732295c8aa3b1aee66f54174e3eb Mon Sep 17 00:00:00 2001 From: Muskan Sharma Date: Fri, 4 Sep 2026 19:33:47 +0000 Subject: [PATCH 4/5] Restore nnx_add_scan_axis in maxtext_utils_nnx.py --- src/maxtext/utils/maxtext_utils_nnx.py | 43 +++++--------------------- 1 file changed, 8 insertions(+), 35 deletions(-) diff --git a/src/maxtext/utils/maxtext_utils_nnx.py b/src/maxtext/utils/maxtext_utils_nnx.py index 1c50a72a0e..097d144be9 100644 --- a/src/maxtext/utils/maxtext_utils_nnx.py +++ b/src/maxtext/utils/maxtext_utils_nnx.py @@ -256,25 +256,11 @@ def _op(x): if not isinstance(x, nnx.Variable): return x - # Scanned stacks record their own axis name, such as "dense_layers" or "moe_layers", - # so prefer it over the caller's default. Otherwise the name never matches and the - # check below strips a real logical axis instead of the scan axis. - axis_name = x.get_metadata().get(nnx.PARTITION_NAME, name) - def remove_fn(l): - removed = axis_name in l - if removed: - l.remove(axis_name) - if len(l) > x.get_value().ndim: - if removed: - raise ValueError( - f"Sharding names {l} still exceed value rank {x.get_value().ndim} after removing scan axis " - f"{axis_name!r}; the partition metadata is inconsistent." - ) - raise ValueError( - f"Scan axis {axis_name!r} not found in sharding names {l} for a rank-{x.get_value().ndim} value; " - "the partition metadata is inconsistent." - ) + if name in l: + l.remove(name) + while len(l) > x.get_value().ndim: + l.pop(0) return l return nnx_update_sharding_meta(x, remove_fn) @@ -282,31 +268,18 @@ def remove_fn(l): return jax.tree.map(_op, tree, is_leaf=lambda x: isinstance(x, nnx.Variable)) -def nnx_add_and_sync_scan_axis(tree, name="layers", pos=0): - """Restores the scan axis on each variable's value and sharding metadata. - - jax.lax.scan stacks its outputs with the scan axis at position 0. For each - variable this moves that axis to the variable's own param_scan_axis (falling - back to pos when the metadata is absent) and inserts the matching axis name at - the same position, so the value and its sharding metadata stay aligned. - """ +def nnx_add_scan_axis(tree, name="layers", pos=0): + """Adds the given scan axis to the PartitionSpec at the specified position.""" def _op(x): if not isinstance(x, nnx.Variable): return x - axis_name = x.get_metadata().get(nnx.PARTITION_NAME, name) - target = x.get_metadata().get("param_scan_axis", pos) - - val = x.get_value() - if target != 0 and hasattr(val, "ndim") and val.ndim > target: - x = x.replace(value=jnp.moveaxis(val, 0, target)) - def add_fn(l): - if axis_name not in l: + if name not in l: while len(l) < x.get_value().ndim - 1: l.append(None) - l.insert(target, axis_name) + l.insert(pos, name) else: while len(l) < x.get_value().ndim: l.append(None) From 9db4b4f6b6fb6f9fd6e397c5e08e159412c04a11 Mon Sep 17 00:00:00 2001 From: Muskan Sharma Date: Fri, 4 Sep 2026 20:03:30 +0000 Subject: [PATCH 5/5] add changs from cl --- run_custom_qwen3_next_on_xpk.sh | 209 ++++++++++++++++++++++++++++++++ 1 file changed, 209 insertions(+) create mode 100644 run_custom_qwen3_next_on_xpk.sh diff --git a/run_custom_qwen3_next_on_xpk.sh b/run_custom_qwen3_next_on_xpk.sh new file mode 100644 index 0000000000..9b7e897bd8 --- /dev/null +++ b/run_custom_qwen3_next_on_xpk.sh @@ -0,0 +1,209 @@ +#!/bin/bash +set -e + +# Activate Python virtual environment +source /usr/local/google/home/muskansh/maxtext_env/bin/activate + +# --- Environment Variables --- +export PROJECT_ID="tpu-prod-env-one-vm" +export CLUSTER_NAME="v6e-256-c2b3-b478935789" +export ZONE="us-central2" + +# --- Configuration & Automated Image Build --- +TIMESTAMP=$(date +%m%d%H%M%S) +export WORKLOAD_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:muskansh_${TIMESTAMP}" +export WORKLOAD_NAME="muskansh-qn80b-${TIMESTAMP}" +export DEVICE_TYPE="v6e-256" +export NUM_SLICES=1 +export PRIORITY="very-high" +export MAX_RESTARTS=0 +export NUM_STEPS=15 +export MODEL_NAME="qwen3-next-80b-a3b" +export BASE_OUTPUT_DIR="/tmp/qwen3-next-80b-profiles/run-${TIMESTAMP}" + +echo "========================================================================" +echo "Building and uploading Docker runner image from /usr/local/google/home/muskansh/maxtext" +echo "Target Image: ${WORKLOAD_IMAGE}" +echo "========================================================================" + +( + cd /usr/local/google/home/muskansh/maxtext && \ + CLOUD_IMAGE_NAME="${WORKLOAD_IMAGE}" \ + BASE_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:latest" \ + bash src/dependencies/scripts/docker_upload_runner.sh +) + +echo "Docker image upload complete: ${WORKLOAD_IMAGE}" + +# --- XLA Flags --- +XLA_FLAGS_ARRAY=( + "--xla_msa_enable_sync_slice_replacement=false" + "--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true" + "--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true" + "--xla_msa_enable_sync_copy_replacement=false" + "--xla_tpu_scoped_vmem_limit_kib=81000" + "--xla_tpu_enable_sparse_core_collective_offload_all_gather=true" + "--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true" + "--xla_tpu_enable_concurrent_sparse_core_offloading=true" + "--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true" + "--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true" + "--xla_tpu_use_single_sparse_core_for_all_gather_offload=true" + "--xla_tpu_sparse_core_all_gather_latency_multiplier=1" + "--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3" + "--xla_tpu_offload_gather_to_sparsecore=true" + "--xla_tpu_dvfs_p_state=7" + "--xla_tpu_disable_sparse_core_collective_offload_remover=true" + "--xla_tpu_use_tc_device_shape_on_sc=true" + "--xla_sc_enable_instruction_fusion=false" + "--xla_sc_disable_megacore_partitioning=true" + "--xla_tpu_enable_async_collective_fusion=true" + "--xla_tpu_overlap_compute_collective_tc=true" + "--xla_tpu_enable_async_collective_fusion_multiple_steps=true" + "--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false" + "--xla_tpu_enable_async_collective_fusion_fuse_reduce_scatter=false" + "--xla_tpu_enable_async_collective_fusion_fuse_all_reduce=false" + "--xla_tpu_enable_latency_hiding_scheduler=true" + "--xla_latency_hiding_scheduler_rerun=10" + "--xla_tpu_all_gather_collective_matmul_mode=post_spmd_conservative" + "--xla_tpu_reduce_scatter_collective_matmul_mode=post_spmd_conservative" + "--xla_latency_hiding_scheduler_enable_selective_resources=true" + "--xla_tpu_enable_scheduler_memory_pressure_tracking=true" + "--xla_tpu_host_transfer_overlap_limit=4" + "--xla_tpu_aggressive_opt_barrier_removal=ENABLED" + "--xla_lhs_prioritize_async_depth_over_stall=ENABLED" + "--xla_should_allow_loop_variant_parameter_in_chain=ENABLED" + "--xla_should_add_loop_invariant_op_in_chain=ENABLED" + "--xla_max_concurrent_host_send_recv=100" + "--xla_tpu_rerun_latency_hiding_scheduler_post_sc_assignment=true" + "--xla_tpu_scheduler_percent_shared_memory_limit=150" +) +export XLA_FLAGS="${XLA_FLAGS_ARRAY[*]}" + +# --- MaxText Workload Overrides --- +MAXTEXT_ARGS_ARRAY=( + "model_name=${MODEL_NAME}" + "base_output_directory=${BASE_OUTPUT_DIR}" + "run_name=${WORKLOAD_NAME}" + "dataset_type=synthetic" + "dataset_name=synthetic" + "dtype=bfloat16" + "allow_split_physical_axes=True" + "ici_expert_parallelism=4" + "use_ring_of_experts=True" + "custom_mesh=hybrid_ring_64x4" + "use_ragged_sort=True" + "use_random_routing=False" + "per_device_batch_size=4" + "num_moe_token_chunks=2" + "opt_type=muon" + "max_target_length=2048" + "ragged_buffer_factor=1.5" + "remat_policy=custom" + "decoder_layer_input=device" + "context=device" + "reuse_example_batch=1" + "ici_fsdp_parallelism=-1" + "steps=${NUM_STEPS}" + "sa_block_q=1024" + "sa_block_kv=1024" + "sa_block_kv_compute=512" + "sa_block_q_dkv=1024" + "sa_block_kv_dkv=1024" + "sa_block_kv_dkv_compute=1024" + "sa_fuse_reciprocal=false" + "use_splash_scheduler=true" + "sa_use_base2_exp=true" + "dq_reduction_steps=3" + "hardware=tpu" + "skip_jax_distributed_system=False" + "attention=flash" + "use_tokamax_splash=True" + "sa_use_fused_bwd_kernel=True" + "sparse_matmul=True" + "megablox=True" + "wi_tile_fwd_batch_seq=512" + "wi_tile_dlhs_batch_seq=512" + "wi_tile_drhs_batch_seq=512" + "wo_tile_fwd_batch_seq=512" + "wo_tile_dlhs_batch_seq=512" + "wo_tile_drhs_batch_seq=512" + "wi_tile_fwd_embed_dim=3072" + "wi_tile_fwd_mlp_dim=1536" + "wi_tile_dlhs_embed_dim=3072" + "wi_tile_dlhs_mlp_dim=1536" + "wi_tile_drhs_embed_dim=3072" + "wi_tile_drhs_mlp_dim=1536" + "wo_tile_fwd_embed_dim=3072" + "wo_tile_fwd_mlp_dim=1536" + "wo_tile_dlhs_embed_dim=3072" + "wo_tile_dlhs_mlp_dim=1536" + "wo_tile_drhs_embed_dim=3072" + "wo_tile_drhs_mlp_dim=1536" + "use_tokamax_gmm=True" + "use_gmm_v2=True" + "optimizer_memory_host_offload=False" + "parameter_memory_host_offload=False" + "enable_checkpointing=False" + "async_checkpointing=False" + "tokenizer_type=huggingface" + "tokenizer_path=assets/tokenizers/qwen3-tokenizer" + "override_model_config=true" + "mhc_expansion_rate=4" + "enable_mhc_lite=True" + "use_mhc_pallas_kernel=True" + "mhc_pallas_kernel_fwd_block_size=256" + "mhc_pallas_kernel_bwd_block_size=256" + "use_gdn_kernel=True" + "use_hybrid_gdn=True" + "profiler=xplane" + "profiler_steps=2" + "skip_first_n_steps_for_profiler=1" + "enable_tpu_profiling_options=True" + "upload_all_profiler_results=False" + "enable_tensorboard=False" + "abort_on_nan_loss=False" + "abort_on_inf_loss=False" +) +MAXTEXT_ARGS="${MAXTEXT_ARGS_ARRAY[*]}" + +USER_TOKEN=$(gcloud auth application-default print-access-token 2>/dev/null || gcloud auth print-access-token 2>/dev/null || true) + +# The command to run inside the container +RUN_COMMAND="set -e && \ +export LIBTPU_INIT_ARGS=\"${XLA_FLAGS}\" && \ +export JAX_PLATFORMS='tpu,cpu' && \ +export ENABLE_PJRT_COMPATIBILITY='true' && \ +export JAX_DISTRIBUTED_INITIALIZE_TIMEOUT=1800 && \ +export PYTHONPATH=/deps:/deps/src:/deps/src/maxtext/src && \ +export CLOUDSDK_AUTH_ACCESS_TOKEN='${USER_TOKEN}' && \ +python3 src/maxtext/trainers/pre_train/train.py src/maxtext/configs/base.yml ${MAXTEXT_ARGS} && \ +(python3 -c 'import os, glob; from google.cloud import storage; import google.oauth2.credentials; token = os.environ.get(\"CLOUDSDK_AUTH_ACCESS_TOKEN\"); client = storage.Client(credentials=google.oauth2.credentials.Credentials(token), project=\"tpu-prod-env-one-vm\") if token else storage.Client(project=\"tpu-prod-env-one-vm\"); bucket = client.bucket(\"muskansh-data\"); [bucket.blob(f\"qwen3-next-80b-profiles/{os.path.relpath(p, \"/tmp/qwen3-next-80b-profiles\")}\").upload_from_filename(p) for p in glob.glob(\"/tmp/qwen3-next-80b-profiles/**/*\", recursive=True) if os.path.isfile(p)]' || true) && \ +sleep 3600" + +# --- XPK Workload Creation --- +echo "Creating XPK workload: ${WORKLOAD_NAME} on cluster: ${CLUSTER_NAME}" + +PYTHONPATH=/usr/local/google/home/muskansh/xpk/src python3 -P -m xpk.main workload create \ + --cluster="${CLUSTER_NAME}" \ + --project="${PROJECT_ID}" \ + --zone="${ZONE}" \ + --priority="${PRIORITY}" \ + --max-restarts="${MAX_RESTARTS}" \ + --device-type="${DEVICE_TYPE}" \ + --num-slices="${NUM_SLICES}" \ + --docker-image="${WORKLOAD_IMAGE}" \ + --enable-debug-logs \ + --workload="${WORKLOAD_NAME}" \ + --command="${RUN_COMMAND}" + +LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22${ZONE}%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}" +GKE_URL="https://console.cloud.google.com/kubernetes/service/${ZONE}/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}" + +echo "========================================================================" +echo "πŸ“‹ Pantheon Cloud Logging (Worker 0 Logs):" +echo "${LOGS_URL}" +echo "" +echo "☸️ GKE Workload Details:" +echo "${GKE_URL}" +echo "" +echo "========================================================================"