From eb04fe67e556023daf3655cb2263fc036372cc9b Mon Sep 17 00:00:00 2001 From: Kuo Wei Date: Thu, 6 Aug 2026 07:15:12 +0000 Subject: [PATCH] Add Hy3 (Tencent Hunyuan V3) model support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds support for Hy3 (Tencent Hunyuan V3, tencent/Hy3 on HF, 295B total / 21B active MoE). Hy3 combines standard GQA + QK-Norm attention (as in Qwen3) with a DeepSeek-V3-style aux-loss-free sigmoid+bias routed MoE (1 shared expert) and a dense first layer (first_num_dense_layers). It reuses DeepSeekGenericLayer's dense/MoE scaffolding and moe.RoutedAndSharedMoE rather than introducing new attention math — the new decoder layer (Hy3DenseLayer/Hy3MoELayer in src/maxtext/models/hy3.py) subclasses DeepSeekGenericLayer and overrides self_attention with plain GQA, following the pattern DeepSeek4DecoderLayer uses in deepseek4.py. Key changes: - Wire DecoderBlockType.HY3 through decoders.py/nnx_decoders.py, moe.py routing gates, and types.py validation guards. - Add src/maxtext/models/hy3.py with Hy3DenseLayer and Hy3MoELayer. - Add hy3-tiny.yml and hy3-295b.yml configs. - Register Hy3 in the checkpoint conversion framework (hf_model_configs.py, param_mapping.py, hf_shape.py, globals.py) with native HYV3Config try/fallback. - Update FLOPs/MFU calculations (get_dense_moe_layers in maxtext_utils.py) and param export. - Add unit tests in tests/unit/hy3_vs_reference_test.py and tests/unit/nnx_decoders_test.py. - Add end-to-end user guide: tests/end_to_end/tpu/hy3/Run_Hy3.md. - Handle MoE block name mapping in train.py and document auxiliary load balancing behavior. --- .../supported_models_and_architectures.md | 6 + .../utils/hf_model_configs.py | 54 ++ .../checkpoint_conversion/utils/hf_shape.py | 97 ++++ .../utils/param_mapping.py | 160 ++++++ src/maxtext/common/common_types.py | 1 + src/maxtext/configs/models/hy3-295b.yml | 50 ++ src/maxtext/configs/models/hy3-tiny.yml | 51 ++ src/maxtext/configs/types.py | 16 +- src/maxtext/experimental/rl/grpo_utils.py | 2 +- src/maxtext/layers/decoders.py | 26 +- src/maxtext/layers/moe.py | 18 +- src/maxtext/layers/nnx_decoders.py | 6 +- src/maxtext/models/deepseek.py | 4 +- src/maxtext/models/hy3.py | 224 ++++++++ src/maxtext/trainers/pre_train/train.py | 25 +- .../utils/generate_param_only_checkpoint.py | 6 +- src/maxtext/utils/globals.py | 1 + src/maxtext/utils/maxtext_utils.py | 12 +- tests/end_to_end/tpu/hy3/Run_Hy3.md | 205 ++++++++ tests/unit/configs_test.py | 15 +- tests/unit/hy3_vs_reference_test.py | 496 ++++++++++++++++++ tests/unit/nnx_decoders_test.py | 47 ++ 22 files changed, 1492 insertions(+), 30 deletions(-) create mode 100644 src/maxtext/configs/models/hy3-295b.yml create mode 100644 src/maxtext/configs/models/hy3-tiny.yml create mode 100644 src/maxtext/models/hy3.py create mode 100644 tests/end_to_end/tpu/hy3/Run_Hy3.md create mode 100644 tests/unit/hy3_vs_reference_test.py diff --git a/docs/reference/models/supported_models_and_architectures.md b/docs/reference/models/supported_models_and_architectures.md index 5213dc6dcc..a6a4c45727 100644 --- a/docs/reference/models/supported_models_and_architectures.md +++ b/docs/reference/models/supported_models_and_architectures.md @@ -66,6 +66,11 @@ MaxText is an open-source, high-performance LLM framework written in Python/JAX. - **Variants**: K2 (1T), K2-Thinking (1T), K2.5 (text), K2.6 (text) - **Notes**: DeepSeek V3 architecture; MuonClip optimizer +### Hy3 + +- **Variants**: Hy3 (295B, MoE 21B-A active) +- **Notes**: GQA; **QK-Norm**; RMSNorm; RoPE; DeepSeek-V3-style aux-loss-free sigmoid+bias routed MoE with 1 shared expert; dense first layer; MTP. + ## Parallelism building blocks MaxText supports a wide range of parallelism strategies for scaling training and inference across TPUs and GPUs: @@ -104,6 +109,7 @@ The following summarizes observed runtime efficiency and scaling behaviors of Ma [Qwen3.5 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/qwen3_5.py) - **GPT-OSS**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/gpt_oss/run_gpt_oss.md) | [GPT-OSS Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/gpt_oss.py) - **Kimi**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/kimi/Run_Kimi.md) | [K2 reuses DeepSeek Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/deepseek.py) + - **Hy3**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/hy3/Run_Hy3.md) | [Hy3 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/hy3.py) - **Technical Explanations:** diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py index 89abd56d4c..478a8ef367 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py @@ -1142,6 +1142,59 @@ def __init__(self, **kwargs): deepseek4_284b_config = DeepseekV4Config(**deepseek4_284b_dict) +# Hy3 (Tencent Hunyuan V3, https://huggingface.co/tencent/Hy3) +# transformers ships a native HYV3Config; fall back to a raw-dict PTConfig for +# older pinned versions that predate it (same pattern as Gemma 4 above). +try: + HYV3Config = transformers.HYV3Config # pyrefly: ignore[missing-attribute] +except AttributeError: + + class HYV3Config(PTConfig): # pyrefly: ignore[invalid-inheritance] + model_type = "hy_v3" + + def __init__(self, **kwargs): + self.max_position_embeddings = kwargs.get("max_position_embeddings", 262144) + super().__init__(**kwargs) + + +hy3_295b_config = HYV3Config( + architectures=["HYV3ForCausalLM"], + attention_bias=False, + attention_dropout=0.0, + bos_token_id=120000, + eos_token_id=120025, + pad_token_id=120002, + head_dim=128, + hidden_act="silu", + hidden_size=4096, + intermediate_size=13312, + initializer_range=0.006, + max_position_embeddings=262144, + model_type="hy_v3", + moe_intermediate_size=1536, + expert_hidden_dim=1536, + first_k_dense_replace=1, + moe_router_enable_expert_bias=True, + moe_router_use_sigmoid=True, + route_norm=True, + router_scaling_factor=2.826, + num_attention_heads=64, + num_key_value_heads=8, + num_experts=192, + num_experts_per_tok=8, + num_shared_experts=1, + num_hidden_layers=80, + num_nextn_predict_layers=1, + qk_norm=True, + rms_norm_eps=1e-05, + rope_parameters={"rope_theta": 11158840.0, "rope_type": "default"}, + tie_word_embeddings=False, + torch_dtype="bfloat16", + use_cache=True, + vocab_size=120832, +) + + # from https://huggingface.co/openai/gpt-oss-20b/blob/main/config.json # remove mxfp4 quantization_config, since we are using bf16 gpt_oss_20b_dict = { @@ -1925,6 +1978,7 @@ def __init__(self, **kwargs): "deepseek3-671b": deepseek3_671b_config, "deepseek3.2-671b": deepseek32_671b_config, "deepseek4-284b": deepseek4_284b_config, + "hy3-295b": hy3_295b_config, "gpt-oss-20b": gpt_oss_20b_config, "gpt-oss-120b": gpt_oss_120b_config, "qwen3-omni-30b-a3b": qwen3_omni_30b_a3b_config, diff --git a/src/maxtext/checkpoint_conversion/utils/hf_shape.py b/src/maxtext/checkpoint_conversion/utils/hf_shape.py index 65908d9bce..d8daaad331 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_shape.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_shape.py @@ -1164,6 +1164,102 @@ def QWEN3_VL_HF_WEIGHTS_TO_SHAPE(config): # {maxtext model name: {hf weight name: hf shape}} +def HY3_HF_WEIGHTS_TO_SHAPE(config): + """Returns mapping between HuggingFace Hy3 weights path and their shape. + + Combines QWEN_HF_WEIGHTS_TO_SHAPE's plain GQA + QK-Norm attention block + (Hy3 has no MLA) with DEEPSEEK_HF_WEIGHTS_TO_SHAPE's dense/MoE-split + + shared-expert loop structure (Hy3 has `first_k_dense_replace` and 1 shared + expert, like DeepSeek V3). The MTP layer (index `num_hidden_layers`) is + intentionally not included -- MTP weights are not mapped by the checkpoint + conversion framework for Hy3 today. + + Args: + config (dict): HF configuration dictionary + e.g., https://huggingface.co/tencent/Hy3/blob/main/config.json + + Returns: + dict: A mapping where: + - Keys are HuggingFace model parameter paths + - Values are parameter shape as a list + + To check expected mapping: + from transformers import AutoModelForCausalLM + model_name = "tencent/Hy3" + model = AutoModelForCausalLM.from_pretrained(model_name, dtype="auto", trust_remote_code=True) + for name, val in model.named_parameters(): + print(name, val.shape) + """ + hidden_size = config["hidden_size"] + num_hidden_layers = config["num_hidden_layers"] + vocab_size = config["vocab_size"] + num_attention_heads = config["num_attention_heads"] + num_key_value_heads = config["num_key_value_heads"] + head_dim = config.get("head_dim", hidden_size // num_attention_heads) + intermediate_size = config["intermediate_size"] # For the dense layer + moe_intermediate_size = config["moe_intermediate_size"] # For expert layers + num_experts = config["num_experts"] + num_shared_experts = config.get("num_shared_experts", 0) + first_k_dense = config.get("first_k_dense_replace", 0) + + mapping = { + "model.embed_tokens.weight": [vocab_size, hidden_size], + "model.norm.weight": [hidden_size], + "lm_head.weight": [vocab_size, hidden_size], + } + + for layer_idx in range(num_hidden_layers): + layer_prefix = f"model.layers.{layer_idx}" + layer_mapping = { + f"{layer_prefix}.input_layernorm.weight": [hidden_size], + f"{layer_prefix}.post_attention_layernorm.weight": [hidden_size], + f"{layer_prefix}.self_attn.q_proj.weight": [num_attention_heads * head_dim, hidden_size], + f"{layer_prefix}.self_attn.k_proj.weight": [num_key_value_heads * head_dim, hidden_size], + f"{layer_prefix}.self_attn.v_proj.weight": [num_key_value_heads * head_dim, hidden_size], + f"{layer_prefix}.self_attn.o_proj.weight": [hidden_size, num_attention_heads * head_dim], + f"{layer_prefix}.self_attn.q_norm.weight": [head_dim], + f"{layer_prefix}.self_attn.k_norm.weight": [head_dim], + } + + if layer_idx < first_k_dense: + # Dense MLP layer + layer_mapping.update( + { + f"{layer_prefix}.mlp.gate_proj.weight": [intermediate_size, hidden_size], + f"{layer_prefix}.mlp.up_proj.weight": [intermediate_size, hidden_size], + f"{layer_prefix}.mlp.down_proj.weight": [hidden_size, intermediate_size], + } + ) + else: + # MoE layer: router + expert_bias + shared_mlp + routed experts + layer_mapping.update( + { + f"{layer_prefix}.mlp.router.gate.weight": [num_experts, hidden_size], + f"{layer_prefix}.mlp.expert_bias": [num_experts], + } + ) + if num_shared_experts > 0: + shared_intermediate_size = moe_intermediate_size * num_shared_experts + layer_mapping.update( + { + f"{layer_prefix}.mlp.shared_mlp.gate_proj.weight": [shared_intermediate_size, hidden_size], + f"{layer_prefix}.mlp.shared_mlp.up_proj.weight": [shared_intermediate_size, hidden_size], + f"{layer_prefix}.mlp.shared_mlp.down_proj.weight": [hidden_size, shared_intermediate_size], + } + ) + for expert_j in range(num_experts): + expert_prefix = f"{layer_prefix}.mlp.experts.{expert_j}" + layer_mapping.update( + { + f"{expert_prefix}.gate_proj.weight": [moe_intermediate_size, hidden_size], + f"{expert_prefix}.up_proj.weight": [moe_intermediate_size, hidden_size], + f"{expert_prefix}.down_proj.weight": [hidden_size, moe_intermediate_size], + } + ) + mapping.update(layer_mapping) + return mapping + + def DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE(config): """Returns a dictionary mapping HuggingFace weight names to shapes for DeepSeek V4.""" hidden_size = config["hidden_size"] @@ -1303,6 +1399,7 @@ def DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE(config): "deepseek3-671b": DEEPSEEK_HF_WEIGHTS_TO_SHAPE, "deepseek3.2-671b": DEEPSEEK_HF_WEIGHTS_TO_SHAPE, "deepseek4-284b": DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE, + "hy3-295b": HY3_HF_WEIGHTS_TO_SHAPE, "gpt-oss-20b": GPT_OSS_HF_WEIGHTS_TO_SHAPE, "gpt-oss-120b": GPT_OSS_HF_WEIGHTS_TO_SHAPE, "mixtral-8x7b": MIXTRAL_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 26359cdddc..183c98c1d4 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -1779,6 +1779,164 @@ def DEEPSEEK_NNX_TO_VLLM_PARAM_HOOK_FN(): return {} +def HY3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): + """Generates a parameter mapping from MaxText to HuggingFace Hy3 weight paths. + + Structurally this follows DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING's dense/MoE + two-stack loop (Hy3 also has a `first_k_dense_replace`-driven dense/MoE + split and a DeepSeek-V3-style sigmoid+bias router with a shared expert), + but the attention block is plain GQA + QK-Norm (no MLA), matching + QWEN_MAXTEXT_TO_HF_PARAM_MAPPING's attention keys instead. HF key names + (`mlp.router.gate.weight`, `mlp.expert_bias`, `mlp.shared_mlp.*`) were + confirmed against the real `tencent/Hy3` `model.safetensors.index.json`, + not assumed from DeepSeek/Qwen naming conventions. + + Note: layer `num_hidden_layers` (the MTP layer) is intentionally not + mapped here -- MTP weights are left randomly initialized on conversion. + + Returns: + dict: A mapping where keys are `atomic_mt_key` (single MaxText parameter names). + Values are Hugging Face parameter names in one of four forms: unscanned (string), + scanned (list of strings), unscanned with expert stacking (list of strings), + or scanned with expert stacking (nested list of strings). + """ + num_main_layers = config["num_hidden_layers"] + first_num_dense_layers = config["first_k_dense_replace"] + num_experts = config.get("num_experts", 0) + + mapping = { + "params-token_embedder-embedding": "model.embed_tokens.weight", + "params-decoder-decoder_norm-scale": "model.norm.weight", + "params-decoder-logits_dense-kernel": "lm_head.weight", + } + # Attention keys are shared by both dense and MoE layers: plain GQA + QK-Norm, no bias. + attention_keys = { + "pre_self_attention_layer_norm-scale": "input_layernorm.weight", + "post_self_attention_layer_norm-scale": "post_attention_layernorm.weight", + "self_attention-query-kernel": "self_attn.q_proj.weight", + "self_attention-key-kernel": "self_attn.k_proj.weight", + "self_attention-value-kernel": "self_attn.v_proj.weight", + "self_attention-out-kernel": "self_attn.o_proj.weight", + "self_attention-query_norm-scale": "self_attn.q_norm.weight", + "self_attention-key_norm-scale": "self_attn.k_norm.weight", + } + # Dense layers (layer 0, per first_k_dense_replace=1) + dense_layer_keys = attention_keys | { + "mlp-wi_0-kernel": "mlp.gate_proj.weight", + "mlp-wi_1-kernel": "mlp.up_proj.weight", + "mlp-wo-kernel": "mlp.down_proj.weight", + } + # MoE layers + moe_layer_keys = attention_keys | { + "Hy3MoeBlock_0-shared_experts-wi_0-kernel": "mlp.shared_mlp.gate_proj.weight", + "Hy3MoeBlock_0-shared_experts-wi_1-kernel": "mlp.shared_mlp.up_proj.weight", + "Hy3MoeBlock_0-shared_experts-wo-kernel": "mlp.shared_mlp.down_proj.weight", + "Hy3MoeBlock_0-MoeBlock_0-gate-kernel": "mlp.router.gate.weight", + "Hy3MoeBlock_0-MoeBlock_0-gate-bias": "mlp.expert_bias", + } + # MoE Experts (nested list mapping: [[e0_l0, e0_l1..], [e1_l0, e1_l1..]..]) + moe_expert_keys = { + "Hy3MoeBlock_0-MoeBlock_0-wi_0": "gate_proj.weight", + "Hy3MoeBlock_0-MoeBlock_0-wi_1": "up_proj.weight", + "Hy3MoeBlock_0-MoeBlock_0-wo": "down_proj.weight", + } + + # scan + if scan_layers: + for maxtext_key, hf_key in dense_layer_keys.items(): + mapping[f"params-decoder-dense_layers-{maxtext_key}"] = [ # pyrefly: ignore[bad-assignment] + f"model.layers.{i}.{hf_key}" for i in range(first_num_dense_layers) + ] + + for maxtext_key, hf_key in moe_layer_keys.items(): + mapping[f"params-decoder-moe_layers-{maxtext_key}"] = [ # pyrefly: ignore[bad-assignment] + f"model.layers.{i}.{hf_key}" for i in range(first_num_dense_layers, num_main_layers) + ] + + for maxtext_key, hf_key in moe_expert_keys.items(): + mapping[f"params-decoder-moe_layers-{maxtext_key}"] = [ # pyrefly: ignore[bad-assignment] + [f"model.layers.{i}.mlp.experts.{e}.{hf_key}" for i in range(first_num_dense_layers, num_main_layers)] + for e in range(num_experts) + ] + # unscan + else: + for i in range(first_num_dense_layers): + for maxtext_key, hf_key in dense_layer_keys.items(): + mapping[f"params-decoder-dense_layers_{i}-{maxtext_key}"] = f"model.layers.{i}.{hf_key}" + + for i in range(first_num_dense_layers, num_main_layers): + moe_layer_idx = i - first_num_dense_layers + + for maxtext_key, hf_key in moe_layer_keys.items(): + mapping[f"params-decoder-moe_layers_{moe_layer_idx}-{maxtext_key}"] = f"model.layers.{i}.{hf_key}" + + for maxtext_key, hf_key in moe_expert_keys.items(): + mapping[f"params-decoder-moe_layers_{moe_layer_idx}-{maxtext_key}"] = [ # pyrefly: ignore[bad-assignment] + f"model.layers.{i}.mlp.experts.{e}.{hf_key}" for e in range(num_experts) + ] + return mapping + + +def HY3_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False): + """Creates parameter transformation functions for Hy3.""" + + def reshape_kernel(input_tensor, target_shape): + """Reshapes and transposes kernel weights between MaxText and HF.""" + if saving_to_hf: + flipped_target_shape = np.flip(np.array(target_shape)) + return input_tensor.reshape(flipped_target_shape).T + else: + return input_tensor.T.reshape(target_shape) + + num_main_layers = config["num_hidden_layers"] + first_num_dense_layers = config["first_k_dense_replace"] + + mapping = { + "params-decoder-logits_dense-kernel": reshape_kernel, + } + + attention_need_reshape = { + "self_attention-query-kernel", + "self_attention-key-kernel", + "self_attention-value-kernel", + "self_attention-out-kernel", + } + + dense_need_reshape = attention_need_reshape | { + "mlp-wi_0-kernel", + "mlp-wi_1-kernel", + "mlp-wo-kernel", + } + + moe_need_reshape = attention_need_reshape | { + "Hy3MoeBlock_0-shared_experts-wi_0-kernel", + "Hy3MoeBlock_0-shared_experts-wi_1-kernel", + "Hy3MoeBlock_0-shared_experts-wo-kernel", + "Hy3MoeBlock_0-MoeBlock_0-gate-kernel", + "Hy3MoeBlock_0-MoeBlock_0-wi_0", + "Hy3MoeBlock_0-MoeBlock_0-wi_1", + "Hy3MoeBlock_0-MoeBlock_0-wo", + } + + # scan + if scan_layers: + for key in dense_need_reshape: + mapping[f"params-decoder-dense_layers-{key}"] = reshape_kernel + for key in moe_need_reshape: + mapping[f"params-decoder-moe_layers-{key}"] = reshape_kernel + # unscan + else: + for i in range(first_num_dense_layers): + for key in dense_need_reshape: + mapping[f"params-decoder-dense_layers_{i}-{key}"] = reshape_kernel + for i in range(first_num_dense_layers, num_main_layers): + moe_layer_idx = i - first_num_dense_layers + for key in moe_need_reshape: + mapping[f"params-decoder-moe_layers_{moe_layer_idx}-{key}"] = reshape_kernel + + return mapping + + def GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): """Generates mapping from MaxText gpt-oss to Hugging Face weight paths. @@ -4241,6 +4399,7 @@ def mhc_concat_scale(input_tensors, target_shape=None): "deepseek3-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, "deepseek3.2-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, "deepseek4-284b": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING, + "hy3-295b": HY3_MAXTEXT_TO_HF_PARAM_MAPPING, "gpt-oss-20b": GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING, "gpt-oss-120b": GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -4296,6 +4455,7 @@ def mhc_concat_scale(input_tensors, target_shape=None): "deepseek3.2-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, "deepseek4-tiny": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN, "deepseek4-284b": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN, + "hy3-295b": HY3_MAXTEXT_TO_HF_PARAM_HOOK_FN, "gpt-oss-20b": GPT_OSS_TO_HF_PARAM_HOOK_FN, "gpt-oss-120b": GPT_OSS_TO_HF_PARAM_HOOK_FN, "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_HOOK_FN, diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py index 664ba8b388..464399626a 100644 --- a/src/maxtext/common/common_types.py +++ b/src/maxtext/common/common_types.py @@ -115,6 +115,7 @@ class DecoderBlockType(enum.Enum): OLMO3 = "olmo3" DEEPSEEK4 = "deepseek4" ENVY = "envy" + HY3 = "hy3" class VisionEncoderBlockType(enum.Enum): diff --git a/src/maxtext/configs/models/hy3-295b.yml b/src/maxtext/configs/models/hy3-295b.yml new file mode 100644 index 0000000000..f1bf312779 --- /dev/null +++ b/src/maxtext/configs/models/hy3-295b.yml @@ -0,0 +1,50 @@ +# Copyright 2023–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 Hy3 (Tencent Hunyuan V3, https://huggingface.co/tencent/Hy3) +# 295B total / 21B active MoE. Standard GQA + QK-Norm attention (no MLA, no +# compressed/sparse attention) with a DeepSeek-V3-style aux-loss-free +# sigmoid+bias routed MoE (1 shared expert) and a dense first layer. + +decoder_block: "hy3" + +# --- Core Architectural Parameters --- +base_emb_dim: 4096 +base_mlp_dim: 13312 +base_num_query_heads: 64 +base_num_kv_heads: 8 +head_dim: 128 +base_num_decoder_layers: 80 +vocab_size: 120832 +mlp_activations: ["silu", "linear"] +normalization_layer_epsilon: 1.0e-5 +logits_via_embedding: false +enable_dropout: false + +# --- Dense/MoE split: layer 0 is dense, layers 1-79 are MoE --- +first_num_dense_layers: 1 + +# --- MoE configuration (aux-loss-free sigmoid + bias router, 1 shared expert) --- +num_experts: 192 +num_experts_per_tok: 8 +base_moe_mlp_dim: 1536 +shared_experts: 1 +routed_score_func: "sigmoid" +routed_bias: true +routed_scaling_factor: 2.826 +norm_topk_prob: true + +# --- Attention: plain GQA + QK-Norm (no MLA, no compressed attention) --- +use_qk_norm: true +rope_max_timescale: 11158840 diff --git a/src/maxtext/configs/models/hy3-tiny.yml b/src/maxtext/configs/models/hy3-tiny.yml new file mode 100644 index 0000000000..34826ef3d6 --- /dev/null +++ b/src/maxtext/configs/models/hy3-tiny.yml @@ -0,0 +1,51 @@ +# Copyright 2023–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. + +# Test config for Hy3 (Tencent Hunyuan V3, https://huggingface.co/tencent/Hy3). +# Small dims/layer count for fast local iteration with random init; architecturally +# significant fields (routing, norm epsilon, QK-Norm) keep their real hy3-295b values. + +decoder_block: "hy3" + +# --- Small dims for fast debugging --- +base_emb_dim: 256 +base_mlp_dim: 512 +base_num_query_heads: 8 +base_num_kv_heads: 2 +head_dim: 64 +base_num_decoder_layers: 4 +vocab_size: 1000 + +# --- Standard Defaults --- +enable_dropout: false +logits_via_embedding: false +normalization_layer_epsilon: 1.0e-5 +mlp_activations: ["silu", "linear"] + +# --- Hy3 dense/MoE split --- +first_num_dense_layers: 1 + +# --- MoE configuration (aux-loss-free sigmoid + bias router, 1 shared expert) --- +num_experts: 8 +num_experts_per_tok: 2 +base_moe_mlp_dim: 128 +shared_experts: 1 +routed_score_func: "sigmoid" +routed_bias: true +routed_scaling_factor: 2.826 +norm_topk_prob: true + +# --- Attention: plain GQA + QK-Norm (no MLA, no compressed attention) --- +use_qk_norm: true +rope_max_timescale: 11158840 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index c4a270a567..3c60e0d674 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -233,6 +233,8 @@ class ProfilerType(str, Enum): "deepseek4-284b", "deepseek-custom", "kimi-k2-1t", + "hy3-tiny", + "hy3-295b", "gemma-7b", "gemma-2b", "gemma2-2b", @@ -3261,7 +3263,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de self.tensors_to_offload = [t for t in tensors if getattr(self, t) == "offload"] if self.pipeline_parallel_layers == -1: - if self.decoder_block == DecoderBlockType.DEEPSEEK: + if self.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3): moe_layers = self.num_decoder_layers - self.first_num_dense_layers self.pipeline_parallel_layers = moe_layers else: @@ -3531,8 +3533,16 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de ) if self.decoder_block == DecoderBlockType.GPT_OSS and not self.sparse_matmul and self.capacity_factor != -1: raise ValueError("GPT-OSS MoE only supports dropless (capacity_factor=-1) with dense matmul.") - if self.routed_bias and self.routed_bias_update_rate > 0.0 and self.decoder_block != DecoderBlockType.DEEPSEEK: - raise ValueError("Loss-free load balancing is only supported for the DeepSeek decoder block.") + if ( + self.routed_bias + and self.routed_bias_update_rate > 0.0 + and self.decoder_block + not in ( + DecoderBlockType.DEEPSEEK, + DecoderBlockType.HY3, + ) + ): + raise ValueError("Loss-free load balancing is only supported for the DeepSeek and Hy3 decoder blocks.") if self.model_name.startswith("deepseek4") and self.first_num_hash_layers > 0 and self.use_ring_of_experts: raise ValueError("DeepSeek V4 hash routing is currently not supported with ring of experts.") self.validate_ragged_buffer_factor() diff --git a/src/maxtext/experimental/rl/grpo_utils.py b/src/maxtext/experimental/rl/grpo_utils.py index 946ae552ef..b578b1788b 100644 --- a/src/maxtext/experimental/rl/grpo_utils.py +++ b/src/maxtext/experimental/rl/grpo_utils.py @@ -289,7 +289,7 @@ def pathways_reshard(config, inference_engine, params, source_shardings, source_ source_mesh: The source device mesh. destination_shardings: The sharding specification for the destination. """ - if config.decoder_block == DecoderBlockType.DEEPSEEK: + if config.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3): layer_groups = [ ("dense_layers", config.first_num_dense_layers), ("moe_layers", config.base_num_decoder_layers - config.first_num_dense_layers), diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 42753eb752..2fdd84513d 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -52,6 +52,7 @@ gemma4_small, gpt3, gpt_oss, + hy3, llama2, llama4, mistral, @@ -454,6 +455,11 @@ def get_decoder_layers(self): deepseek.DeepSeekDenseLayerToLinen, deepseek.DeepSeekMoELayerToLinen, ] + case DecoderBlockType.HY3: + return [ + hy3.Hy3DenseLayerToLinen, + hy3.Hy3MoELayerToLinen, + ] case DecoderBlockType.DEEPSEEK4: return ( [deepseek4.DeepSeek4ScannableBlockToLinen] if self.config.scan_layers else [deepseek4.DeepSeek4LayerToLinen] @@ -529,6 +535,7 @@ def get_scannable(normal_cls, scannable_cls): DecoderBlockType.SIMPLE: [simple_layer.SimpleDecoderLayer], DecoderBlockType.SIMPLE_MLP: [simple_layer.SimpleMlpDecoderLayer], DecoderBlockType.DEEPSEEK: [deepseek.DeepSeekDenseLayer, deepseek.DeepSeekMoELayer], + DecoderBlockType.HY3: [hy3.Hy3DenseLayer, hy3.Hy3MoELayer], DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock), DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock), DecoderBlockType.ENVY: get_scannable(envy.EnvyDecoderLayer, envy.EnvyScannableBlock), @@ -569,7 +576,9 @@ def map_fn(path, value): def _build_nnx_pipeline_stage(self, decoder_blocks, rngs): """Creates a single NNX pipeline stage module.""" cfg = self.config - base_stage_cls = decoder_blocks[1] if cfg.decoder_block == DecoderBlockType.DEEPSEEK else decoder_blocks[0] + base_stage_cls = ( + decoder_blocks[1] if cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3) else decoder_blocks[0] + ) if cfg.num_layers_per_pipeline_stage == 1: return base_stage_cls(config=cfg, mesh=self.mesh, quant=self.quant, model_mode=self.model_mode, rngs=rngs) @@ -585,7 +594,7 @@ def get_pipeline_stage_module(self, decoder_blocks): """get pipeline stage module""" def get_layer_to_pipeline(blocks, cfg): - if cfg.decoder_block == DecoderBlockType.DEEPSEEK: + if cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3): return blocks[1] # return the sparse block else: return blocks[0] @@ -627,6 +636,7 @@ def get_norm_layer(self, num_features: int): DecoderBlockType.MIXTRAL, DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4, + DecoderBlockType.HY3, DecoderBlockType.GEMMA, DecoderBlockType.GEMMA2, DecoderBlockType.GEMMA3, @@ -896,8 +906,8 @@ def __call__( if cfg.pipeline_fsdp_ag_once or cfg.pipeline_fsdp_ag_per_repeat else None ) - if cfg.decoder_block == DecoderBlockType.DEEPSEEK: - assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek." + if cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3): + assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek/hy3." dense_layer = RemattedBlockLayers[0] moe_layer = RemattedBlockLayers[1] num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers @@ -942,8 +952,8 @@ def __call__( )(y, *broadcast_args) else: if cfg.scan_layers: - if cfg.decoder_block == DecoderBlockType.DEEPSEEK: - assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek." + if cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3): + assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek/hy3." layer_call_kwargs = { "previous_chunk": previous_chunk, "slot": slot, @@ -1148,8 +1158,8 @@ def __call__( **layer_kwargs, )(y, *current_broadcast_args) else: - if cfg.decoder_block == DecoderBlockType.DEEPSEEK: - assert len(RemattedBlockLayers) == 2, "Unscanned layers must have a length of 2 using deepseek." + if cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3): + assert len(RemattedBlockLayers) == 2, "Unscanned layers must have a length of 2 using deepseek/hy3." dense_layer = RemattedBlockLayers[0] moe_layer = RemattedBlockLayers[1] diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index b310d97fb5..014a4f62cd 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -364,7 +364,7 @@ def __call__(self, inputs: jax.Array, _initializing: bool = False) -> Tuple[jax. output = linears._convert_to_activation_function(self.score_func)(output) # NOTE: deepseek2 has a different pattern - if self.model_name.startswith(("deepseek3", "deepseek4")): + if self.model_name.startswith(("deepseek3", "deepseek4", "hy3")): pre_bias_logits = output if self.use_bias: @@ -717,7 +717,7 @@ def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None): top_k_indices = tid2eid_int[input_ids.astype(jnp.int32)] top_k_weights = jnp.take_along_axis(pre_bias_logits, top_k_indices, axis=-1) # NOTE: deepseek2 has a different pattern - elif self.config.model_name.startswith(("deepseek3", "deepseek4")): + elif self.config.model_name.startswith(("deepseek3", "deepseek4", "hy3")): top_k_weights, top_k_indices = self.deepseek_routing(gate_logits, pre_bias_logits) elif self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4: router_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1) @@ -726,7 +726,11 @@ def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None): else: top_k_weights, top_k_indices = jax.lax.top_k(gate_logits, self.num_experts_per_tok) - if self.config.decoder_block in (ctypes.DecoderBlockType.DEEPSEEK, ctypes.DecoderBlockType.DEEPSEEK4): + if self.config.decoder_block in ( + ctypes.DecoderBlockType.DEEPSEEK, + ctypes.DecoderBlockType.DEEPSEEK4, + ctypes.DecoderBlockType.HY3, + ): top_k_weights = self.deepseek_scale_weights(top_k_weights) else: if self.config.decoder_block not in (ctypes.DecoderBlockType.LLAMA4, ctypes.DecoderBlockType.GEMMA4): @@ -1583,7 +1587,7 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): gate_logits_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) # NOTE: deepseek2 has a different pattern - if self.config.model_name.startswith(("deepseek3", "deepseek4")): + if self.config.model_name.startswith(("deepseek3", "deepseek4", "hy3")): pre_bias_logits_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) else: # pre_bias_logits is None for non-deepseek3/4 models, including deepseek2 @@ -2345,7 +2349,7 @@ def sparse_matmul_route_and_compute( gate_logits_axes = (batch_logical_axis, "activation_norm_length", None) # NOTE: deepseek2 has a different pattern - if self.config.model_name.startswith(("deepseek3", "deepseek4")): + if self.config.model_name.startswith(("deepseek3", "deepseek4", "hy3")): pre_bias_logits_axes = (batch_logical_axis, "activation_norm_length", None) else: pre_bias_logits_axes = None @@ -2644,8 +2648,8 @@ def dense_matmul( # gate_logits: batch, length, expert gate_logits = self._maybe_shard_with_logical(gate_logits, ("activation_batch_moe", "activation_length_moe", None)) # NOTE: deepseek2 has a different pattern - if self.config.model_name.startswith(("deepseek3", "deepseek4")): - # pre_bias_logits is None for non-deepseek3/4 models, including deepseek2 + if self.config.model_name.startswith(("deepseek3", "deepseek4", "hy3")): + # pre_bias_logits is None for non-deepseek3/4/hy3 models, including deepseek2 pre_bias_logits = self._maybe_shard_with_logical( pre_bias_logits, ("activation_batch_moe", "activation_length_moe", None) ) diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 895ea27c14..4a2f1cbf2d 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -56,6 +56,7 @@ gemma4_small, gpt3, gpt_oss, + hy3, llama2, llama4, mistral, @@ -432,7 +433,8 @@ def __init__( ) self.scanned_layers = None - self.is_deepseek = self.config.decoder_block == DecoderBlockType.DEEPSEEK + # DeepSeek and Hy3 both use a two-stack dense/MoE layer split driven by `first_num_dense_layers`. + self.is_deepseek = self.config.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3) self.is_deepseek4 = self.config.decoder_block == DecoderBlockType.DEEPSEEK4 self.is_gemma3 = self.config.decoder_block == DecoderBlockType.GEMMA3 self.is_gemma4 = self.config.decoder_block == DecoderBlockType.GEMMA4 @@ -1118,6 +1120,7 @@ def get_deepseek(): DecoderBlockType.SIMPLE: [simple_layer.SimpleDecoderLayer], DecoderBlockType.SIMPLE_MLP: [simple_layer.SimpleMlpDecoderLayer], DecoderBlockType.DEEPSEEK: get_deepseek(), + DecoderBlockType.HY3: [hy3.Hy3DenseLayer, hy3.Hy3MoELayer], DecoderBlockType.DEEPSEEK4: get_scannable(deepseek4.DeepSeek4DecoderLayer, deepseek4.DeepSeek4ScannableBlock), DecoderBlockType.GPT_OSS: get_scannable(gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock), DecoderBlockType.QWEN3_NEXT: get_scannable(qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock), @@ -1271,6 +1274,7 @@ def get_norm_layer(self, num_features: int, rngs: nnx.Rngs): DecoderBlockType.MIXTRAL, DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4, + DecoderBlockType.HY3, DecoderBlockType.GEMMA, DecoderBlockType.GEMMA2, DecoderBlockType.GEMMA3, diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py index 0ad8978e7f..9d14462b54 100644 --- a/src/maxtext/models/deepseek.py +++ b/src/maxtext/models/deepseek.py @@ -139,8 +139,8 @@ def __init__( self.engram_layer_norm = None self.engram = None - # DeepSeek V4 natively overrides this block with CompressedAttention. - if self.config.decoder_block != DecoderBlockType.DEEPSEEK4: + # DeepSeek V4 and Hy3 natively override this block with their own attention module. + if self.config.decoder_block not in (DecoderBlockType.DEEPSEEK4, DecoderBlockType.HY3): self.self_attention = attention_mla.MLA( config=self.config, num_query_heads=self.config.num_query_heads, diff --git a/src/maxtext/models/hy3.py b/src/maxtext/models/hy3.py new file mode 100644 index 0000000000..9612bdb576 --- /dev/null +++ b/src/maxtext/models/hy3.py @@ -0,0 +1,224 @@ +# Copyright 2023–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. + +"""Hy3 (Tencent Hunyuan V3) model definition. + +Hy3 combines standard GQA attention with QK-Norm (as in Qwen3) with a +DeepSeek-V3-style aux-loss-free sigmoid+bias MoE router that has a shared +expert and a dense first layer (`first_num_dense_layers`). Unlike DeepSeek +V3/V4, Hy3 uses no Multi-Head Latent Attention and no compressed/sparse +attention, so this reuses `DeepSeekGenericLayer`'s generic dense/MoE-split +scaffolding (norms, dropout, sharding, post_process) but overrides +`self_attention` with plain GQA, following the same subclassing pattern +`DeepSeek4DecoderLayer` uses in `deepseek4.py`. +""" + +from typing import Optional + +from flax import nnx +from jax.ad_checkpoint import checkpoint_name +from jax.sharding import Mesh + +from maxtext.common.common_types import Config +from maxtext.layers import initializers +from maxtext.layers import linears +from maxtext.layers import moe +from maxtext.layers import nnx_wrappers +from maxtext.layers import quantizations +from maxtext.layers.attentions import Attention +from maxtext.models import deepseek + + +def _build_hy3_attention(config: Config, mesh: Mesh, model_mode: str, quant, dummy_inputs_shape, rngs: nnx.Rngs): + """Builds the plain GQA + QK-Norm attention module shared by Hy3's dense and MoE layers.""" + return Attention( + config=config, + num_query_heads=config.num_query_heads, + num_kv_heads=config.num_kv_heads, + head_dim=config.head_dim, + max_target_length=config.max_target_length, + max_prefill_predict_length=config.max_prefill_predict_length, + attention_kernel=config.attention, + inputs_q_shape=dummy_inputs_shape, + inputs_kv_shape=dummy_inputs_shape, + mesh=mesh, + dtype=config.dtype, + weight_dtype=config.weight_dtype, + dropout_rate=config.dropout_rate, + float32_qk_product=config.float32_qk_product, + float32_logits=config.float32_logits, + quant=quant, + kv_quant=quantizations.configure_kv_quant(config), + use_ragged_attention=config.use_ragged_attention, + ragged_block_size=config.ragged_block_size, + use_qk_norm=config.use_qk_norm, + query_pre_attn_scalar=config.head_dim**-0.5, + model_mode=model_mode, + name="self_attention", + rngs=rngs, + ) + + +class Hy3DenseLayer(deepseek.DeepSeekGenericLayer): + """Hy3 dense decoder layer, used only for the first `config.first_num_dense_layers` layers.""" + + def __init__( + self, + config: Config, + model_mode: str, + mesh: Mesh, + rngs: nnx.Rngs, + quant: Optional[quantizations.AqtQuantization] = None, + layer_idx: int = -1, + ) -> None: + super().__init__(config, model_mode, mesh, rngs, quant, layer_idx) + self.self_attention = _build_hy3_attention(self.config, mesh, model_mode, quant, self.dummy_inputs_shape, rngs) + self.mlp = linears.MlpBlock( + in_features=self.dummy_inputs_shape[-1], + intermediate_dim=self.config.mlp_dim, + activations=self.config.mlp_activations, + intermediate_dropout_rate=self.config.dropout_rate, + dtype=self.config.dtype, + weight_dtype=self.config.weight_dtype, + config=self.config, + quant=quant, + model_mode=model_mode, + mesh=mesh, + rngs=self.rngs, + ) + + def mlp_op(self, x, deterministic, *args, **kwargs): + mlp = self.mlp(x, deterministic, intermediate_sharding=self.mlp_intermediate_sharding, out_sharding=self.out_sharding) + return self.with_logical_constraint(mlp) + + def __call__( + self, + inputs, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=None, + slot: None | int = None, + kv_cache=None, + attention_metadata=None, + decoder_input_tokens=None, + ): + if isinstance(inputs, tuple): + inputs = inputs[0] + x = self.with_logical_constraint(inputs) + x = checkpoint_name(x, "decoder_layer_input") + + if self.is_engram_enabled: + engram_output = self.engram_op(x, decoder_input_tokens) + x = x + engram_output + + hidden_states, intermediate_inputs = self.self_attention_with_norm_op( + x, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + ) + + mlp_lnx = self.mlp_op(hidden_states, deterministic) + layer_output = mlp_lnx + intermediate_inputs + layer_output = self.dropout_op(layer_output, deterministic=deterministic) + + return self.post_process(layer_output, None, None, kv_cache) + + +Hy3DenseLayerToLinen = nnx_wrappers.to_linen_class( + Hy3DenseLayer, + base_metadata_fn=initializers.variable_to_logically_partitioned, +) + + +class Hy3MoELayer(deepseek.DeepSeekGenericLayer): + """Hy3 MoE decoder layer: GQA+QK-Norm attention with a DeepSeek-V3-style + sigmoid+bias routed MoE block (1 shared expert).""" + + def __init__( + self, + config: Config, + model_mode: str, + mesh: Mesh, + rngs: nnx.Rngs, + quant: Optional[quantizations.AqtQuantization] = None, + layer_idx: int = -1, + ) -> None: + super().__init__(config, model_mode, mesh, rngs, quant, layer_idx) + self.self_attention = _build_hy3_attention(self.config, mesh, model_mode, quant, self.dummy_inputs_shape, rngs) + self.Hy3MoeBlock_0 = moe.RoutedAndSharedMoE( + config=self.config, + mesh=mesh, + kernel_init=initializers.nd_dense_init(self.config.dense_init_scale, "fan_in", "truncated_normal"), + kernel_axes=("embed", None), + dtype=self.config.dtype, + weight_dtype=self.config.weight_dtype, + quant=quant, + rngs=self.rngs, + ) + + def mlp_op(self, x, deterministic, *args, **kwargs): + mlp_lnx, load_balance_loss, moe_bias_updates = self.Hy3MoeBlock_0( + x, intermediate_sharding=self.mlp_intermediate_sharding, out_sharding=self.out_sharding + ) + return self.with_logical_constraint(mlp_lnx), load_balance_loss, moe_bias_updates + + def __call__( + self, + inputs, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=None, + slot: None | int = None, + kv_cache=None, + attention_metadata=None, + decoder_input_tokens=None, + ): + if isinstance(inputs, tuple): + inputs = inputs[0] + x = self.with_logical_constraint(inputs) + x = checkpoint_name(x, "decoder_layer_input") + + if self.is_engram_enabled: + engram_output = self.engram_op(x, decoder_input_tokens) + x = x + engram_output + + hidden_states, intermediate_inputs = self.self_attention_with_norm_op( + x, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + ) + + mlp_lnx, load_balance_loss, moe_bias_updates = self.mlp_op(hidden_states, deterministic) + layer_output = mlp_lnx + intermediate_inputs + layer_output = self.dropout_op(layer_output, deterministic=deterministic) + + return self.post_process(layer_output, load_balance_loss, moe_bias_updates, kv_cache) + + +Hy3MoELayerToLinen = nnx_wrappers.to_linen_class( + Hy3MoELayer, + base_metadata_fn=initializers.variable_to_logically_partitioned, +) diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 138bfe1ec7..0a743e4951 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -39,6 +39,7 @@ from flax.linen import partitioning as nn_partitioning from flax.nnx import variablelib +from maxtext.common.common_types import DecoderBlockType from maxtext.configs import pyconfig from maxtext.utils.globals import EPS from maxtext.utils import elastic_utils @@ -73,6 +74,24 @@ VertexTensorboardManager, _vertex_tb_is_stub = vertex_tensorboard_modules() +# Aux-loss-free load balancing (DeepSeek V3-style sigmoid+bias routing) sows its +# router-bias update under a module attribute name that differs per model family, +# since each family names its RoutedAndSharedMoE instance differently. +# NOTE: as of writing, this update path is a no-op in scanned mode for every family +# below (nnx_decoders.py's scan application strips nnx.Intermediate state before it +# reaches here, so `moe_bias_updates` is always None -- see nnx_decoders.py:1089-ish +# `nnx.filter_state(scanned_state, nnx.Not((nnx.RngState, nnx.Intermediate)))`) and +# raises AttributeError in unscanned mode (this code assumes the single stacked +# `moe_layers` attribute scanning produces, which doesn't exist when unscanned -- +# unscanned layers are named `moe_layers_0`, `moe_layers_1`, etc. instead). Both are +# pre-existing gaps that predate Hy3 and affect DeepSeek V3 too; this table only +# prevents Hy3 from crashing on a name mismatch on top of that, it does not make +# the update path functional. +_MOE_BLOCK_ATTR_BY_DECODER_BLOCK = { + DecoderBlockType.DEEPSEEK: "DeepSeekMoeBlock_0", + DecoderBlockType.HY3: "Hy3MoeBlock_0", +} + def get_first_step(model, state): if isinstance(model, nn.Module): @@ -525,7 +544,8 @@ def move(path, value): # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family if config.routed_bias and config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None: - target_path = ("params", "decoder", "moe_layers", "DeepSeekMoeBlock_0", "MoeBlock_0", "gate", "bias") + moe_block_attr = _MOE_BLOCK_ATTR_BY_DECODER_BLOCK.get(config.decoder_block, "DeepSeekMoeBlock_0") + target_path = ("params", "decoder", "moe_layers", moe_block_attr, "MoeBlock_0", "gate", "bias") # Updates the shape to be aligned with state. moe_bias_updates = jnp.array(moe_bias_updates[0]).transpose() new_state = maxtext_utils.update_state_param(new_state, target_path, moe_bias_updates) @@ -557,7 +577,8 @@ def move(path, value): # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family if config.routed_bias and config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None: - target_bias = new_state.model.decoder.moe_layers.DeepSeekMoeBlock_0.MoeBlock_0.gate.bias + moe_block_attr = _MOE_BLOCK_ATTR_BY_DECODER_BLOCK.get(config.decoder_block, "DeepSeekMoeBlock_0") + target_bias = getattr(new_state.model.decoder.moe_layers, moe_block_attr).MoeBlock_0.gate.bias target_bias.value = target_bias.value + jnp.array(moe_bias_updates[0]).transpose() lm_loss = xent_sum / (total_weights + EPS) diff --git a/src/maxtext/utils/generate_param_only_checkpoint.py b/src/maxtext/utils/generate_param_only_checkpoint.py index 3d9ce27f72..824e3f4c4e 100644 --- a/src/maxtext/utils/generate_param_only_checkpoint.py +++ b/src/maxtext/utils/generate_param_only_checkpoint.py @@ -106,7 +106,7 @@ def slice_ith(input_layers): jax.tree_util.tree_map(lambda x: x.delete(), layers) - if config.decoder_block == DecoderBlockType.DEEPSEEK: + if config.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3): # Unroll dense and MoE layers separately unroll_layer_group(config.first_num_dense_layers, layer_name="dense_layers") unroll_layer_group(config.num_decoder_layers - config.first_num_dense_layers, layer_name="moe_layers") @@ -188,7 +188,7 @@ def _slice_leaf(x, sharding): decoder_shardings.pop(layer_name) jax.tree_util.tree_map(lambda x: x.delete() if hasattr(x, "delete") else None, layers) - if config.decoder_block == DecoderBlockType.DEEPSEEK: + if config.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3): unroll_layer_group(config.first_num_dense_layers, layer_name="dense_layers") unroll_layer_group(config.num_decoder_layers - config.first_num_dense_layers, layer_name="moe_layers") else: @@ -385,7 +385,7 @@ def _slice_leaf(x, spec): del decoder_annotations[layer_name] jax.tree_util.tree_map(lambda x: x.delete() if hasattr(x, "delete") else None, layers) - if config.decoder_block == DecoderBlockType.DEEPSEEK: + if config.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3): unroll_layer_group(config.first_num_dense_layers, layer_name="dense_layers") unroll_layer_group(config.num_decoder_layers - config.first_num_dense_layers, layer_name="moe_layers") else: diff --git a/src/maxtext/utils/globals.py b/src/maxtext/utils/globals.py index c06f4f4f10..83bae2e0fe 100644 --- a/src/maxtext/utils/globals.py +++ b/src/maxtext/utils/globals.py @@ -79,6 +79,7 @@ "deepseek3-671b": "deepseek-ai/DeepSeek-V3", "deepseek3.2-671b": "deepseek-ai/DeepSeek-V3.2", "deepseek4-284b": "deepseek-ai/DeepSeek-V4-Flash", + "hy3-295b": "tencent/Hy3", "gpt-oss-20b": "openai/gpt-oss-20b", "gpt-oss-120b": "openai/gpt-oss-120b", "qwen3-omni-30b-a3b": "Qwen/Qwen3-Omni-30B-A3B-Instruct", diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 769e86b4bd..c4bf9b944b 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -738,7 +738,7 @@ def calculate_routed_and_shared_ffn_tflops_per_device(config): def get_dense_moe_layers(config): """Helper function to calculate number of dense and moe layers""" - if config.decoder_block == DecoderBlockType.DEEPSEEK: + if config.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3): num_dense_layers = config.first_num_dense_layers num_moe_layers = config.num_decoder_layers - config.first_num_dense_layers return num_dense_layers, num_moe_layers @@ -1152,7 +1152,12 @@ def calculate_tflops_training_per_device(config, log=True): DecoderBlockType.QWEN3_5, DecoderBlockType.GEMMA4, DecoderBlockType.DEEPSEEK4, + DecoderBlockType.HY3, ): + # Hy3 has DeepSeek's routed + shared + leading-dense structure. The + # generic fallback below sizes experts with mlp_dim (the dense width) + # instead of moe_mlp_dim and skips the shared expert, inflating + # reported TFLOP/s. Training itself is unaffected; only MFU reporting is. total_ffn_flops = calculate_routed_and_shared_ffn_tflops_per_device(config) is_ffn_flops_already_total = True elif config.decoder_block == DecoderBlockType.QWEN3_CUSTOM_MOE: @@ -1246,7 +1251,9 @@ def calculate_tflops_training_per_device(config, log=True): attention_tflops, learnable_weight_tflops = calculate_deepseek4_tflops_training_per_device( config, total_ffn_flops_all_layers, embedding_flops ) - elif config.decoder_block == DecoderBlockType.DEEPSEEK: + elif config.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.HY3): + # total_ffn_flops_all_layers is already summed over layers by the helper + # above; the generic branch would multiply by num_decoder_layers again. learnable_weight_tflops = ( (total_ffn_flops_all_layers + (qkv_flops + projection_flops) * config.num_decoder_layers + embedding_flops) * 3 @@ -1316,6 +1323,7 @@ def calculate_tflops_training_per_device(config, log=True): DecoderBlockType.LLAMA4, DecoderBlockType.QWEN3_NEXT, DecoderBlockType.GEMMA4, + DecoderBlockType.HY3, ): shared_flops = ( calculate_ffn_mamtul_tflops_per_device(config, get_shared_expert_mlp_dim(config)) * config.shared_experts diff --git a/tests/end_to_end/tpu/hy3/Run_Hy3.md b/tests/end_to_end/tpu/hy3/Run_Hy3.md new file mode 100644 index 0000000000..8cb8bd5b85 --- /dev/null +++ b/tests/end_to_end/tpu/hy3/Run_Hy3.md @@ -0,0 +1,205 @@ + + +# Hy3 (Tencent Hunyuan V3) + +Hy3 is an open-weights Mixture-of-Experts (MoE) model released by Tencent ([tencent/Hy3](https://huggingface.co/tencent/Hy3)). +* **Architecture**: 295B total parameters (~21B active parameters per token), 80 decoder layers with a dense 1st layer (`first_num_dense_layers: 1`) followed by 79 MoE layers. +* **Attention**: Grouped-Query Attention (GQA) with QK-Norm (RMSNorm on Query and Key head vectors) and RoPE. +* **MoE Routing**: Auxiliary-loss-free routing with Sigmoid activation and bias, 192 routed experts + 1 shared expert (selecting top-8 routed experts per token). +* **Supported Configs**: `hy3-tiny` (testing/smoke checks), `hy3-295b` (full-scale model). + +### Known limitation: MoE load balancing does not currently function during training + +Hy3's aux-loss-free routing (`routed_bias=true`) has two optional training-time +load-balancing mechanisms, controlled by `routed_bias_update_rate` (EMA-style +router-bias update) and `load_balance_loss_weight` (gradient-based auxiliary +loss). **Neither currently works when `scan_layers=true`**: the scanned-layer +application in `nnx_decoders.py` discards `nnx.Intermediate` state (the sown +`moe_bias_updates`/`moe_lb_loss` values) before it reaches `train.py`, so both +mechanisms silently no-op -- no error, but the values also never take effect. +`routed_bias_update_rate > 0.0` additionally raises `AttributeError` when +`scan_layers=false`, since that code path assumes the single stacked +`moe_layers` attribute scanning produces, which doesn't exist for unscanned +layers (`moe_layers_0`, `moe_layers_1`, ... instead). + +This is **not specific to Hy3** -- it reproduces identically with DeepSeek V3's +own `deepseek3-tiny` config, since Hy3 shares that code path. It's a +pre-existing MaxText gap, not something introduced by this model's onboarding. +Until it's fixed upstream, leave `routed_bias_update_rate` and +`load_balance_loss_weight` at their defaults (`0.0`, matches `hy3-tiny.yml`/ +`hy3-295b.yml`) -- inference and plain next-token-loss training are unaffected, +this only concerns the two optional load-balancing signals. + +--- + +## 1. Checkpoint Conversion + +### Step 1: Download Model Weights from Hugging Face +Hy3's checkpoint is ~598GB, so make sure the target disk/host has enough free space +before downloading. + +```bash +hf download tencent/Hy3 --local-dir /tmp/hy3-hf +``` + +Alternatively, `to_maxtext.py` defaults to `--lazy_load_tensors=True`, which fetches +tensors on demand directly from the `tencent/Hy3` HF Hub repo -- you can skip this +manual download step entirely and pass `model_name=hy3-295b` without +`--hf_model_path`/`--eager_load_method` in Step 2. This is the only practical option +for a partial/smoke conversion (e.g. a config truncated to a handful of layers via +`base_num_decoder_layers=N override_model_config=True`), since it only downloads the +specific shards needed instead of the full repo. + +### Step 2: Convert to MaxText Orbax Format +Use MaxText's unified `to_maxtext` conversion tool to produce an Orbax checkpoint. + +* **For Training / Fine-tuning (Scanned format, `scan_layers=true`)**: +```bash +python3 -m maxtext.checkpoint_conversion.to_maxtext \ + src/maxtext/configs/base.yml \ + model_name=hy3-295b \ + scan_layers=true \ + attention=dot_product \ + base_output_directory=${BASE_OUTPUT_PATH} \ + hf_access_token=${HF_TOKEN} \ + hardware=cpu \ + skip_jax_distributed_system=True \ + --hf_model_path=/tmp/hy3-hf \ + --eager_load_method=safetensors \ + --save_dtype=bfloat16 +``` + +* **For Decoding / Inference (Unscanned format, `scan_layers=false`)**: +```bash +python3 -m maxtext.checkpoint_conversion.to_maxtext \ + src/maxtext/configs/base.yml \ + model_name=hy3-295b \ + scan_layers=false \ + attention=dot_product \ + base_output_directory=${BASE_OUTPUT_PATH} \ + hf_access_token=${HF_TOKEN} \ + hardware=cpu \ + skip_jax_distributed_system=True \ + --hf_model_path=/tmp/hy3-hf \ + --eager_load_method=safetensors \ + --save_dtype=bfloat16 +``` + +--- + +## 2. Pre-training + +### Smoke Test (Local / Single TPU / CPU with `hy3-tiny`) +```bash +python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \ + model_name=hy3-tiny \ + steps=10 \ + per_device_batch_size=1 \ + dataset_type=synthetic \ + base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ + run_name=hy3_tiny_smoke_test +``` + +### Full-Scale Pre-training (`hy3-295b` on Multi-Slice TPU v5p) +Example training run on TPU v5p-256 / v5p-512: +```bash +python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \ + model_name=hy3-295b \ + base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ + run_name=hy3_295b_pretraining \ + per_device_batch_size=1 \ + max_target_length=4096 \ + ici_fsdp_parallelism=64 \ + ici_expert_parallelism=4 \ + megablox=true \ + sparse_matmul=true \ + attention=flash \ + dtype=bfloat16 \ + weight_dtype=bfloat16 \ + dataset_type=synthetic \ + steps=100 +``` + +--- + +## 3. Fine-tuning + +After converting the checkpoint to scanned Orbax format, you can fine-tune with your own dataset: +```bash +python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \ + model_name=hy3-295b \ + load_parameters_path=${CONVERTED_ORBAX_PATH} \ + scan_layers=true \ + base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ + run_name=hy3_295b_sft \ + per_device_batch_size=1 \ + max_target_length=4096 \ + dataset_type=huggingface \ + hf_path=${HF_DATASET_PATH} \ + learning_rate=2e-5 \ + steps=1000 +``` + +--- + +## 4. Inference / Text Generation (Decode) + +To perform text generation using the unscanned checkpoint: +```bash +python3 -m maxtext.inference.decode src/maxtext/configs/base.yml \ + model_name=hy3-295b \ + load_parameters_path=${CONVERTED_ORBAX_UNSCANNED_PATH} \ + scan_layers=false \ + max_predict_length=256 \ + prompt="Hello, who are you?" +``` + +--- + +## 5. Verification: Forward Pass Logit Check + +Verify numerical equivalence between MaxText and Hugging Face reference: +```bash +python3 tests/utils/forward_pass_logit_checker.py \ + src/maxtext/configs/base.yml \ + --run_hf_model=True \ + --hf_model_path=/tmp/hy3-hf \ + model_name=hy3-295b \ + scan_layers=false \ + weight_dtype=float32 \ + dtype=float32 \ + activations_in_float32=true \ + matmul_precision=float32 \ + float32_logits=true \ + float32_qk_product=true +``` + +**Note on full-scale (80-layer) runs:** `--run_hf_model=True` loads the real +Hugging Face PyTorch reference model on the host CPU (not on TPU chips), so it +needs enough host RAM to hold the full model -- roughly 590GB for Hy3's 295B +parameters in bf16 (more in float32, as used by the command above). A standard +single TPU-VM host will not have enough RAM for this; you'd need either a +large-memory VM or a sharded/distributed PyTorch loading setup. For a +resource-friendly sanity check, truncate both the MaxText config and the HF +reference config to a handful of layers (e.g. `base_num_decoder_layers=2 +override_model_config=True` on the MaxText side, and +`AutoConfig.from_pretrained(..., num_hidden_layers=2)` plus a golden-logits jsonl +generated from that truncated HF model on the reference side, since +`--run_hf_model=True` does not support truncating the live-loaded reference model) +-- since Hy3 has no per-layer-varying architecture, a truncated-layer comparison +still exercises every distinct code path (attention, dense MLP, MoE routing, +shared expert). diff --git a/tests/unit/configs_test.py b/tests/unit/configs_test.py index 2a7bd0f660..16e1eabbdd 100644 --- a/tests/unit/configs_test.py +++ b/tests/unit/configs_test.py @@ -266,7 +266,20 @@ def test_kimi_configs(config_file): run_config_validation(config_file) -# --- Test Group 9: Inference-specific Configs --- +# --- Test Group 9: Hy3 Model Family --- + +HY3_CONFIGS = [ + os.path.join(CONFIGS_DIR, "models", "hy3-tiny.yml"), + os.path.join(CONFIGS_DIR, "models", "hy3-295b.yml"), +] + + +@pytest.mark.parametrize("config_file", HY3_CONFIGS) +def test_hy3_configs(config_file): + run_config_validation(config_file) + + +# --- Test Group 10: Inference-specific Configs --- INFERENCE_CONFIGS = [ os.path.join(CONFIGS_DIR, "inference", "inference.yml"), diff --git a/tests/unit/hy3_vs_reference_test.py b/tests/unit/hy3_vs_reference_test.py new file mode 100644 index 0000000000..98c0b06731 --- /dev/null +++ b/tests/unit/hy3_vs_reference_test.py @@ -0,0 +1,496 @@ +# 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 +# +# http://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. + +"""Tests validating Hy3 (Tencent Hunyuan V3) MaxText components against PyTorch references.""" + +import os +from types import SimpleNamespace +from typing import Optional, Tuple +import unittest + +from flax import nnx +import jax +from jax.experimental import mesh_utils +import jax.numpy as jnp +from jax.sharding import Mesh +import numpy as np +import torch +from torch import nn +import torch.nn.functional as F + +from maxtext.common.common_types import ( + Config, + DecoderBlockType, + MODEL_MODE_TRAIN, +) +from maxtext.configs import pyconfig +from maxtext.layers import attentions, linears, moe, normalizations +from maxtext.models import hy3 +from maxtext.utils import maxtext_utils +from maxtext.utils.globals import MAXTEXT_REPO_ROOT + + +# ============================================================================== +# PyTorch Reference Implementations for Hy3 +# ============================================================================== + + +def to_jax(pt_tensor: torch.Tensor) -> jax.Array: + """Converts PyTorch tensor to JAX array.""" + return jnp.asarray(pt_tensor.detach().cpu().numpy()) + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + """Rotates half the hidden dimensions of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Applies Rotary Position Embedding (RoPE) to query and key tensors.""" + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class Hy3RMSNorm_PT(nn.Module): + """PyTorch reference RMSNorm.""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + variance = x.pow(2).mean(-1, keepdim=True) + hidden_states = x * torch.rsqrt(variance + self.eps) + return self.weight * hidden_states + + +class Hy3RotaryEmbedding_PT(nn.Module): + """PyTorch reference Rotary Embedding.""" + + def __init__(self, dim: int, max_position_embeddings: int = 4096, base: float = 10000.0): + super().__init__() + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32) / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + # position_ids: (batch, seq_len) + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Hy3Attention_PT(nn.Module): + """PyTorch reference implementation for Hy3 Attention (GQA + QK-Norm).""" + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + rms_norm_eps: float = 1e-6, + rope_theta: float = 10000.0, + max_position_embeddings: int = 4096, + ): + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_attention_heads + self.num_kv_heads = num_key_value_heads + self.head_dim = head_dim + self.num_key_value_groups = self.num_heads // self.num_kv_heads + self.scaling = head_dim**-0.5 + + self.q_proj = nn.Linear(hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, hidden_size, bias=False) + + self.q_norm = Hy3RMSNorm_PT(self.head_dim, eps=rms_norm_eps) + self.k_norm = Hy3RMSNorm_PT(self.head_dim, eps=rms_norm_eps) + self.rotary_emb = Hy3RotaryEmbedding_PT(self.head_dim, max_position_embeddings, rope_theta) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + bsz, q_len, _ = hidden_states.shape + + query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + + # QK-Norm per head + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + # Apply RoPE + cos, sin = self.rotary_emb(value_states, position_ids) + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + # Repeat KV for GQA + if self.num_key_value_groups > 1: + key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1) + value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_output = torch.matmul(attn_weights, value_states) + + attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, self.num_heads * self.head_dim) + return self.o_proj(attn_output) + + +class Hy3MLP_PT(nn.Module): + """PyTorch reference implementation for Hy3 Dense SwiGLU MLP.""" + + def __init__(self, hidden_size: int, intermediate_size: int): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class Hy3MoE_PT(nn.Module): + """PyTorch reference implementation for Hy3 MoE with shared expert and sigmoid+bias routing.""" + + def __init__( + self, + hidden_size: int, + moe_intermediate_size: int, + shared_intermediate_size: int, + num_experts: int, + top_k: int, + routed_scaling_factor: float = 1.0, + ): + super().__init__() + self.num_experts = num_experts + self.top_k = top_k + self.routed_scaling_factor = routed_scaling_factor + + # Router: gate with bias + self.gate = nn.Linear(hidden_size, num_experts, bias=True) + + # Routed Experts + self.experts = nn.ModuleList([Hy3MLP_PT(hidden_size, moe_intermediate_size) for _ in range(num_experts)]) + + # Shared Expert + self.shared_expert = Hy3MLP_PT(hidden_size, shared_intermediate_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # hidden_states: (bsz, seq_len, hidden_size) + orig_shape = hidden_states.shape + x_flat = hidden_states.view(-1, orig_shape[-1]) + + # Shared expert output + shared_out = self.shared_expert(x_flat) + + # Routing logits with bias + router_logits = self.gate(x_flat) + scores = torch.sigmoid(router_logits) + + # Top-K selection + topk_weights, topk_indices = torch.topk(scores, self.top_k, dim=-1) + + # Renormalize weights + topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20) + topk_weights = topk_weights * self.routed_scaling_factor + + # Expert computation + routed_out = torch.zeros_like(x_flat) + for i in range(self.top_k): + idx = topk_indices[:, i] + weight = topk_weights[:, i : i + 1] + for exp_id in range(self.num_experts): + mask = idx == exp_id + if mask.any(): + exp_input = x_flat[mask] + exp_out = self.experts[exp_id](exp_input) + routed_out[mask] += exp_out * weight[mask] + + total_out = (routed_out + shared_out).view(orig_shape) + return total_out + + +class Hy3DenseLayer_PT(nn.Module): + """PyTorch reference implementation for Hy3 Dense Decoder Layer.""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + rms_norm_eps: float = 1e-6, + rope_theta: float = 10000.0, + ): + super().__init__() + self.input_layernorm = Hy3RMSNorm_PT(hidden_size, eps=rms_norm_eps) + self.self_attn = Hy3Attention_PT(hidden_size, num_heads, num_kv_heads, head_dim, rms_norm_eps, rope_theta=rope_theta) + self.post_attention_layernorm = Hy3RMSNorm_PT(hidden_size, eps=rms_norm_eps) + self.mlp = Hy3MLP_PT(hidden_size, intermediate_size) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # Self Attention with pre-norm & residual + normed_attn_in = self.input_layernorm(hidden_states) + attn_out = self.self_attn(normed_attn_in, position_ids, attention_mask) + hidden_states = hidden_states + attn_out + + # MLP with post-attention norm & residual + normed_mlp_in = self.post_attention_layernorm(hidden_states) + mlp_out = self.mlp(normed_mlp_in) + hidden_states = hidden_states + mlp_out + + return hidden_states + + +class Hy3MoELayer_PT(nn.Module): + """PyTorch reference implementation for Hy3 MoE Decoder Layer.""" + + def __init__( + self, + hidden_size: int, + moe_intermediate_size: int, + shared_intermediate_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + num_experts: int, + top_k: int, + rms_norm_eps: float = 1e-6, + rope_theta: float = 10000.0, + ): + super().__init__() + self.input_layernorm = Hy3RMSNorm_PT(hidden_size, eps=rms_norm_eps) + self.self_attn = Hy3Attention_PT(hidden_size, num_heads, num_kv_heads, head_dim, rms_norm_eps, rope_theta=rope_theta) + self.post_attention_layernorm = Hy3RMSNorm_PT(hidden_size, eps=rms_norm_eps) + self.moe_block = Hy3MoE_PT( + hidden_size=hidden_size, + moe_intermediate_size=moe_intermediate_size, + shared_intermediate_size=shared_intermediate_size, + num_experts=num_experts, + top_k=top_k, + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # Self Attention with pre-norm & residual + normed_attn_in = self.input_layernorm(hidden_states) + attn_out = self.self_attn(normed_attn_in, position_ids, attention_mask) + hidden_states = hidden_states + attn_out + + # MoE block with post-attention norm & residual + normed_moe_in = self.post_attention_layernorm(hidden_states) + moe_out = self.moe_block(normed_moe_in) + hidden_states = hidden_states + moe_out + + return hidden_states + + +# ============================================================================== +# Test Suite +# ============================================================================== + + +class Hy3VsReferenceTest(unittest.TestCase): + """Unit tests comparing MaxText Hy3 components against PyTorch reference.""" + + def setUp(self): + super().setUp() + torch.manual_seed(42) + torch.set_grad_enabled(False) + + self.batch_size = 2 + self.seq_len = 8 + self.hidden_size = 128 + self.head_dim = 32 + self.num_heads = 4 + self.num_kv_heads = 2 + self.mlp_dim = 256 + self.moe_mlp_dim = 128 + self.num_experts = 4 + self.num_experts_per_tok = 2 + self.rms_norm_eps = 1e-6 + + # Base test configuration for MaxText (must be initialized first before JAX operations) + base_config_path = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "src", "maxtext", "configs", "base.yml") + ) + self.jax_config = pyconfig.initialize( + ["", base_config_path], + model_name="hy3-tiny", + override_model_config=True, + base_emb_dim=self.hidden_size, + head_dim=self.head_dim, + base_num_query_heads=self.num_heads, + base_num_kv_heads=self.num_kv_heads, + base_mlp_dim=self.mlp_dim, + base_moe_mlp_dim=self.moe_mlp_dim, + num_experts=self.num_experts, + num_experts_per_tok=self.num_experts_per_tok, + shared_experts=1, + routed_scaling_factor=1.0, + routed_score_func="sigmoid", + routed_bias=True, + use_qk_norm=True, + attention="dot_product", + matmul_precision="highest", + max_target_length=self.seq_len, + dtype="float32", + weight_dtype="float32", + float32_logits=True, + float32_qk_product=True, + dropout_rate=0.0, + ) + + # JAX mesh + devices_array = maxtext_utils.create_device_mesh(self.jax_config) + self.mesh = Mesh(devices_array, self.jax_config.mesh_axes) + + # Initial inputs + self.pt_input = torch.randn(self.batch_size, self.seq_len, self.hidden_size, dtype=torch.float32) + self.pt_positions = torch.arange(self.seq_len, dtype=torch.long).unsqueeze(0).expand(self.batch_size, -1) + self.jax_input = to_jax(self.pt_input) + self.jax_positions = to_jax(self.pt_positions) + self.jax_segment_ids = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32) + + def test_hy3_dense_mlp_vs_reference(self): + """Validates Hy3 Dense MLP (SwiGLU) against PyTorch reference.""" + pt_mlp = Hy3MLP_PT(self.hidden_size, self.mlp_dim) + pt_out = pt_mlp(self.pt_input) + + # MaxText MlpBlock + rngs = nnx.Rngs(0) + jax_mlp = linears.MlpBlock( + in_features=self.hidden_size, + intermediate_dim=self.mlp_dim, + activations=["silu", "linear"], + intermediate_dropout_rate=0.0, + dtype=jnp.float32, + weight_dtype=jnp.float32, + config=self.jax_config, + model_mode=MODEL_MODE_TRAIN, + mesh=self.mesh, + rngs=rngs, + ) + + # Copy weights + # MaxText wi_0 is gate, wi_1 is up, wo is down + jax_mlp.wi_0.kernel.value = to_jax(pt_mlp.gate_proj.weight.t()) + jax_mlp.wi_1.kernel.value = to_jax(pt_mlp.up_proj.weight.t()) + jax_mlp.wo.kernel.value = to_jax(pt_mlp.down_proj.weight.t()) + + jax_out = jax_mlp(self.jax_input, deterministic=True) + + np.testing.assert_allclose(to_jax(pt_out), jax_out, rtol=1e-3, atol=1e-3) + + def test_hy3_dense_layer_vs_reference(self): + """Validates full Hy3DenseLayer forward pass against PyTorch reference.""" + pt_layer = Hy3DenseLayer_PT( + hidden_size=self.hidden_size, + intermediate_size=self.mlp_dim, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + rms_norm_eps=self.rms_norm_eps, + rope_theta=float(self.jax_config.rope_max_timescale), + ) + + # Create causal mask for PyTorch + mask = torch.triu(torch.ones(self.seq_len, self.seq_len, dtype=torch.bool), diagonal=1) + causal_mask = torch.zeros(self.seq_len, self.seq_len, dtype=torch.float32).masked_fill(mask, -1e9) + causal_mask = causal_mask.unsqueeze(0).unsqueeze(0) # (1, 1, seq_len, seq_len) + pt_out = pt_layer(self.pt_input, self.pt_positions, attention_mask=causal_mask) + + # MaxText Hy3DenseLayer + rngs = nnx.Rngs(0) + jax_layer = hy3.Hy3DenseLayer( + config=self.jax_config, + model_mode=MODEL_MODE_TRAIN, + mesh=self.mesh, + rngs=rngs, + layer_idx=0, + ) + + # Copy weights: Norms + jax_layer.pre_self_attention_layer_norm.scale.value = to_jax(pt_layer.input_layernorm.weight) + jax_layer.post_self_attention_layer_norm.scale.value = to_jax(pt_layer.post_attention_layernorm.weight) + + # Copy weights: Attention QKV & Out + jax_layer.self_attention.query.kernel.value = to_jax( + pt_layer.self_attn.q_proj.weight.t().view(self.hidden_size, self.num_heads, self.head_dim) + ) + jax_layer.self_attention.key.kernel.value = to_jax( + pt_layer.self_attn.k_proj.weight.t().view(self.hidden_size, self.num_kv_heads, self.head_dim) + ) + jax_layer.self_attention.value.kernel.value = to_jax( + pt_layer.self_attn.v_proj.weight.t().view(self.hidden_size, self.num_kv_heads, self.head_dim) + ) + jax_layer.self_attention.out.kernel.value = to_jax( + pt_layer.self_attn.o_proj.weight.t().view(self.num_heads, self.head_dim, self.hidden_size) + ) + jax_layer.self_attention.query_norm.scale.value = to_jax(pt_layer.self_attn.q_norm.weight) + jax_layer.self_attention.key_norm.scale.value = to_jax(pt_layer.self_attn.k_norm.weight) + + # Copy weights: MLP + jax_layer.mlp.wi_0.kernel.value = to_jax(pt_layer.mlp.gate_proj.weight.t()) + jax_layer.mlp.wi_1.kernel.value = to_jax(pt_layer.mlp.up_proj.weight.t()) + jax_layer.mlp.wo.kernel.value = to_jax(pt_layer.mlp.down_proj.weight.t()) + + jax_out = jax_layer( + inputs=self.jax_input, + decoder_segment_ids=self.jax_segment_ids, + decoder_positions=self.jax_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + if isinstance(jax_out, tuple): + jax_out = jax_out[0] + + np.testing.assert_allclose(to_jax(pt_out), jax_out, rtol=1e-3, atol=1e-3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index fcd5acb5cc..53d2c33439 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -1080,6 +1080,53 @@ def make_block(): self.assertEqual(len(updated_kvs), num_layers) np.testing.assert_allclose(y_external, y_scanned, rtol=1e-5, atol=1e-5) + def test_hy3_decoder_forward(self): + """Test NNXDecoder with hy3 block (1 dense layer + 3 MoE layers).""" + cfg = _make_config( + model_name="hy3-tiny", + decoder_block="hy3", + scan_layers=False, + num_decoder_layers=4, + first_num_dense_layers=1, + base_emb_dim=128, + base_num_query_heads=4, + base_num_kv_heads=2, + head_dim=32, + base_mlp_dim=256, + base_moe_mlp_dim=128, + num_experts=4, + num_experts_per_tok=2, + shared_experts=1, + routed_scaling_factor=1.0, + routed_score_func="sigmoid", + routed_bias=True, + use_qk_norm=True, + sparse_matmul=False, + vocab_size=256, + ) + mesh = _make_mesh(cfg) + decoder = NNXDecoder( + config=cfg, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=self.rngs, + ) + shared_embedding = self._make_shared_embedding(cfg) + ids, segment_ids, positions = self._make_token_inputs(cfg) + + logits, _, _ = decoder( + shared_embedding, + ids, + positions, + decoder_segment_ids=segment_ids, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + self.assertEqual( + logits.shape, + (cfg.global_batch_size_to_train_on, cfg.max_target_length, cfg.vocab_size), + ) + @pytest.mark.tpu_only class TestGemma4SmallNNXDecoder(unittest.TestCase):