Merge pull request #4039 from AI-Hypercomputer:yujiedeng/repo-fix

PiperOrigin-RevId: 933555873

Author: Google-ML-Automation

src/maxtext/checkpoint_conversion/utils/hf_model_configs.py

@@ -718,6 +718,44 @@
     vocab_size=151936,
 )
 
+# Matches Qwen/Qwen3-30B-A3B-Base
+qwen3_30b_a3b_base_config = transformers.Qwen3MoeConfig(
+    architectures=["Qwen3MoeForCausalLM"],
+    attention_bias=False,
+    attention_dropout=0.0,
+    bos_token_id=151643,
+    decoder_sparse_step=1,
+    eos_token_id=151643,
+    head_dim=128,
+    hidden_act="silu",
+    hidden_size=2048,
+    initializer_range=0.02,
+    intermediate_size=6144,
+    max_position_embeddings=32768,
+    max_window_layers=48,
+    mlp_only_layers=[],
+    model_type="qwen3_moe",
+    moe_intermediate_size=768,
+    norm_topk_prob=True,
+    num_attention_heads=32,
+    num_experts=128,
+    num_experts_per_tok=8,
+    num_hidden_layers=48,
+    num_key_value_heads=4,
+    output_router_logits=False,
+    rms_norm_eps=1e-06,
+    rope_scaling=None,
+    rope_theta=1000000,
+    router_aux_loss_coef=0.001,
+    sliding_window=None,
+    tie_word_embeddings=False,
+    torch_dtype="bfloat16",
+    transformers_version="4.51.0",
+    use_cache=True,
+    use_sliding_window=False,
+    vocab_size=151936,
+)
+
 qwen3_235b_a22b_thinking_2507_config = transformers.Qwen3MoeConfig(
     architectures=["Qwen3MoeForCausalLM"],
     attention_bias=False,
@@ -1579,7 +1617,7 @@ def __init__(self, **kwargs):
     "llama3.1-70b": llama31_70b_config,
     "llama3.1-405b": llama31_405b_config,
     "qwen3-30b-a3b": qwen3_30b_a3b_thinking_2507_config,
-    "qwen3-30b-a3b-base": qwen3_30b_a3b_thinking_2507_config,
+    "qwen3-30b-a3b-base": qwen3_30b_a3b_base_config,
     "qwen3-235b-a22b": qwen3_235b_a22b_thinking_2507_config,
     "qwen3-480b-a35b": qwen3_coder_480b_a35b_config,
     "deepseek2-16b": deepseek2_16b_config,

src/maxtext/configs/models/qwen3-30b-a3b-base.yml

@@ -34,7 +34,7 @@ base_moe_mlp_dim: 768
 norm_topk_prob: true
 
 # RoPE Settings
-rope_max_timescale: 10_000_000
+rope_max_timescale: 1_000_000
 
 # General Model Settings
 enable_dropout: false