[Fix] Run GLM-5.2 shared MTP for 7 logical steps - #2072
Open
Knlife wants to merge 3 commits into
Open
Conversation
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.
Motivation
The official GLM-5.2 release uses shared recurrent MTP steps. The official blog[1] reports 7 MTP steps for both training and inference, while the released checkpoint retains
num_nextn_predict_layers=1. XTuner'sGlm52MoEConfig.from_hf()currently maps this field directly toMTPConfig.num_layers, so the canonical SFT recipe constructs 1 physical shared MTP layer but executes it for only one logical step.[1] https://huggingface.co/blog/zai-org/glm-52-blog
[2] https://huggingface.co/zai-org/GLM-5.2/raw/main/config.json
Changes
Override
examples/v1/config/sft_glm5p2.pywithPreserve
num_nextn_predict_layers=1so checkpoint construction and strict loading continue to use 1 physical MTP layer. Add a regression test for the recipe configuration. The existing shared-weight implementation inMTPBlockand GLM-5.2 model construction remains unchanged.Scope and compatibility
This change is limited to the canonical GLM-5.2 SFT recipe. Generic Hugging Face conversion, existing CE loss behavior, MTP loss scaling, and
strict_load=Truebehavior are unchanged. Other autotest and RL configurations, including explicit MTP4 cases, are intentionally left unchanged. This PR does not implement or modify the end-to-end TV loss objective.Validation
python -m py_compileandgit diff --checkpassed. A dependency-free recipe smoke test passed. Full pytest and CUDA/strict-load integration tests require the project runtime dependencies and checkpoint, so they are left to CI.