GLM-5.2 131k: dsv4-next bridge pin + vectorized FP8 dequant - #12
Closed
JackRao123 wants to merge 11 commits into
Closed
GLM-5.2 131k: dsv4-next bridge pin + vectorized FP8 dequant#12JackRao123 wants to merge 11 commits into
JackRao123 wants to merge 11 commits into
Conversation
…karound Re-add the TODO comments noting the cache_position forward docstrings are a local workaround to be removed once this branch includes NVIDIA-NeMo/Megatron-Bridge@0a21da47 or a later upstream fix. The docstrings themselves now come from upstream; only the workaround markers are local. Signed-off-by: Jerry Hong <jerry.hong@baseten.co> Co-authored-by: Cursor <cursoragent@cursor.com>
Enable the GLM-5.2 fused DSA defaults in the Bridge GLM provider without carrying temporary Megatron-LM submodule pointer hops.
Point the Megatron-LM submodule at the rewritten trainers-main tip, where the Baseten-specific Megatron-LM patches are a clean top stack.
…Kimi-K2 transformers' get_cached_module_file() copies a remote-code module plus only its direct (1-level) relative imports into the transformers_modules cache, but get_class_in_module() resolves relative imports recursively. Remote code with a >1-level relative-import chain (Kimi-K2: modeling_kimi_k25 -> modeling_deepseek -> configuration_deepseek) leaves the transitive file uncopied and loading fails with FileNotFoundError for the deepest dependency. Make check_imports() return the full transitive set so every relatively-imported sibling is materialized.
…nfig (#10) GLM-5.2 declares indexer_rope_interleave: true — its DSA indexer uses GPT-J interleaved RoPE, which serving engines honor (vLLM: is_neox_style = not flag). The bridge never plumbed it, so Megatron's indexer defaulted to the DeepSeek-V3.2 non-interleaved convention and its top-k selection drifted from inference progressively with sequence length (trainer<->vLLM logprob KL 0.010 -> 0.084 between 9k and 15k tokens; 0.009 flat after the fix). Bumps 3rdparty/Megatron-LM to pick up dsa_indexer_rope_interleave (basetenlabs/Megatron-LM#6).
#10 landed with the submodule pointing at the pre-squash PR-branch head (377e4f5e); repoint to the squash-merged commit on trainers-main (2229188e, identical tree) so the gitlink doesn't depend on the PR branch existing.
JackRao123
changed the title
Extracted from the devbox working-changes commit; the Megatron-LM pin bump it also carried is dropped (superseded by the final pin commit).
dequantize_mxfp4_e2m1_packed ran on CPU (int64 nibble-unpack + torch.stack + repeat_interleave + multiply) over every expert shard at load — ~150s of get_model on DeepSeek-V4-Flash, redone every boot. Under DSV4_GPU_DEQUANT=1, move the packed weight + scale to CUDA before dequant so it runs on GPU (~10-50x faster); weights are copied to GPU right after anyway, and per-param usage keeps GPU memory bounded. Validated: boot 358s -> 177s, loss unchanged. CPU path unchanged when env unset. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PDQHqx2q8YGuUCVKEYN4Wr Signed-off-by: Jerry Hong <jerry.hong@baseten.co>
The per-block Python loop made GLM-5.2-FP8's 800B load CPU-bound (8 workers pegged ~50 min). Two repeat_interleaves + one multiply, verified bit-exact against the loop. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…sed DSA contiguous-CP adapter) Feature-branch pin while GLM-5.2 131k dev continues; repin to the merged Megatron-LM trainers-main SHA before this PR merges. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
JackRao123
force-pushed
the
jackrao/glm-cp-bridge
branch
from
July 8, 2026 23:55
2d3bdf4 to
57f1a0a
Compare
The qk_rope_head_dim workaround (transformers GlmMoeDsaConfig collapses the qk head-dim split, corrupting kv_a_proj shapes: 704 != 576) only read <base_model>/config.json as a filesystem path, so it silently no-opped when base_model is a hub repo id and the 800B weight load failed. Resolve the raw config through hf_hub_download when the local read misses — transformers has already cached config.json by the time the bridge runs, so this works offline (HF_HUB_OFFLINE) too. Warn loudly when neither path resolves. Prod launches GLM-5.2-FP8 by HF id, so this unblocks the registry row in trainers#592. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Author
|
Superseded by #17 (jackrao/glm-cp-bridge-20260907, based on the upstream-rebased trainers-main-20260907 / #16): the FP8-dequant vectorization and the hub-cache raw-config fix were cherry-picked there, and the contiguous-CP Megatron-LM pin this PR carried (LM#7) is obsolete on the upstream fused-DSA stack. trainers#592 now pins #17's head. |
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
2 participants
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.
Purpose
One diff against
trainers-mainshowing every megatron-bridge change needed for GLM-5.2 131k CP. Retargeted per Jack; the diff includes thetrainers-main-next-dsv4bridge line this stack sits on (jerry's — MXFP4 GPU dequant, FP8 split recompute, Megatron pin bumps), plus my commits:dequantize_fp8_blockwise's per-128×128-block Python loop made the 800B GLM-5.2-FP8 load CPU-bound (~60 CPU-min/rank). Tworepeat_interleaves + one multiply — bit-exact (torch.equal-verified incl. non-divisible edges), ~17×/tensor; end-to-end load 60 → ~15 min.0ce18605(FP8 dequant-on-load, qk_rope_head_dim workaround, IndexShare wiring) — needed because this line's pin predates it; disappears on the next rebase. This part of the diff may show as no-op vs trainers-main (it IS mainline's file).Validation
Carried the overnight GLM-5.2 131k profiling stack (basetenlabs/trainers#592): 131k LoRA-SFT fwd-bwd on 4×8 B200 via TP1/PP1/EP32/CP32.
🤖 Generated with Claude Code