Merge pull request #2487 from AI-Hypercomputer:parambole/maxtext_qwen3_next_v1

PiperOrigin-RevId: 821851030

Author: Google-ML-Automation

src/MaxText/common_types.py

@@ -86,6 +86,7 @@ class DecoderBlockType(enum.Enum):
   GEMMA3 = "gemma3"
   QWEN3 = "qwen3"
   QWEN3_MOE = "qwen3_moe"
+  QWEN3_NEXT = "qwen3_next"
   GPT3 = "gpt3"
   GPT_OSS = "gpt_oss"
   SIMPLE = "simple"

src/MaxText/configs/base.yml

@@ -860,3 +860,19 @@ subslice_shape: ""
 # NNX
 enable_nnx: false
 shard_optimizer_over_data: False
+
+################################## Qwen3-Next Specific Configs ##################################
+# Kernel size for the 1D convolution in the Gated Delta Net
+gdn_conv_kernel_dim: 4
+# Head dimension for the key/query in the Gated Delta Net
+gdn_key_head_dim: 128
+# Head dimension for the value in the Gated Delta Net
+gdn_value_head_dim: 128
+# Number of key/query heads in the Gated Delta Net
+gdn_num_key_heads: 16
+# Number of value heads in the Gated Delta Net
+gdn_num_value_heads: 32
+# Chunk size for the parallel scan algorithm in the Gated Delta Net.
+gdn_chunk_size: 64
+# Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.
+use_qk_norm_in_gdn: True

src/MaxText/configs/models/qwen3-next-80b-a3b.yml

@@ -0,0 +1,47 @@
+# 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.yml
+
+# Set the decoder block to our new implementation
+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
+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
+num_experts_per_tok: 10
+norm_topk_prob: True
+
+# Qwen3-Next Specific Parameters for Linear Attention (Gated Delta Net)
+inhomogeneous_layer_cycle_interval: 4
+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
+
+# RoPE Settings
+rope_max_timescale: 10000000

src/MaxText/layers/decoders.py

@@ -401,6 +401,8 @@ def get_decoder_layers(self):
         return [qwen3.Qwen3DecoderLayerToLinen]
       case DecoderBlockType.QWEN3_MOE:
         return [qwen3.Qwen3MoeDecoderLayerToLinen]
+      case DecoderBlockType.QWEN3_NEXT:
+        return [qwen3.Qwen3NextScannableBlockToLinen] if self.config.scan_layers else [qwen3.Qwen3NextDecoderLayerToLinen]
       case DecoderBlockType.SIMPLE:
         return [simple_layer.SimpleDecoderLayerToLinen]
       case DecoderBlockType.SIMPLE_MLP:
@@ -452,6 +454,7 @@ def get_norm_layer(self, num_features: int):
         DecoderBlockType.GEMMA3,
         DecoderBlockType.QWEN3,
         DecoderBlockType.QWEN3_MOE,
+        DecoderBlockType.QWEN3_NEXT,
         DecoderBlockType.GPT_OSS,
         DecoderBlockType.SIMPLE,
         DecoderBlockType.SIMPLE_MLP,
@@ -824,6 +827,8 @@ def __call__(
                   "is_nope_layer": llama4.determine_is_nope_layer(lyr, self.config.nope_layer_interval),
                   "is_moe_layer": llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step),
               }
+            if cfg.decoder_block == DecoderBlockType.QWEN3_NEXT:
+              layer_kwargs = {"layer_idx": lyr}
             if cfg.decoder_block == DecoderBlockType.GPT_OSS:
               layer_kwargs = {"attention_type": gpt_oss.get_attention_type(layer_id=lyr)}
             layer = RemattedBlockLayer(

src/MaxText/layers/normalizations.py

@@ -25,6 +25,7 @@
 from MaxText import max_utils
 from MaxText.layers import nnx_wrappers
 from MaxText.layers.initializers import Initializer, variable_to_logically_partitioned
+from MaxText.common_types import Array
 
 
 class RMSNorm(nnx.Module):
@@ -93,3 +94,19 @@ def rms_norm(
       metadata_fn=variable_to_logically_partitioned,
   )
   return module
+
+
+def l2norm(x: Array, dim: int = -1, eps: float = 1e-6) -> Array:
+  """L2 normalization function. Normalizes a vector to have a length of 1.
+
+  Args:
+    x: Input array.
+    dim: The axis or axes along which to normalize. Defaults to the last axis.
+    eps: Small epsilon to prevent division by zero.
+
+  Returns:
+    L2 normalized array with the same shape as x.
+  """
+
+  inv_norm = jax.lax.rsqrt((x * x).sum(axis=dim, keepdims=True) + jnp.array(eps, dtype=x.dtype))
+  return x * inv_norm