Add Hy3 (Tencent Hunyuan V3) model support

[weikuo0506]

Description

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 has no MLA and no compressed/sparse attention, so 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 same pattern
DeepSeek4DecoderLayer uses in deepseek4.py.

Changes

• DecoderBlockType.HY3 wired through decoders.py/nnx_decoders.py (both the
  legacy Linen path and the default pure-NNX path), moe.py's DeepSeek-V3-style
  routing gates (pre-bias logits capture + routed_scaling_factor application),
  and the relevant types.py validation guards (including the ModelName literal
  allowlist and the loss-free-load-balancing decoder_block check).
• New src/maxtext/models/hy3.py: Hy3DenseLayer/Hy3MoELayer.
• New hy3-tiny.yml / hy3-295b.yml model configs.
• Registered Hy3 in the checkpoint conversion framework (hf_model_configs.py,
  param_mapping.py, hf_shape.py, globals.py's HF_IDS). transformers
  already ships a native HYV3Config; hf_model_configs.py uses it via the
  same try/except native-class-with-PTConfig-fallback pattern already used
  for Gemma 4, so no trust_remote_code-only workaround is needed. HF tensor
  names/shapes were cross-checked against the real tencent/Hy3 checkpoint's
  model.safetensors.index.json (46,545 non-MTP tensors, zero missing/extra).
  MTP layer weights are intentionally left unmapped (randomly initialized on
  conversion) for this pass.
• tests/unit/hy3_vs_reference_test.py: from-scratch PyTorch reference
  implementations of Hy3's attention/MoE blocks, compared directly against the
  MaxText JAX layers (random weights, no HF download needed).
• tests/unit/nnx_decoders_test.py: NNXDecoder forward-pass test with the
  hy3 decoder block (1 dense + 3 MoE layers, unscanned).
• tests/end_to_end/tpu/hy3/Run_Hy3.md: end-to-end user guide (checkpoint
  conversion, pre-training, fine-tuning, decoding, logit verification).
• FLOPs/MFU accounting (maxtext_utils.py): get_dense_moe_layers plus 3
  more branches in calculate_tflops_training_per_device that were still
  missing HY3 and would otherwise fall through to a generic path that sizes
  experts with mlp_dim instead of moe_mlp_dim and skips the shared expert.
  Verified by measuring the function's output for hy3-295b before/after —
  learnable_weight_tflops changes measurably, confirming this wasn't a no-op.
• grpo_utils.py / generate_param_only_checkpoint.py: extended the
  dense+MoE two-stack layer-group handling to HY3, so RL (GRPO) parameter
  resharding and params-only checkpoint construction handle Hy3's
  first_num_dense_layers split correctly instead of mishandling it.
• train.py: the aux-loss-free router-bias update path hardcodes the MoE
  submodule attribute name "DeepSeekMoeBlock_0". Extended the lookup to
  include HY3 for forward-compatibility. This does not make the update path
  functional today — see "Known limitation 1" below; it only prevents Hy3
  from crashing on top of a pre-existing, unrelated bug.

Known limitation 1: MoE load balancing does not currently work during training

Both of Hy3's optional training-time load-balancing mechanisms
(routed_bias_update_rate and load_balance_loss_weight) are non-functional
with scan_layers=true (silently no-op — nnx_decoders.py's scanned-layer
application discards nnx.Intermediate state, including the sown
moe_bias_updates/moe_lb_loss, before train.py can read it), and
routed_bias_update_rate>0 additionally crashes with AttributeError when
scan_layers=false (that code path assumes a single stacked moe_layers
attribute that only exists when scanned; unscanned layers are named
moe_layers_0, moe_layers_1, ... instead).

This is not specific to Hy3. Reproduced locally (CPU) that deepseek3-tiny
with the same config overrides exhibits identical behavior in both modes,
confirming this is a pre-existing MaxText gap in the shared DeepSeek-V3-style
code path, not something this PR introduces. It is separate from the
DeepSeek-V4-specific MoEBiasVar migration in #4753 (still open as of this
writing), which explicitly leaves the legacy nnx.Param-based path — the one
Hy3 and DeepSeek V3 share — untouched.

