diff --git a/qwen3.8-27b_handover.md b/qwen3.8-27b_handover.md new file mode 100644 index 0000000000..a8fbc8ca11 --- /dev/null +++ b/qwen3.8-27b_handover.md @@ -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. diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py index 5dc3315d77..642fd30054 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py @@ -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, diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 40e2508e58..8a09885a38 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -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, diff --git a/src/maxtext/configs/models/qwen3.8-27b.yml b/src/maxtext/configs/models/qwen3.8-27b.yml new file mode 100644 index 0000000000..d5c738dcb3 --- /dev/null +++ b/src/maxtext/configs/models/qwen3.8-27b.yml @@ -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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 353656c665..35fe060b23 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -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", diff --git a/src/maxtext/inference/kvcache.py b/src/maxtext/inference/kvcache.py index e8475e5e33..3c7b492af4 100644 --- a/src/maxtext/inference/kvcache.py +++ b/src/maxtext/inference/kvcache.py @@ -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( diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py index dc24527558..6cd0f09bec 100644 --- a/src/maxtext/models/qwen3_5.py +++ b/src/maxtext/models/qwen3_5.py @@ -25,6 +25,7 @@ from flax import nnx from maxtext.common.common_types import Config, Array +from maxtext.layers.linears import MlpBlock from maxtext.layers import initializers as max_initializers from maxtext.layers import nnx_wrappers from maxtext.layers.normalizations import Qwen3NextRMSNorm @@ -168,7 +169,7 @@ def __init__( config=cfg, inputs_shape=dummy_inputs_shape, mesh=self.mesh, dtype=cfg.dtype, model_mode=model_mode, rngs=rngs ) - # Second LayerNorm, applied before the MoE block. + # Second LayerNorm, applied before the MoE/MLP block. self.post_attention_layernorm = Qwen3NextRMSNorm( num_features=cfg.emb_dim, epsilon=cfg.normalization_layer_epsilon, @@ -177,8 +178,23 @@ def __init__( rngs=rngs, ) - # Instantiate our `Qwen3_5SparseMoEBlock`. - self.mlp = Qwen3_5SparseMoEBlock(config=cfg, mesh=self.mesh, quant=self.quant, rngs=rngs) + # Conditionally instantiate MoE or dense MLP. + if getattr(cfg, "num_experts", 1) > 1: + self.mlp = Qwen3_5SparseMoEBlock(config=cfg, mesh=self.mesh, quant=self.quant, rngs=rngs) + else: + self.mlp = MlpBlock( + in_features=cfg.emb_dim, + intermediate_dim=cfg.mlp_dim, + activations=cfg.mlp_activations, + intermediate_dropout_rate=cfg.dropout_rate, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + config=cfg, + mesh=mesh, + quant=quant, + model_mode=model_mode, + rngs=rngs, + ) def __call__( self, @@ -226,26 +242,26 @@ def __call__( hidden_states = residual + attention_output hidden_states = nn.with_logical_constraint(hidden_states, self.activation_axis_names) - # Prepare for the MoE block by capturing the new residual + # Prepare for the MoE/MLP block by capturing the new residual residual = hidden_states - # Second LayerNorm, applied before the MoE block. + # Second LayerNorm, applied before the MoE/MLP block. hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = nn.with_logical_constraint(hidden_states, self.activation_axis_names) - # Instantiate and call our `Qwen3_5SparseMoEBlock`. - mlp_output, load_balance_loss = self.mlp( - hidden_states, - deterministic=deterministic, - forced_routed_experts=forced_routed_experts, - ) - - # We sow the load balancing loss so it can be collected and added to the total loss - # during training. - if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: - self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss) + # Instantiate and call our MLP / MoE block. + if getattr(self.config, "num_experts", 1) > 1: + mlp_output, load_balance_loss = self.mlp( + hidden_states, + deterministic=deterministic, + forced_routed_experts=forced_routed_experts, + ) + if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: + self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss) + else: + mlp_output = self.mlp(hidden_states, deterministic=deterministic) - # Final residual connection (after the MoE block) + # Final residual connection (after the MoE/MLP block) layer_output = residual + mlp_output layer_output = nn.with_logical_constraint( layer_output, diff --git a/src/maxtext/utils/globals.py b/src/maxtext/utils/globals.py index 0c73e6f020..b4639c2f97 100644 --- a/src/maxtext/utils/globals.py +++ b/src/maxtext/utils/globals.py @@ -85,6 +85,7 @@ "qwen3-next-80b-a3b": "Qwen/Qwen3-Next-80B-A3B-Instruct", "qwen3.5-397b-a17b": "Qwen/Qwen3.5-397B-A17B", "qwen3.5-35b-a3b": "Qwen/Qwen3.5-35B-A3B", + "qwen3.8-27b": "Qwen/Qwen3.8-27B", "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1", "mistral-7b": "mistralai/Mistral-7B-v0.1", "mixtral-8x22b": "mistralai/Mixtral-8x22B-Instruct-v0.1",