Add Hy3 (Tencent Hunyuan V3) model support

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:**
 

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,

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,

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,

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):