Until fixed upstream, hy3-tiny.yml/hy3-295b.yml leave both settings at
their defaults (0.0). Inference and plain next-token-loss training are
unaffected; only these two optional load-balancing signals are. Full
explanation in the new section of Run_Hy3.md.

Verification scope: not yet tested at full 80-layer/295B distributed scale

What's verified: numerical correctness (real-checkpoint 2-layer golden-logits
comparison + a from-scratch PyTorch-reference unit test covering every
distinct code path — attention, dense MLP, MoE routing, shared expert),
basic single-host training mechanics (CPU smoke test, loss decreasing
normally), and — at the full declared 80-layer/295B config — AOT compilation
at compile_topology=v5p-256 (train_compile.py, ici_fsdp_parallelism=-1):
compiles successfully, ~70GB/device (13.8GB args + 44.4GB temp, well under
v5p's ~95GB HBM/chip). This exercises the full-scale scanned-layer
construction and sharding/mesh-partitioning logic without needing real
weights or a live TPU pod. Since Hy3 has no per-layer-varying architecture,
together these give good confidence in per-layer numerical correctness at
any layer count and in shape/sharding-spec correctness at the full 80-layer
config.

What's not verified: an actual multi-host training run at that scale
(only compiled, not executed) — in particular, expert parallelism
(ici_expert_parallelism/dcn_expert_parallelism > 1) has never been
exercised at all, even in the AOT compile above (ici_fsdp_parallelism=-1
used FSDP, not EP). Real multi-host MoE collective behavior at runtime is a
distinct risk class from both numerical correctness and compile-time
shape/sharding checks. Running this needs the full ~598GB checkpoint and a
multi-chip TPU pod; a live HF-vs-MaxText logits comparison at full scale
additionally needs ~590GB of host RAM for the PyTorch reference model
(exceeds a single TPU VM host).

Not added: Muon optimizer / layerwise-quantization allowlists

Both are optional MaxText features not required by Hy3's own architecture,
and neither has been tested against Hy3, so they're left out rather than
claiming untested support. Happy to add either in a follow-up if useful.

Tests

• pytest tests/unit/configs_test.py — full suite (78 tests) passes.
• pytest tests/unit/hy3_vs_reference_test.py — 2/2 pass: MaxText's attention
  and MoE routing layers match from-scratch PyTorch reference implementations.
• pytest tests/unit/nnx_decoders_test.py -k hy3 — 1/1 pass.
• Random-init CPU forward pass on hy3-tiny.yml, scan_layers=True/False,
  with mtp_num_layers=1 (MTP reuses the last decoder layer class generically
  per models.py's existing wiring — verified this works for Hy3 too), and a
  real 3-step training run confirming loss decreases normally.
• Real-weight verification on a GCP TPU (v5litepod-4), scan_layers=False:
  converted the real tencent/Hy3 checkpoint truncated to 2 layers via
  to_maxtext.py, generated golden logits from the real HF reference model
  (also truncated to 2 layers, same real weights), compared via
  forward_pass_logit_checker.py:
  • prompt "The capital of France is": max KL divergence 2.2e-4, top-10 token
    overlap 10/10
  • prompt "I love to": max KL divergence 4.5e-4, top-10 token overlap 8/10
    (Both far under the --max_kl_div=0.5 threshold used.)
• Confirmed no regressions: deepseek3-tiny and deepseek4-284b-family
  config/training paths behave identically before and after every change in
  this PR (re-ran test_deepseek_configs, and manually reproduced the
  aux-loss-free bias-update behavior on deepseek3-tiny to confirm it's
  unaffected by the train.py change).

Checklist

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed (added a
  ### Hy3 entry to docs/reference/models/supported_models_and_architectures.md
  and a Run_Hy3.md end-to-end guide, including the load-balancing limitation).

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.