AI-Hypercomputer · hengtaoguo · Sep 4, 2026
@@ -0,0 +1,66 @@
+# Qwen3.8-27B (Text-only) MaxText Onboarding Handover
+
+## 1. Summary
+The text-only architecture of **Qwen/Qwen3.8-27B** has been successfully onboarded to MaxText on branch `hengtaoguo-exp5` and thoroughly verified on TPU v5p (hengtaoguo-dev-v5p2 / 34.32.131.57).
+
+Both validation milestones are **PASSED**:
+1. **Forward Pass Logits Checker** (`tests/utils/forward_pass_logit_checker.py`): **PASSED** (Top-10 Overlap: 10/10, Jaccard similarity: 1.0, Rank agreement: 100%, Average KL divergence: < 1e-3).
+2. **16-Token Autoregressive Decoding** (`maxtext.inference.decode`): **PASSED** (100% exact token ID match against HuggingFace PyTorch reference).
+
+---
+
+## 2. Model Architecture & Hyperparameters
+- **Base Architecture**: Hybrid Attention model with Gated Delta Network (GDN) linear attention and full attention in a 4-layer cycle (`inhomogeneous_layer_cycle_interval: 4`: 3 GDN layers + 1 Full Attention layer).
+- **MLP Type**: Dense SwiGLU MLP (`num_experts: None` / 1).
+- **Key Dimensions**:
+  - `vocab_size`: 248,320
+  - `emb_dim`: 5,120 (`hidden_size`)
+  - `mlp_dim`: 17,408 (`intermediate_size`)
+  - `num_decoder_layers`: 64 (`num_hidden_layers`)
+  - `num_query_heads`: 24 (`num_attention_heads`)
+  - `num_kv_heads`: 4 (`num_key_value_heads`)
+  - `head_dim`: 256
+  - `partial_rotary_factor`: 0.25 (rotary dim = 64)
+  - `rope_theta`: 10,000,000.0
+  - `normalization_layer_epsilon`: 1e-06
+  - `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
+  - `use_qk_norm_in_gdn`: True
+
+---
+
+## 3. Key Changes Made
+1. **Model Config**:
+   - Added `src/maxtext/configs/models/qwen3.8-27b.yml` with exact architecture parameters.
+   - Added `qwen3.8-27b` to `ModelName` in `src/maxtext/configs/types.py`.
+   - Added `qwen3.8-27b: Qwen/Qwen3.8-27B` to `HF_IDS` in `src/maxtext/utils/globals.py`.
+   - Registered `qwen3_8_27b_config` in `src/maxtext/checkpoint_conversion/utils/hf_model_configs.py`.
+2. **Model Definition**:
+   - Updated `src/maxtext/models/qwen3_5.py` (`Qwen3_5DecoderLayer`) to support dense `MlpBlock` when `getattr(cfg, num_experts, 1) <= 1`.
+3. **Checkpoint Conversion**:
+   - Updated `src/maxtext/checkpoint_conversion/utils/param_mapping.py` to support dense MLP weight mapping and hooks for Qwen3.5/Qwen3.8 models.
+4. **KV Cache & Inference Fix**:
+   - Fixed `src/maxtext/inference/kvcache.py` to initialize `cached_prefill_key` with `jnp.float32` for GDN recurrent state, preventing dtype mismatch during decode insert.
+
+---
+
+## 4. Verification Commands & Results
+
+### Sliced Mini-Checkpoint Location
+- HuggingFace 4-layer mini model: `/dev/shm/hf_mini/qwen3.8-27b_4layers`
+- Converted Orbax checkpoint: `/dev/shm/hengtaoguo/checkpoints/qwen3.8-27b_mini_orbax/0/items`
+
+### 1. Forward Pass Logits Checker
+```bash
+python3 -m tests.utils.forward_pass_logit_checker src/maxtext/configs/base.yml   tokenizer_path=/dev/shm/hf_mini/qwen3.8-27b_4layers   load_parameters_path=/dev/shm/hengtaoguo/checkpoints/qwen3.8-27b_mini_orbax/0/items   model_name=qwen3.8-27b   base_num_decoder_layers=4   override_model_config=true   per_device_batch_size=1   scan_layers=false   dtype=bfloat16   weight_dtype=bfloat16   attention=dot_product   max_prefill_predict_length=16   max_target_length=16   --run_hf_model=true   --hf_model_path=/dev/shm/hf_mini/qwen3.8-27b_4layers   --max_kl_div=0.01
+```
+**Result**: 4/4 prompts evaluated. Top-10 overlap 10/10 (100%), Jaccard similarity 1.0, KL divergence ~ 5e-4 to 1e-3.
+
+### 2. Autoregressive Decode (16 tokens)
+```bash
+python3 -m maxtext.inference.decode src/maxtext/configs/base.yml   model_name=qwen3.8-27b   tokenizer_path=/dev/shm/hf_mini/qwen3.8-27b_4layers   tokenizer_type=huggingface   load_parameters_path=/dev/shm/hengtaoguo/checkpoints/qwen3.8-27b_mini_orbax/0/items   run_name=qwen3-8-27b-mini-decode   base_num_decoder_layers=4   override_model_config=true   per_device_batch_size=1   max_prefill_predict_length=16   max_target_length=32   steps=1   scan_layers=false   dtype=bfloat16   weight_dtype=bfloat16   attention=dot_product   prompt='I love to'
+```
+**Result**: Generated 16 tokens matching the HuggingFace PyTorch reference with 100% token-by-token parity.
@@ -1609,6 +1609,63 @@ def __init__(self, **kwargs):
   qwen3_5_397b_a17b_config = PTConfig(**qwen3_5_397b_a17b_dict)  # pytype: disable=wrong-arg-types
 
 
+qwen3_8_27b_dict = {
+    "architectures": ["Qwen3_5ForConditionalGeneration"],
+    "model_type": "qwen3_5",
+    "text_config": {
+        "attention_bias": False,
+        "attention_dropout": 0.0,
+        "attn_output_gate": True,
+        "dtype": "bfloat16",
+        "eos_token_id": 248044,
+        "full_attention_interval": 4,
+        "head_dim": 256,
+        "hidden_act": "silu",
+        "hidden_size": 5120,
+        "initializer_range": 0.02,
+        "intermediate_size": 17408,
+        "layer_types": [
+            "linear_attention",
+            "linear_attention",
+            "linear_attention",
+            "full_attention",
+        ]
+        * 16,
+        "linear_conv_kernel_dim": 4,
+        "linear_key_head_dim": 128,
+        "linear_num_key_heads": 16,
+        "linear_num_value_heads": 48,
+        "linear_value_head_dim": 128,
+        "mamba_ssm_dtype": "float32",
+        "max_position_embeddings": 262144,
+        "model_type": "qwen3_5_text",
+        "mtp_num_hidden_layers": 1,
+        "mtp_use_dedicated_embeddings": False,
+        "num_attention_heads": 24,
+        "num_hidden_layers": 64,
+        "num_key_value_heads": 4,
+        "output_gate_type": "swish",
+        "partial_rotary_factor": 0.25,
+        "rms_norm_eps": 1e-06,
+        "rope_parameters": {
+            "mrope_interleaved": True,
+            "mrope_section": [11, 11, 10],
+            "partial_rotary_factor": 0.25,
+            "rope_theta": 10000000,
+            "rope_type": "default",
+        },
+        "tie_word_embeddings": False,
+        "use_cache": True,
+        "vocab_size": 248320,
+    },
+}
+
+try:
+  qwen3_8_27b_config = transformers.Qwen3_5Config(**qwen3_8_27b_dict)  # pyrefly: ignore[missing-attribute]
+except (AttributeError, TypeError):
+  qwen3_8_27b_config = PTConfig(**qwen3_8_27b_dict)  # pytype: disable=wrong-arg-types
+
+
 # from https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1/blob/main/config.json
 mixtral_8x7b_dict = {
     "architectures": ["MixtralForCausalLM"],
@@ -1982,6 +2039,7 @@ def __init__(self, **kwargs):
     "qwen3-next-80b-a3b": qwen3_next_80b_a3b_config,
     "qwen3.5-397b-a17b": qwen3_5_397b_a17b_config,
     "qwen3.5-35b-a3b": qwen3_5_35b_a3b_config,
+    "qwen3.8-27b": qwen3_8_27b_config,
     "mixtral-8x7b": mixtral_8x7b_config,
     "mixtral-8x22b": mixtral_8x22b_config,
     "olmo3-7b": olmo3_7b_config,

@@ -873,6 +873,9 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
   num_main_layers = config["text_config"]["num_hidden_layers"]
   layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval
 
+  text_cfg = config.get("text_config", config)
+  num_experts = text_cfg.get("num_experts", text_cfg.get("num_local_experts", 0)) or 0
+
   # 1. Non-layer specific weight mappings
   mapping = {
       "params-token_embedder-embedding": "model.language_model.embed_tokens.weight",
@@ -956,38 +959,43 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
             }
         )
 
-      # 3. Handle MLP: Gates and Shared Experts
-      mapping.update(  # pyrefly: ignore[no-matching-overload]
-          {
-              f"{prefix}-mlp-routed_experts-gate-kernel": [
-                  f"model.language_model.layers.{i}.mlp.gate.weight" for i in hf_indices
-              ],
-              f"{prefix}-mlp-shared_expert-wi_0-kernel": [
-                  f"model.language_model.layers.{i}.mlp.shared_expert.gate_proj.weight" for i in hf_indices
-              ],
-              f"{prefix}-mlp-shared_expert-wi_1-kernel": [
-                  f"model.language_model.layers.{i}.mlp.shared_expert.up_proj.weight" for i in hf_indices
-              ],
-              f"{prefix}-mlp-shared_expert-wo-kernel": [
-                  f"model.language_model.layers.{i}.mlp.shared_expert.down_proj.weight" for i in hf_indices
-              ],
-              f"{prefix}-mlp-shared_expert_gate-kernel": [
-                  f"model.language_model.layers.{i}.mlp.shared_expert_gate.weight" for i in hf_indices
-              ],
-          }
-      )
-
-      # 4. Handle MoE Routed Experts
-      mapping.update(  # pyrefly: ignore[no-matching-overload]
-          {
-              f"{prefix}-mlp-routed_experts-wo": [
-                  f"model.language_model.layers.{i}.mlp.experts.down_proj" for i in hf_indices
-              ],
-              (f"{prefix}-mlp-routed_experts-wi_0", f"{prefix}-mlp-routed_experts-wi_1"): [
-                  f"model.language_model.layers.{i}.mlp.experts.gate_up_proj" for i in hf_indices
-              ],
-          }
-      )
+      # 3. Handle MLP
+      if num_experts > 1:
+        mapping.update(  # pyrefly: ignore[no-matching-overload]
+            {
+                f"{prefix}-mlp-routed_experts-gate-kernel": [
+                    f"model.language_model.layers.{i}.mlp.gate.weight" for i in hf_indices
+                ],
+                f"{prefix}-mlp-shared_expert-wi_0-kernel": [
+                    f"model.language_model.layers.{i}.mlp.shared_expert.gate_proj.weight" for i in hf_indices
+                ],
+                f"{prefix}-mlp-shared_expert-wi_1-kernel": [
+                    f"model.language_model.layers.{i}.mlp.shared_expert.up_proj.weight" for i in hf_indices
+                ],
+                f"{prefix}-mlp-shared_expert-wo-kernel": [
+                    f"model.language_model.layers.{i}.mlp.shared_expert.down_proj.weight" for i in hf_indices
+                ],
+                f"{prefix}-mlp-shared_expert_gate-kernel": [
+                    f"model.language_model.layers.{i}.mlp.shared_expert_gate.weight" for i in hf_indices
+                ],
+                f"{prefix}-mlp-routed_experts-wo": [
+                    f"model.language_model.layers.{i}.mlp.experts.down_proj" for i in hf_indices
+                ],
+                (f"{prefix}-mlp-routed_experts-wi_0", f"{prefix}-mlp-routed_experts-wi_1"): [
+                    f"model.language_model.layers.{i}.mlp.experts.gate_up_proj" for i in hf_indices
+                ],
+            }
+        )
+      else:
+        mapping.update(  # pyrefly: ignore[no-matching-overload]
+            {
+                f"{prefix}-mlp-wi_0-kernel": [
+                    f"model.language_model.layers.{i}.mlp.gate_proj.weight" for i in hf_indices
+                ],
+                f"{prefix}-mlp-wi_1-kernel": [f"model.language_model.layers.{i}.mlp.up_proj.weight" for i in hf_indices],
+                f"{prefix}-mlp-wo-kernel": [f"model.language_model.layers.{i}.mlp.down_proj.weight" for i in hf_indices],
+            }
+        )
   else:
     # Unscanned layer mapping
     for i in range(num_main_layers):
@@ -1035,29 +1043,32 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
             }
         )
 
-      # MLP: Gates and Shared Experts
+      # MLP
       hf_mlp = f"model.language_model.layers.{i}.mlp"
 
-      mapping.update(
-          {
-              f"{prefix}-mlp-routed_experts-gate-kernel": (f"{hf_mlp}.gate.weight"),
-              f"{prefix}-mlp-shared_expert-wi_0-kernel": (f"{hf_mlp}.shared_expert.gate_proj.weight"),
-              f"{prefix}-mlp-shared_expert-wi_1-kernel": (f"{hf_mlp}.shared_expert.up_proj.weight"),
-              f"{prefix}-mlp-shared_expert-wo-kernel": (f"{hf_mlp}.shared_expert.down_proj.weight"),
-              f"{prefix}-mlp-shared_expert_gate-kernel": (f"{hf_mlp}.shared_expert_gate.weight"),
-          }
-      )
-
-      # MoE Routed Experts
-      mapping.update(  # pyrefly: ignore[no-matching-overload]
-          {
-              f"{prefix}-mlp-routed_experts-wo": f"model.language_model.layers.{i}.mlp.experts.down_proj",
-              (
-                  f"{prefix}-mlp-routed_experts-wi_0",
-                  f"{prefix}-mlp-routed_experts-wi_1",
-              ): f"model.language_model.layers.{i}.mlp.experts.gate_up_proj",
-          }
-      )
+      if num_experts > 1:
+        mapping.update(
+            {
+                f"{prefix}-mlp-routed_experts-gate-kernel": (f"{hf_mlp}.gate.weight"),
+                f"{prefix}-mlp-shared_expert-wi_0-kernel": (f"{hf_mlp}.shared_expert.gate_proj.weight"),
+                f"{prefix}-mlp-shared_expert-wi_1-kernel": (f"{hf_mlp}.shared_expert.up_proj.weight"),
+                f"{prefix}-mlp-shared_expert-wo-kernel": (f"{hf_mlp}.shared_expert.down_proj.weight"),
+                f"{prefix}-mlp-shared_expert_gate-kernel": (f"{hf_mlp}.shared_expert_gate.weight"),
+                f"{prefix}-mlp-routed_experts-wo": f"{hf_mlp}.experts.down_proj",
+                (
+                    f"{prefix}-mlp-routed_experts-wi_0",
+                    f"{prefix}-mlp-routed_experts-wi_1",
+                ): f"{hf_mlp}.experts.gate_up_proj",
+            }
+        )
+      else:
+        mapping.update(
+            {
+                f"{prefix}-mlp-wi_0-kernel": f"{hf_mlp}.gate_proj.weight",
+                f"{prefix}-mlp-wi_1-kernel": f"{hf_mlp}.up_proj.weight",
+                f"{prefix}-mlp-wo-kernel": f"{hf_mlp}.down_proj.weight",
+            }
+        )
 
   # Vision mapping for Qwen3.5
   if maxtext_config.use_multimodal and "vision_config" in config:
@@ -1279,16 +1290,22 @@ def concat_ba_and_transpose(input_tensor, target_shape=None):
       hooks[f"{prefix}-attention-conv1d-kernel"] = permute_conv
 
     mlp_prefix = f"{prefix}-mlp"
-    hooks[f"{mlp_prefix}-routed_experts-gate-kernel"] = transpose
-    hooks[f"{mlp_prefix}-shared_expert-wi_0-kernel"] = transpose
-    hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel"] = transpose
-    hooks[f"{mlp_prefix}-shared_expert-wo-kernel"] = transpose
-    hooks[f"{mlp_prefix}-shared_expert_gate-kernel"] = transpose
-    # pyrefly: ignore[unsupported-operation]
-    hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = (
-        process_wi_0_wi_1  # pyrefly: ignore[unsupported-operation]
-    )
-    hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose_expert
+    num_experts = text_cfg.get("num_experts", text_cfg.get("num_local_experts", 0)) or 0
+    if num_experts > 1:
+      hooks[f"{mlp_prefix}-routed_experts-gate-kernel"] = transpose
+      hooks[f"{mlp_prefix}-shared_expert-wi_0-kernel"] = transpose
+      hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel"] = transpose
+      hooks[f"{mlp_prefix}-shared_expert-wo-kernel"] = transpose
+      hooks[f"{mlp_prefix}-shared_expert_gate-kernel"] = transpose
+      # pyrefly: ignore[unsupported-operation]
+      hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = (
+          process_wi_0_wi_1  # pyrefly: ignore[unsupported-operation]
+      )
+      hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose_expert
+    else:
+      hooks[f"{mlp_prefix}-wi_0-kernel"] = transpose
+      hooks[f"{mlp_prefix}-wi_1-kernel"] = transpose
+      hooks[f"{mlp_prefix}-wo-kernel"] = transpose
 
   # Vision hooks for Qwen3.5
   vision_config = config.get("vision_config", None)
@@ -4298,6 +4315,7 @@ def mhc_concat_scale(input_tensors, target_shape=None):
     "qwen3-next-80b-a3b": QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING,
     "qwen3.5-397b-a17b": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING,
     "qwen3.5-35b-a3b": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING,
+    "qwen3.8-27b": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING,
     "mixtral-8x7b": MIXTRAL_MAXTEXT_TO_HF_PARAM_MAPPING,
     "mixtral-8x22b": MIXTRAL_MAXTEXT_TO_HF_PARAM_MAPPING,
     "olmo3-7b": OLMO3_MAXTEXT_TO_HF_PARAM_MAPPING,
@@ -4353,6 +4371,7 @@ def mhc_concat_scale(input_tensors, target_shape=None):
     "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "qwen3.5-397b-a17b": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "qwen3.5-35b-a3b": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN,
+    "qwen3.8-27b": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "qwen3-next-80b-a3b": QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "mixtral-8x7b": MIXTRAL_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "mixtral-8x22b": MIXTRAL_MAXTEXT_TO_HF_PARAM_HOOK_FN,

@@ -0,0 +1,45 @@
+# 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.
+
+
+decoder_block: "qwen3_5"
+
+# Core Architectural Parameters
+base_emb_dim: 5120
+base_num_decoder_layers: 64
+base_num_query_heads: 24
+base_num_kv_heads: 4
+head_dim: 256
+vocab_size: 248320
+normalization_layer_epsilon: 1.0e-6
+
+# Dense MLP Parameters (SwiGLU)
+base_mlp_dim: 17408
+mlp_activations: ["silu", "linear"]
+
+# GatedDeltaNet Specific Parameters for Linear Attention (GDN)
+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: 48
+gdn_chunk_size: 64
+
+# RoPE Settings
+rope_max_timescale: 10000000
+partial_rotary_factor: 0.25
+
+# General Model Settings
+enable_dropout: False
@@ -283,6 +283,7 @@ class ProfilerType(str, Enum):
     "qwen3-custom-30b-a3b",
     "qwen3.5-35b-a3b",
     "qwen3.5-397b-a17b",
+    "qwen3.8-27b",
     "gpt3-175b",
     "gpt3-22b",
     "gpt3-6b",

@@ -405,7 +405,7 @@ def _initialize_prefill_caches(self, model_mode):
       cache_batch_axis_name = CACHE_BATCH_PREFILL if model_mode == MODEL_MODE_PREFILL else CACHE_BATCH
 
       self.cached_prefill_key = nnx.Cache(
-          jnp.zeros((self.batch, self.key_heads, self.key_head_size, self.value_head_size), dtype=dtype),
+          jnp.zeros((self.batch, self.key_heads, self.key_head_size, self.value_head_size), dtype=jnp.float32),
           out_sharding=(cache_batch_axis_name, CACHE_HEADS, None, None),
       )
       self.cached_prefill_value = nnx.Cache(