GLM-5.2: full/shared indexer typing for glm_moe_dsa (DSA schedule + interleaved indexer rope)#1463
Draft
machiabeli wants to merge 2 commits into
Draft
GLM-5.2: full/shared indexer typing for glm_moe_dsa (DSA schedule + interleaved indexer rope)#1463machiabeli wants to merge 2 commits into
machiabeli wants to merge 2 commits into
Conversation
…ad + dense path) The stub subclassed deepseek_v32.py, whose Attention instantiates an Indexer on EVERY layer. GLM-5.2's checkpoint only carries indexer weights on the 21 "full" layers (config indexer_types = 21 full + 57 shared, 78 total), so load failed with 285 missing indexer params (57 shared x 5 tensors). Fix, entirely by subclassing in glm_moe_dsa.py (deepseek_v32.py UNTOUCHED): - ModelArgs: parse indexer_types, index_share_for_mtp_iteration, index_topk_freq, index_skip_topk_offset. Defensive: absent indexer_types => all layers "full" (vanilla DeepSeek-V3.2 behavior identical). - GlmMoeDsaAttention: full layers own an Indexer (parent semantics); shared layers `del self.indexer` so the param tree matches the checkpoint exactly (0 missing / 0 unexpected, verified by dry tree-diff: 2489 == 2489 keys). Threads prev_topk_indices and returns topk for the next layer when the next layer is "shared" (mirrors HF next_skip_topk). Below index_topk (2048) tokens the indexer short-circuits to None on every layer => full dense path. - GlmMoeDsaModel: rebuilds the decoder stack with indexer-typed layers and threads prev_topk through the layer loop (wired now; M2 verifies >2048). - make_cache: full layers => CacheList(KVCache, KVCache); shared layers => CacheList(KVCache) (no indexer slot). M2 readiness: prev_topk threading is fully wired through attention + the model loop; only the >2048 sparse path remains to be verified. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014EmGUUAsEwp5sUwBBrhVVc
…parse path) - indexer_rope_interleave defaults True for GLM-5.2 (the arch keeps the pre-ml-explore#1431 interleaved convention; validated against the real 743B mxfp4 checkpoint — engaged-DSA generation at 3k context is coherent and an independent C/Metal engine reproduces the top-k selection exactly with interleaved rope). - test_all_models entry with indexer_types=["full","shared","full","shared"] (exercises shared-layer construction + asymmetric CacheList at B=1/B=2). - test_glm_moe_dsa_shared_schedule: full layers own indexers, shared layers don't (checkpoint ships no weights for them), interleaved rope flag, CacheList(kv,indexer) vs CacheList(kv), and the ENGAGED path — prefill L > index_topk then a decode step threading prev_topk full->shared, finite outputs. The sparse (>topk) path is additionally cross-validated externally: a C reference matched this implementation's engaged logits at 1e-6 on full AND shared schedules (tiny config), and per-layer parity probes on the real checkpoint at 3015-token context selected identical top-2048 sets.
machiabeli
force-pushed
the
glm52-shared-indexer
branch
from
0fffab6 to
f2a8ed5
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
1 participant
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.
GLM-5.2 ships a per-layer DSA indexer schedule (
indexer_types: 21fulllayers own a lightning indexer; 57sharedlayers reuse the previous full layer's top-k and ship no indexer weights). The currentglm_moe_dsa.pysubclasses DeepSeek-V3.2 directly, so every layer expects indexer weights — strict loading of the real checkpoint fails (#1418), and pastindex_topktokens the schedule semantics are lost (#1453).What this adds
GlmMoeDsaAttention/ decoder / model honoringindexer_types: full layers construct theIndexer; shared layersdelit (param tree matches the checkpoint — 0 missing / 0 unexpected on strict load of the real 743B mxfp4 checkpoint, 21 full + 57 shared) and consumeprev_topk_indicesthreaded layer-to-layer (mirrors HF'snext_skip_topk).make_cache:CacheList(kv, indexer)on full layers,CacheList(kv)on shared.indexer_rope_interleave: bool = Truefor this arch — carries forward the merged Fix deepseek indexer #1431 resolution (GLM-5.2 keeps interleaved; per Fix deepseek indexer #1431 discussion and [glm-mode-dsa] Indexer uses interleaved rope huggingface/transformers#46842, the HF GLM-5.2 reference is the suspected outlier here). Independent validation below agrees with interleaved.test_all_modelsentry with a["full","shared","full","shared"]schedule (B=1/B=2), plustest_glm_moe_dsa_shared_schedulecovering schedule construction, the cache asymmetry, the rope flag, and the engaged sparse path (prefill pastindex_topk, then a decode step threadingprev_topk).Validation beyond the unit tests
load_modelof the real GLM-5.2 743B mxfp4 checkpoint: 78 layers, zero random weights, schedule correct.fullandfull/sharedschedules at tiny config, and per-layer parity probes on the real checkpoint at 3015-token context reproduce the top-2048 selections exactly (interleaved rope), with coherent long-context generation.Relation to #1410 / #1412 / #1419
#1410 (@pcuenca) implements the same IndexShare core and has broad tester traction; #1410's remaining open item is numerical-parity evidence, which is exactly what the engaged-path tests and the independent-engine validation here provide. Happy to consolidate in whichever direction the maintainers prefer — these tests and the validation story are offered either way.
🤖 Generated with Claude Code
https://claude.ai/code/session_011YpPndSoRJZRbZeCNs8KYL