Add Hy3 (Tencent Hunyuan V3) model support - #4704
Draft
weikuo0506 wants to merge 3 commits into
Draft
Conversation
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, so it reuses DeepSeekGenericLayer's dense/MoE scaffolding and moe.RoutedAndSharedMoE rather than introducing new attention math. - Add DecoderBlockType.HY3 and wire it 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, and the relevant types.py validation guards (including the ModelName literal allowlist). - New src/maxtext/models/hy3.py: Hy3DenseLayer/Hy3MoELayer, subclassing DeepSeekGenericLayer and overriding self_attention with plain GQA (matching deepseek4.py's precedent for swapping attention modules). - New hy3-tiny.yml / hy3-295b.yml model configs. - Register Hy3 in the checkpoint conversion framework (hf_model_configs.py, param_mapping.py, hf_shape.py). The HF tensor names/shapes were verified 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. - Verified locally: full configs_test.py suite (78 tests, no regressions), and a random-init CPU forward pass on hy3-tiny.yml across scan_layers True/False and with mtp_num_layers=1. - Real-weight golden-logits verification against the full 295B checkpoint (requires TPU) is intentionally not done yet in this branch. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Missed during the initial Phase 6b checkpoint-conversion registration pass -- to_maxtext.py's model_name -> HF repo id lookup consults this dict separately from HF_MODEL_CONFIGS/PARAM_MAPPING/HOOK_FNS. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
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. |
Running the repo's pre-commit hooks (pyink, pylint, codespell, mdformat, yamllint) per CONTRIBUTING.md flagged a formatting-only reflow of the HY3 loss-free-load-balancing guard added earlier. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds support for Hy3 (Tencent Hunyuan V3,
tencent/Hy3on 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 itreuses
DeepSeekGenericLayer's dense/MoE scaffolding andmoe.RoutedAndSharedMoErather than introducing new attention math — the new decoder layer (
Hy3DenseLayer/Hy3MoELayerinsrc/maxtext/models/hy3.py) subclassesDeepSeekGenericLayerandoverrides
self_attentionwith plain GQA, following the same patternDeepSeek4DecoderLayeruses indeepseek4.py.Changes
DecoderBlockType.HY3wired throughdecoders.py/nnx_decoders.py(both thelegacy Linen path and the default pure-NNX path),
moe.py's DeepSeek-V3-stylerouting gates (pre-bias logits capture +
routed_scaling_factorapplication),and the relevant
types.pyvalidation guards (including theModelNameliteralallowlist and the loss-free-load-balancing decoder_block check).
src/maxtext/models/hy3.py:Hy3DenseLayer/Hy3MoELayer.hy3-tiny.yml/hy3-295b.ymlmodel configs.hf_model_configs.py,param_mapping.py,hf_shape.py,globals.py'sHF_IDS).transformersalreadyships a native
HYV3Config, so notrust_remote_code-only config workaround wasneeded. HF tensor names/shapes were cross-checked against the real
tencent/Hy3checkpoint's
model.safetensors.index.json(46,545 non-MTP tensors, zeromissing/extra). MTP layer weights are intentionally left unmapped (randomly
initialized on conversion) for this pass.
get_dense_moe_layersinmaxtext_utils.py) forHy3's dense/MoE layer split.
Status against the Model Bringup checklist
Self-assessed against docs/guides/model_bringup.md:
whole-model golden-logit comparison so far, not an isolated per-layer
vs_referenceunit test intests/unit/).get_dense_moe_layers).(
hy3-tiny.yml,scan_layers=True), not with real weights.### Hy3entry todocs/reference/models/supported_models_and_architectures.md, but have notwritten a
Run_Hy3.md-style end-to-end guide or posted an announcement.Known limitation: full-scale (295B, 80-layer) verification not done
Golden-logits verification was run against the real checkpoint, truncated to the
first 2 layers (1 dense + 1 MoE, with real routing/bias/shared-expert weights) —
see Tests below. The full 80-layer / 295B checkpoint was not verified end-to-end:
it requires downloading the full ~598GB checkpoint and a multi-chip TPU pod, and
running the real HF PyTorch reference model for comparison needs ~590GB of host RAM
(exceeding a single TPU VM host), which would need either a large-memory VM or a
sharded/distributed PyTorch loading setup. Since Hy3 has no per-layer-varying
architecture (unlike e.g. DeepSeek V4's per-layer
compress_ratios), the 2-layerresult exercises every distinct code path (attention, dense MLP, MoE routing,
shared expert) and should generalize, but this is flagged here for reviewers.
Tests
pytest tests/unit/configs_test.py— full suite (78 tests) passes, including thenew
test_hy3_configsand no regressions intest_deepseek_configs/test_qwen_configs.hy3-tiny.yml, bothscan_layers=True/Falseandwith
mtp_num_layers=1(MTP reuses the last decoder layer class generically, permodels.py's existing wiring — verified this works for Hy3 too).scan_layers=False: convertedthe real
tencent/Hy3checkpoint truncated to 2 layers viato_maxtext.py,generated golden logits from the real HF reference model (also truncated to 2
layers, same real weights), and compared via
forward_pass_logit_checker.py:overlap 10/10
(Both far under the
--max_kl_div=0.5threshold used; the guide's stricter1e-4target is for the full float32 comparison mode, not attempted here.)Checklist
### Hy3entry todocs/reference/models/supported_models_and_architectures.md;have not yet added a
Run_Hy3.mde2e guide since the full-scale run is pending).