Cherry Pick Qwen3-Next Model Provider + Pretraining Support (#1027)

docs/models/llm/qwen.md

@@ -1,6 +1,7 @@
 # Qwen
 
-Qwen2/2.5/3 models are supported via the Bridge with QK layernorm handling (Qwen3) and bias in QKV (Qwen2).
+Qwen2/2.5/3/3-Next models are supported via the Bridge with QK layernorm handling (Qwen3) and bias in QKV (Qwen2).
+Qwen3-Next models also support Gated-Delta Networks, Zero-Centered RMSNorm and QK Layernorm weight decay during training.
 
 ## Conversion with 🤗 Hugging Face
 
@@ -47,6 +48,7 @@ cfg = qwen3_8b_pretrain_config(
   - Qwen2 recipes: [bridge.recipes.qwen.qwen2](../../apidocs/bridge/bridge.recipes.qwen.qwen2.md)
   - Qwen3 recipes: [bridge.recipes.qwen.qwen3](../../apidocs/bridge/bridge.recipes.qwen.qwen3.md)
   - Qwen3 MoE recipes: [bridge.recipes.qwen.qwen3_moe](../../apidocs/bridge/bridge.recipes.qwen.qwen3_moe.md)
+  - Qwen3-Next recipes: [bridge.recipes.qwen.qwen3_next](../../apidocs/bridge/bridge.recipes.qwen.qwen3_next.md)
 
 ## Finetuning recipes
 - Coming soon
@@ -55,6 +57,7 @@ cfg = qwen3_8b_pretrain_config(
 - Qwen2: `https://huggingface.co/Qwen/Qwen2-7B`
 - Qwen2.5: `https://huggingface.co/Qwen/Qwen2.5-7B`
 - Qwen3: `https://huggingface.co/Qwen/Qwen3-7B`
+- Qwen3-Next `https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct` and `https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Thinking`
 
 ## Related docs
 - Recipe usage and customization: [Recipe usage](../../recipe-usage.md)

src/megatron/bridge/models/qwen/qwen_provider.py

@@ -15,15 +15,23 @@
 import logging
 from dataclasses import dataclass
 from functools import partial
-from typing import TYPE_CHECKING, Callable, List, Optional, Union
+from typing import TYPE_CHECKING, Callable, Optional
 
 import torch
 import torch.nn.functional as F
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
+from megatron.core.transformer.spec_utils import ModuleSpec
 
 from megatron.bridge.models.gpt_provider import GPTModelProvider
 
 
+try:
+    import transformer_engine  # type: ignore  # noqa: F401
+
+    HAVE_TE = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_TE = False
+
 try:
     import transformer_engine  # type: ignore  # noqa: F401
 
@@ -425,42 +433,42 @@ class Qwen3MoEModelProvider235B_A22B(Qwen3MoEModelProvider):
 class Qwen3NextModelProvider(Qwen3MoEModelProvider):
     """Base provider for Qwen 3 Next Models."""
 
-    transformer_layer_spec: Union["ModuleSpec", Callable[["GPTModelProvider"], "ModuleSpec"]] = partial(
+    transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = partial(
         get_gpt_decoder_block_spec, use_transformer_engine=HAVE_TE
     )
 
-    normalization: str = "RMSNorm"
-    layernorm_zero_centered_gamma: bool = True  # Zero-Centered RMSNorm
-    rotary_percent: float = 0.25  # Qwen3-Next applies RoPE only to the first 25% of dims
-    attention_output_gate: bool = True  # Qwen3-Next applies output gate to standard attention
-    kv_channels: Optional[int] = 256
-    seq_length: int = 262144  # Qwen3-Next-80B-A3B supports up to 256k sequence length
+    layernorm_zero_centered_gamma: bool = True  # Zero-centered RMSNorm
+    kv_channels: int | None = 256
+    num_query_groups: int = 2
+    seq_length: int = 262144  # 256k tokens
     rotary_base: float = 10000000.0
+    rotary_percent: float = 0.25  # 25% of the hidden size is used for RoPE
+    attention_output_gate: bool = True  # Gated Attention
 
-    # GDN specific parameters
-    linear_attention_type: str = "gated_delta_net"
-    linear_attention_freq: Union[int, List[int]] = 4  # 1 standard attention layer per 4 layers
+    # MoE specific parameters
+    num_moe_experts: int = 512
+    moe_router_topk: int = 10  # 10 routed experts per token
+    moe_shared_expert_gate: bool = True  # Qwen3-Next uses a gate for the shared expert
+    moe_router_dtype: str = "fp32"
+    moe_router_load_balancing_type: str = "global_aux_loss"  # Qwen3-Next uses global aux loss for load balancing
+
+    # Linear Attention specific parameters
+    linear_attention_type: str = "gated_delta_net"  # Gated Delta Net used in 75% of the model layers
+    linear_attention_freq: int | list[int] = 4  # 1 gated standard attention layer per 4 layers
     linear_conv_kernel_dim: int = 4
     linear_key_head_dim: int = 128
     linear_value_head_dim: int = 128
     linear_num_key_heads: int = 16
     linear_num_value_heads: int = 32
 
-    # MoE specific parameters
-    num_moe_experts: int = 512
-    moe_router_topk: int = 10
-    moe_ffn_hidden_size: int = 512
-    moe_shared_expert_intermediate_size: int = 512  # 512 * 1 shared expert
-    moe_shared_expert_gate: bool = True
-
     # Checkpointing
     hetereogenous_dist_checkpoint: bool = True
 
 
 @dataclass
 class Qwen3NextModelProvider80B_A3B(Qwen3NextModelProvider):
     """
-    Provider for Qwen 3 Next: https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct
+    Provider for Qwen 3 Next 80B-A3B: https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct and https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Thinking
     """
 
     num_layers: int = 48
@@ -469,4 +477,6 @@ class Qwen3NextModelProvider80B_A3B(Qwen3NextModelProvider):
     num_query_groups: int = 2
     ffn_hidden_size: int = 5120
     moe_ffn_hidden_size: int = 512
-    mtp_num_layers: int = 0  # Set to 1 if need MTP
+    moe_shared_expert_intermediate_size: int = 512
+    mtp_num_layers: Optional[int] = 0
+    mtp_loss_scaling_factor: Optional[float] = None

src/megatron/bridge/recipes/qwen/__init__.py

@@ -49,6 +49,11 @@
     qwen3_235b_a22b_pretrain_config,
 )
 
+# Qwen3-Next models
+from .qwen3_next import (
+    qwen3_next_80b_a3b_pretrain_config,
+)
+
 
 __all__ = [
     # Qwen2 models
@@ -79,4 +84,6 @@
     # Qwen3 MoE models
     "qwen3_30b_a3b_pretrain_config",
     "qwen3_235b_a22b_pretrain_config",
+    # Qwen3-Next models
+    "qwen3_next_80b_a3b_pretrain_config",
 ]