diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py index 89abd56d4c..54aaea7075 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py @@ -1931,6 +1931,8 @@ 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.5-35b-a3b-fp8": qwen3_5_35b_a3b_config, + "qwen3.5-35b-fp8": qwen3_5_35b_a3b_config, "mixtral-8x7b": mixtral_8x7b_config, "mixtral-8x22b": mixtral_8x22b_config, "olmo3-7b": olmo3_7b_config, diff --git a/src/maxtext/checkpoint_conversion/utils/hf_shape.py b/src/maxtext/checkpoint_conversion/utils/hf_shape.py index 85dd1d6ea0..274ac536d8 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_shape.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_shape.py @@ -1317,6 +1317,8 @@ def DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE(config): "mixtral-8x7b": MIXTRAL_HF_WEIGHTS_TO_SHAPE, "mixtral-8x22b": MIXTRAL_HF_WEIGHTS_TO_SHAPE, "qwen3.5-35b-a3b": QWEN3_5_HF_WEIGHTS_TO_SHAPE, + "qwen3.5-35b-a3b-fp8": QWEN3_5_HF_WEIGHTS_TO_SHAPE, + "qwen3.5-35b-fp8": QWEN3_5_HF_WEIGHTS_TO_SHAPE, "qwen3.5-397b-a17b": QWEN3_5_HF_WEIGHTS_TO_SHAPE, "qwen3-next-80b-a3b": QWEN3_NEXT_HF_WEIGHTS_TO_SHAPE, } diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 94e96173f1..0c9a6125e5 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -912,6 +912,18 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals f"{prefix}-attention-attention-out-kernel": [ f"model.language_model.layers.{i}.self_attn.o_proj.weight" for i in hf_indices ], + f"{prefix}-attention-attention-query-kernel_scale": [ + f"model.language_model.layers.{i}.self_attn.q_proj.weight_scale" for i in hf_indices + ], + f"{prefix}-attention-attention-key-kernel_scale": [ + f"model.language_model.layers.{i}.self_attn.k_proj.weight_scale" for i in hf_indices + ], + f"{prefix}-attention-attention-value-kernel_scale": [ + f"model.language_model.layers.{i}.self_attn.v_proj.weight_scale" for i in hf_indices + ], + f"{prefix}-attention-attention-out-kernel_scale": [ + f"model.language_model.layers.{i}.self_attn.o_proj.weight_scale" for i in hf_indices + ], f"{prefix}-attention-attention-query_norm-scale": [ f"model.language_model.layers.{i}.self_attn.q_norm.weight" for i in hf_indices ], @@ -940,6 +952,20 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals ) for i in hf_indices ], + f"{prefix}-attention-in_proj_qkvz-kernel_scale": [ + ( + f"model.language_model.layers.{i}.linear_attn.in_proj_qkv.weight_scale", + f"model.language_model.layers.{i}.linear_attn.in_proj_z.weight_scale", + ) + for i in hf_indices + ], + f"{prefix}-attention-in_proj_ba-kernel_scale": [ + ( + f"model.language_model.layers.{i}.linear_attn.in_proj_b.weight_scale", + f"model.language_model.layers.{i}.linear_attn.in_proj_a.weight_scale", + ) + for i in hf_indices + ], f"{prefix}-attention-conv1d-kernel": [ f"model.language_model.layers.{i}.linear_attn.conv1d.weight" for i in hf_indices ], @@ -953,6 +979,9 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals f"{prefix}-attention-out_proj-kernel": [ f"model.language_model.layers.{i}.linear_attn.out_proj.weight" for i in hf_indices ], + f"{prefix}-attention-out_proj-kernel_scale": [ + f"model.language_model.layers.{i}.linear_attn.out_proj.weight_scale" for i in hf_indices + ], } ) @@ -974,6 +1003,18 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals 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-shared_expert-wi_0-kernel_scale": [ + f"model.language_model.layers.{i}.mlp.shared_expert.gate_proj.weight_scale" for i in hf_indices + ], + f"{prefix}-mlp-shared_expert-wi_1-kernel_scale": [ + f"model.language_model.layers.{i}.mlp.shared_expert.up_proj.weight_scale" for i in hf_indices + ], + f"{prefix}-mlp-shared_expert-wo-kernel_scale": [ + f"model.language_model.layers.{i}.mlp.shared_expert.down_proj.weight_scale" for i in hf_indices + ], + f"{prefix}-mlp-shared_expert_gate-kernel_scale": [ + f"model.language_model.layers.{i}.mlp.shared_expert_gate.weight_scale" for i in hf_indices + ], } ) @@ -986,6 +1027,12 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals (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 ], + f"{prefix}-mlp-routed_experts-wo-kernel_scale": [ + f"model.language_model.layers.{i}.mlp.experts.down_proj.weight_scale" for i in hf_indices + ], + (f"{prefix}-mlp-routed_experts-wi_0-kernel_scale", f"{prefix}-mlp-routed_experts-wi_1-kernel_scale"): [ + f"model.language_model.layers.{i}.mlp.experts.gate_up_proj.weight_scale" for i in hf_indices + ], } ) else: @@ -1009,6 +1056,10 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals f"{prefix}-attention-attention-key-kernel": f"model.language_model.layers.{i}.self_attn.k_proj.weight", f"{prefix}-attention-attention-value-kernel": f"model.language_model.layers.{i}.self_attn.v_proj.weight", f"{prefix}-attention-attention-out-kernel": f"model.language_model.layers.{i}.self_attn.o_proj.weight", + f"{prefix}-attention-attention-query-kernel_scale": f"model.language_model.layers.{i}.self_attn.q_proj.weight_scale", + f"{prefix}-attention-attention-key-kernel_scale": f"model.language_model.layers.{i}.self_attn.k_proj.weight_scale", + f"{prefix}-attention-attention-value-kernel_scale": f"model.language_model.layers.{i}.self_attn.v_proj.weight_scale", + f"{prefix}-attention-attention-out-kernel_scale": f"model.language_model.layers.{i}.self_attn.o_proj.weight_scale", f"{prefix}-attention-attention-query_norm-scale": f"model.language_model.layers.{i}.self_attn.q_norm.weight", f"{prefix}-attention-attention-key_norm-scale": f"model.language_model.layers.{i}.self_attn.k_norm.weight", } @@ -1027,11 +1078,20 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals f"model.language_model.layers.{i}.linear_attn.in_proj_b.weight", f"model.language_model.layers.{i}.linear_attn.in_proj_a.weight", ), + f"{prefix}-attention-in_proj_qkvz-kernel_scale": ( + f"model.language_model.layers.{i}.linear_attn.in_proj_qkv.weight_scale", + f"model.language_model.layers.{i}.linear_attn.in_proj_z.weight_scale", + ), + f"{prefix}-attention-in_proj_ba-kernel_scale": ( + f"model.language_model.layers.{i}.linear_attn.in_proj_b.weight_scale", + f"model.language_model.layers.{i}.linear_attn.in_proj_a.weight_scale", + ), f"{prefix}-attention-conv1d-kernel": f"model.language_model.layers.{i}.linear_attn.conv1d.weight", f"{prefix}-attention-A_log": f"model.language_model.layers.{i}.linear_attn.A_log", f"{prefix}-attention-dt_bias": f"model.language_model.layers.{i}.linear_attn.dt_bias", f"{prefix}-attention-norm-rms_norm-scale": f"model.language_model.layers.{i}.linear_attn.norm.weight", f"{prefix}-attention-out_proj-kernel": f"model.language_model.layers.{i}.linear_attn.out_proj.weight", + f"{prefix}-attention-out_proj-kernel_scale": f"model.language_model.layers.{i}.linear_attn.out_proj.weight_scale", } ) @@ -1045,6 +1105,10 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals 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-shared_expert-wi_0-kernel_scale": (f"{hf_mlp}.shared_expert.gate_proj.weight_scale"), + f"{prefix}-mlp-shared_expert-wi_1-kernel_scale": (f"{hf_mlp}.shared_expert.up_proj.weight_scale"), + f"{prefix}-mlp-shared_expert-wo-kernel_scale": (f"{hf_mlp}.shared_expert.down_proj.weight_scale"), + f"{prefix}-mlp-shared_expert_gate-kernel_scale": (f"{hf_mlp}.shared_expert_gate.weight_scale"), } ) @@ -1056,6 +1120,11 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals 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", + f"{prefix}-mlp-routed_experts-wo-kernel_scale": f"model.language_model.layers.{i}.mlp.experts.down_proj.weight_scale", + ( + f"{prefix}-mlp-routed_experts-wi_0-kernel_scale", + f"{prefix}-mlp-routed_experts-wi_1-kernel_scale", + ): f"model.language_model.layers.{i}.mlp.experts.gate_up_proj.weight_scale", } ) @@ -1251,6 +1320,13 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): interleaved = np.concatenate([b_r, a_r], axis=1) return interleaved.reshape(-1, b_m.shape[-1]).T + def reshape_scale(input_tensor, target_shape=None): + if target_shape is None: + return input_tensor + if input_tensor.ndim == 2: + return input_tensor.transpose().reshape(target_shape) + return input_tensor.reshape(target_shape) + # Initialize Hooks hooks = { "params-decoder-logits_dense-kernel": transpose, @@ -1272,11 +1348,15 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): if is_full_attention_layer: for key in ["query", "key", "value", "out"]: hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attention-attention-{key}-kernel_scale"] = reshape_scale else: hooks[f"{prefix}-attention-in_proj_qkvz-kernel"] = concat_qkvz_and_transpose hooks[f"{prefix}-attention-in_proj_ba-kernel"] = concat_ba_and_transpose hooks[f"{prefix}-attention-out_proj-kernel"] = transpose hooks[f"{prefix}-attention-conv1d-kernel"] = permute_conv + hooks[f"{prefix}-attention-in_proj_qkvz-kernel_scale"] = reshape_scale + hooks[f"{prefix}-attention-in_proj_ba-kernel_scale"] = reshape_scale + hooks[f"{prefix}-attention-out_proj-kernel_scale"] = reshape_scale mlp_prefix = f"{prefix}-mlp" hooks[f"{mlp_prefix}-routed_experts-gate-kernel"] = transpose @@ -1284,11 +1364,19 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): 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 + hooks[f"{mlp_prefix}-shared_expert-wi_0-kernel_scale"] = reshape_scale + hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel_scale"] = reshape_scale + hooks[f"{mlp_prefix}-shared_expert-wo-kernel_scale"] = reshape_scale + hooks[f"{mlp_prefix}-shared_expert_gate-kernel_scale"] = reshape_scale # 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 + hooks[f"{mlp_prefix}-routed_experts-wo-kernel_scale"] = reshape_scale + hooks[(f"{mlp_prefix}-routed_experts-wi_0-kernel_scale", f"{mlp_prefix}-routed_experts-wi_1-kernel_scale")] = ( + reshape_scale + ) # Vision hooks for Qwen3.5 vision_config = config.get("vision_config", None) @@ -4263,6 +4351,8 @@ 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.5-35b-a3b-fp8": QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING, + "qwen3.5-35b-fp8": 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, @@ -4317,6 +4407,8 @@ 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.5-35b-a3b-fp8": QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN, + "qwen3.5-35b-fp8": 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.5-35b-a3b-fp8.yml b/src/maxtext/configs/models/qwen3.5-35b-a3b-fp8.yml new file mode 100644 index 0000000000..c169babc4c --- /dev/null +++ b/src/maxtext/configs/models/qwen3.5-35b-a3b-fp8.yml @@ -0,0 +1,75 @@ +# 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. + +# model config for qwen3.5-35b-a3b-fp8 (FP8 weight-only storage with dynamic dequantization) + +decoder_block: "qwen3_5" + +# Core Architectural Parameters +base_emb_dim: 2048 +base_num_decoder_layers: 40 +base_num_query_heads: 16 +base_num_kv_heads: 2 +head_dim: 256 +vocab_size: 248320 +normalization_layer_epsilon: 1.0e-6 + +# MoE Specific Parameters +# Set base_mlp_dim to match base_moe_mlp_dim to pass validation for fully MoE models. +base_mlp_dim: 512 +base_moe_mlp_dim: 512 +num_experts: 256 +shared_experts: 1 +num_experts_per_tok: 8 +norm_topk_prob: True + +# 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: 32 +gdn_chunk_size: 64 + +# RoPE Settings +rope_max_timescale: 10000000 +partial_rotary_factor: 0.25 + +# General Model Settings +enable_dropout: False + +# Vision Encoder Configuration (need to set use_multimodal=true) +vision_encoder_block: "qwen3_5" +# Based on Qwen3.5 MoE Vision Model Config +image_size_for_vit: 768 +hidden_size_for_vit: 1152 +intermediate_size_for_vit: 4304 +num_attention_heads_for_vit: 16 +num_hidden_layers_for_vit: 27 +num_channels_for_vit: 3 +patch_size_for_vit: 16 +temporal_patch_size_for_vit: 2 +spatial_merge_size_for_vit: 2 +out_hidden_size_for_vit: 2048 # Projects to decoder emb_dim (2048) +num_position_embeddings_for_vit: 2304 +deepstack_visual_indexes_for_vit: [] # No deepstack for Qwen3.5 VL +rope_theta_for_vit: 10000 + +# MRoPE Settings (Multi-dimensional RoPE for multimodal) +use_mrope: true +mrope_section: [11, 11, 10] + +weight_dtype: "float8_e4m3fn" +dtype: "bfloat16" diff --git a/src/maxtext/utils/globals.py b/src/maxtext/utils/globals.py index 30f6e65124..633e6e2261 100644 --- a/src/maxtext/utils/globals.py +++ b/src/maxtext/utils/globals.py @@ -85,6 +85,8 @@ "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.5-35b-a3b-fp8": "Qwen/Qwen3.5-35B-A3B-FP8", + "qwen3.5-35b-fp8": "Qwen/Qwen3.5-35B-A3B-FP8", "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1", "mistral-7b": "mistralai/Mistral-7B-v0.1", "mixtral-8x22b": "mistralai/Mixtral-8x22B-Instruct-v0.1", diff --git a/tests/unit/configs_test.py b/tests/unit/configs_test.py index 2a7bd0f660..693267022c 100644 --- a/tests/unit/configs_test.py +++ b/tests/unit/configs_test.py @@ -246,6 +246,8 @@ def test_mistral_configs(config_file): os.path.join(CONFIGS_DIR, "models", "qwen3-480b-a35b.yml"), os.path.join(CONFIGS_DIR, "models", "qwen3-next-80b-a3b.yml"), os.path.join(CONFIGS_DIR, "models", "qwen3-omni-30b-a3b.yml"), + os.path.join(CONFIGS_DIR, "models", "qwen3.5-35b-a3b.yml"), + os.path.join(CONFIGS_DIR, "models", "qwen3.5-35b-a3b-fp8.yml"), ]