[Feat] Add native GLM-5.2 MoE training support#1967
Open
jayhenry wants to merge 29 commits into
Open
Conversation
jayhenry
commented
Jul 22, 2026
jayhenry
force-pushed
the
glm52-upstream-pr
branch
from
July 23, 2026 14:57
5947f73 to
1fbb09d
Compare
jayhenry
force-pushed
the
glm52-upstream-pr
branch
from
July 23, 2026 15:43
719e8be to
b17eaf0
Compare
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.
Summary
This supersedes the implementation in #1956 with a smaller, refactored series based directly on current
upstream/main.Validation
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.