Skip to content

[Fix] Run GLM-5.2 shared MTP for 7 logical steps - #2072

Open
Knlife wants to merge 3 commits into
InternLM:mainfrom
Knlife:fix/mtp
Open

[Fix] Run GLM-5.2 shared MTP for 7 logical steps#2072
Knlife wants to merge 3 commits into
InternLM:mainfrom
Knlife:fix/mtp

Conversation

@Knlife

@Knlife Knlife commented Sep 5, 2026

Copy link
Copy Markdown

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's Glm52MoEConfig.from_hf() currently maps this field directly to MTPConfig.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.py with

model_cfg.mtp_config = MTPConfig(
    num_layers=7,
    share_weights=True,
    detach_mtp_lm_head_weight=False,
    detach_mtp_inputs=False,
    loss_scaling_factor=0.1,
)

Preserve num_nextn_predict_layers=1 so 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 in MTPBlock and 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=True behavior 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_compile and git diff --check passed. 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants