Skip to content

[Feat] Add native GLM-5.2 MoE training support#1967

Open
jayhenry wants to merge 29 commits into
InternLM:mainfrom
jayhenry:glm52-upstream-pr
Open

[Feat] Add native GLM-5.2 MoE training support#1967
jayhenry wants to merge 29 commits into
InternLM:mainfrom
jayhenry:glm52-upstream-pr

Conversation

@jayhenry

@jayhenry jayhenry commented Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add GLM-5.2 chat templating and tokenization, native MoE/MTP model configuration, Hugging Face checkpoint mapping, and SFT integration.
  • Add the SparseMLA public operator interface with PyTorch reference, TileLang, and cuDNN backends, plus cross-layer DSA top-k sharing and cache lifecycle management.
  • Support FP8, torch.compile, activation and DSA top-k offload, EP/SP, all-to-all and DeepEP dispatchers, intra-layer micro-batching, and shared-weight MTP index sharing.
  • Add model, operator, distributed engine, sequence-parallel, checkpoint load/save, and end-to-end SFT regression coverage.

This supersedes the implementation in #1956 with a smaller, refactored series based directly on current upstream/main.

Validation

  • GLM-5.2 tokenization: 7 passed
  • GLM config/model/MTP/micro-batch suites: 42 passed
  • DSA and SparseMLA suites: 38 passed
  • MoE train engine: 11 passed
  • SFT smoke: 1 passed
  • Focused physical-MTP indexer and cache lifecycle regression: 8 passed
  • Pre-commit hooks passed, including ruff, formatting, mypy, codespell, and Pydantic checks.

An 8 GPUs GLM-5.2 30B MTP run was validated for 10 steps with 16K packs, FP8, EP4, cuDNN DSA, compile, activation offload, DSA top-k offload, and intra-layer micro-batch=2. Steps 2-10 averaged 14.62K global tokens/s with 122.665 GiB peak memory per rank; final LLM/MTP losses were 10.062593 and 0.955561 and aligned with the pre-refactor baseline. SP1/SP2/SP4/SP8 configurations were also validated.

Known limitation

Shared-weight MTP with more than one logical prediction depth and more than one intra-layer micro-batch remains unsupported with reentrant checkpointing because backward replay can revisit a released autograd graph. The native GLM-5.2 configuration uses one checkpoint-backed MTP layer; this experimental combination is outside this PR scope.

Comment thread xtuner/v1/data_proto/sequence_context.py Outdated
Comment thread xtuner/v1/model/moe/glm52.py
Comment thread xtuner/v1/module/attention/dsa_topk_sharing.py Outdated
Comment thread xtuner/v1/module/attention/dsa_topk_sharing.py Outdated
Comment thread xtuner/v1/model/moe/moe.py Outdated
Comment thread xtuner/v1/model/moe/moe.py Outdated
@jayhenry
jayhenry force-pushed the glm52-upstream-pr branch from 5947f73 to 1fbb09d Compare July 23, 2026 14:57
@jayhenry
jayhenry force-pushed the glm52-upstream-pr branch from 719e8be to b17eaf0 Compare July 23, 2026 15:43
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.

1 participant