Merge pull request #2428 from AI-Hypercomputer:rbierneni-qwen3-dense-checkpoint

PiperOrigin-RevId: 818747460

Author: Google-ML-Automation

src/MaxText/configs/models/qwen3-4b-thinking-2507.yml

@@ -0,0 +1,36 @@
+# Copyright 2023–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.
+
+# model config for qwen3-4b-Thinking-2507
+
+base_emb_dim: 2560
+base_num_query_heads: 32
+base_num_kv_heads: 8
+base_mlp_dim: 9728
+base_num_decoder_layers: 36
+head_dim: 128
+mlp_activations: ["silu", "linear"] # "hidden_act": "silu" implies SwiGLU
+vocab_size: 151936
+
+decoder_block: "qwen3"
+
+normalization_layer_epsilon: 1.0e-6
+rope_max_timescale: 5000000
+
+use_qk_norm: True
+
+logits_via_embedding: True # from "tie_word_embeddings": true
+normalize_embedding_logits: False
+enable_dropout: False # deterministic for testing
+

src/MaxText/pyconfig.py

@@ -384,6 +384,7 @@ def validate_model_name(s: str) -> bool:
       "gemma3-27b",
       "qwen3-0.6b",
       "qwen3-4b",
+      "qwen3-4b-thinking-2507",
       "qwen3-8b",
       "qwen3-14b",
       "qwen3-32b",

src/MaxText/utils/ckpt_conversion/utils/hf_model_configs.py

@@ -478,6 +478,7 @@
     "gemma3-27b": gemma3_27b_config,
     "qwen3-0.6b": qwen3_0_6b_config,
     "qwen3-4b": qwen3_4b_config,
+    "qwen3-4b-thinking-2507": qwen3_4b_config,
     "qwen3-8b": qwen3_8b_config,
     "qwen3-14b": qwen3_14b_config,
     "qwen3-32b": qwen3_32b_config,

src/MaxText/utils/ckpt_conversion/utils/hf_shape.py

@@ -371,6 +371,7 @@ def LLAMA31_HF_WEIGHTS_TO_SHAPE(config):
     "gemma3-27b": GEMMA3_HF_WEIGHTS_TO_SHAPE,
     "qwen3-0.6b": QWEN3_HF_WEIGHTS_TO_SHAPE,
     "qwen3-4b": QWEN3_HF_WEIGHTS_TO_SHAPE,
+    "qwen3-4b-thinking-2507": QWEN3_HF_WEIGHTS_TO_SHAPE,
     "qwen3-8b": QWEN3_HF_WEIGHTS_TO_SHAPE,
     "qwen3-14b": QWEN3_HF_WEIGHTS_TO_SHAPE,
     "qwen3-32b": QWEN3_HF_WEIGHTS_TO_SHAPE,

src/MaxText/utils/ckpt_conversion/utils/param_mapping.py

@@ -997,6 +997,7 @@ def from_hf():
     "gemma3-27b": GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING,
     "qwen3-0.6b": QWEN3_MAXTEXT_TO_HF_PARAM_MAPPING,
     "qwen3-4b": QWEN3_MAXTEXT_TO_HF_PARAM_MAPPING,
+    "qwen3-4b-thinking-2507": QWEN3_MAXTEXT_TO_HF_PARAM_MAPPING,
     "qwen3-8b": QWEN3_MAXTEXT_TO_HF_PARAM_MAPPING,
     "qwen3-14b": QWEN3_MAXTEXT_TO_HF_PARAM_MAPPING,
     "qwen3-32b": QWEN3_MAXTEXT_TO_HF_PARAM_MAPPING,
@@ -1017,6 +1018,7 @@ def from_hf():
     "gemma3-27b": GEMMA3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "qwen3-0.6b": QWEN3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "qwen3-4b": QWEN3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
+    "qwen3-4b-thinking-2507": QWEN3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "qwen3-8b": QWEN3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "qwen3-14b": QWEN3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "qwen3-32b": QWEN3_MAXTEXT_TO_HF_PARAM_HOOK_FN,

src/MaxText/utils/ckpt_conversion/utils/utils.py

@@ -60,6 +60,7 @@
     "gemma3-27b": "google/gemma-3-27b-it",
     "qwen3-0.6b": "Qwen/Qwen3-0.6B",
     "qwen3-4b": "Qwen/Qwen3-4B",
+    "qwen3-4b-thinking-2507": "Qwen/Qwen3-4B-Thinking-2507",
     "qwen3-8b": "Qwen/Qwen3-8B",
     "qwen3-14b": "Qwen/Qwen3-14B",
     "qwen3-32b": "Qwen/Qwen3-32B",