GLM-5/5.1/5.2 DSA + LoRA enablement (DSA option: megatron-bridge-native and glm-native) - #15
Closed
yushengsu-thu wants to merge 23 commits into
Closed
GLM-5/5.1/5.2 DSA + LoRA enablement (DSA option: megatron-bridge-native and glm-native)#15yushengsu-thu wants to merge 23 commits into
yushengsu-thu wants to merge 23 commits into
Conversation
Older megatron-core (e.g. the radixark/miles image's 0.16.0rc0) only wires "gated_delta_net" in get_experimental_attention_variant_module_spec and raises ValueError for "dsa", and its get_dsa_module_spec_for_backend omits the metainfo the variant layer-builder reads. The GLM-5/5.1 bridge sets experimental_attention_variant="dsa" + transformer_layer_spec to that builder, so the model fails to build on such a core (LoRA and full-FT bridge paths alike). Wrap transformer_layer_spec in _build_glm5_dsa_block_spec: PREFER megatron-core's native handling, and only when it raises for "dsa" back-fill via the shipped DSA builder + set metainfo["fuse_input_layernorm"]=False (MLA-based DSA keeps a separate, non-fused input layernorm, like the deepseek_v4 dsv4 spec). On newer megatron-core (which handles "dsa" natively + sets metainfo) this is a transparent no-op, so the helper self-disables and can be deleted once the runtime core is bumped. Same spirit as the other miles-compat backports (mimo.config.role, training.config, parse_hybrid_pattern). Verified e2e: GLM-5.1 6-layer GRPO LoRA via bridge with no caller-side patch -> Job succeeded + PEFT adapter saved. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…dim) GLM-5.2 sets rope_theta=8e6, and under transformers>=5.12 the parsed GlmMoeDsaConfig reports qk_rope_head_dim as head_dim (192) instead of the config.json value (64). The base config-mapping then sized MLA linear_kv_down_proj as kv_lora_rank + 192 = 704, contradicting the checkpoint (kv_a_proj_with_mqa = kv_lora_rank + qk_rope_head_dim = 576 = 512 + 64). - rotary_base: read rope_theta whether nested in rope_parameters or flat. - qk_pos_emb_head_dim: re-read qk_rope_head_dim straight from config.json so MLA rope/kv dims match the weights. No-op when the parse is already correct; GLM-5.1 is unaffected (its head_dim already equals qk_rope_head_dim = 64). Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
GLM-5.2 keeps GLM-5.1's glm_moe_dsa arch but only "computing"/anchor layers carry the lightning indexer and compute the sparse top-k; "skip" layers reuse the most recent computing layer's top-k (HF config index_topk_freq>1 + index_skip_topk_offset). megatron-core's DSA is per-layer only, so this adds a Bridge-owned CrossLayerDSAttention(DSAttention): anchors publish topk_indices to a per-microbatch holder (packed_seq_params for thd, thread-local for bshd), skip layers drop their indexer (matching the subset checkpoint) and reuse the source anchor's top-k. get_glm5_crosslayer_dsa_spec calls megatron-core's exact get_dsa_module_spec_for_backend and only swaps core_attention.module. Feature-gated in _build_glm5_dsa_block_spec on dsa_index_topk_freq>1, so GLM-5.1 (no freq -> 1) keeps the existing per-layer path unchanged. No megatron-core edits. Validated: GLM-5.2 7-layer train-only e2e (build + 98G subset-ckpt load + cross-layer fwd/bwd + LoRA adapter) and GLM-5.1 6-layer full e2e regression both reach TRAIN EXIT 0. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
… layer (build time)
The per-microbatch top-k holder used for DSA cross-layer index sharing does NOT cross
pipeline boundaries, so a skip layer's source computing layer must live in the same PP
stage. Previously a bad (virtual) pipeline split that started a stage on a skip layer was
only caught by the runtime guard in CrossLayerDSAttention.forward (first forward of that
layer). This adds a build-time check mirroring slime's get_glm5_spec: it fails at model
construction with a precise message ("stage starts at global layer_number=X which is a skip
layer whose source computing layer=Y is on a previous stage").
- cross_layer_dsa.py: new assert_pp_stage_starts_on_computing_layer(config, vp_stage); uses
get_transformer_layer_offset + is_skip_topk_layer. No-op unless dsa_index_topk_freq>1, and
silently returns if the layout can't be determined (runtime guard remains the backstop).
- glm5_bridge.py: _build_glm5_dsa_block_spec calls it (gated on dsa + freq>1) before building
the block, forwarding vp_stage.
GLM-5.1 (freq=1) and valid PP=1 layouts are unaffected (no-op). Verified: raises on a stage
starting at a skip layer, no-op on a computing-layer start / GLM-5.1, and the real PP=1
GLM-5.2 build still passes.
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…older not recompute-safe) The per-microbatch top-k holder used for DSA cross-layer index sharing rides on packed_seq_params in the thd layout (closure-captured by the activation-checkpoint custom_forward, so it survives recompute). In the bshd layout packed_seq_params is None and the holder falls back to a process thread-local dict, which is NOT recompute-safe: under activation recompute a skip layer's recompute can read a stale anchor top-k (the dict is not captured per microbatch), silently corrupting gradients. CrossLayerDSAttention now records whether activation recompute is configured (self._recompute_active, from config.recompute_granularity) and, on a cross-layer forward with packed_seq_params is None, raises a clear AssertionError directing the user to --qkv-format thd (recompute-safe) or to disable activation recompute. No-op for thd, for no-recompute, and for GLM-5.1 (index_topk_freq=1). Training currently uses thd, so this only guards the unsafe bshd + recompute + cross-layer combination. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…gatron) megatron-core's DSAIndexer applies RoPE to the LAST qk_pos_emb_head_dim of each index head (split([D-rope, rope])), but the HF/DeepSeek (glm_moe_dsa) checkpoint stores the rope dims in the FIRST half. The previous name->name mapping loaded wq_b/wk/k_norm unchanged, so the indexer rotated the wrong dimensions -- its index scores were structurally decorrelated from the reference (HF<->bridge pearson ~0.48 vs slime ~0.70), flipping ~50% of the sparse top-k key selection at long context. Fix: _IndexerRopeHalfSwapMapping swaps the two halves of each index head's dsa_indexer_head_dim when loading the indexer wq_b, wk, k_norm.weight and k_norm.bias (self-inverse on export); weights_proj (per-head scalar) is untouched. k_norm is included because it is applied to the (swapped) key BEFORE RoPE, so its per-dim scale/bias must be swapped consistently. Mirrors the slime mbridge reference (THUDM/slime#2093 slime_plugins/mbridge/deepseek_v32.py: "training uses last half for rope while DeepSeek uses first half"). Covers GLM-5.1 (index_topk_freq=1) and GLM-5.2 (cross-layer). Validated against slime: all 5 indexer weights become bit-identical, indexer score pearson(slime,bridge) -> 0.98-1.0, and a 3-way HF-reference check shows the fixed bridge tracks HF identically to slime. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…lighting_indexer) Vendor slime's fused DSA kernels into glm_moe_dsa/fused/ for the selectable 'slime' attention backend: SparseMLA (sparse-MLA attention) + lighting_indexer (DSA indexer), each with fwd/bwd TileLang kernels. Self-contained (the miles-only indexer replay hook removed); imported lazily so the default backend stays dependency-free. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…-capable New SlimeMLASelfAttention(MLASelfAttention): default backend delegates to super() (bit-identical); the 'slime' backend runs slime's exact projection/rope/absorb/indexer numerics on the vendored fused kernels (attention bit-matches slime). LoRA-capable: folds the kv_b_proj LoRA delta into the absorb weight. Import-guard raises a clear error if the optional fused deps (tilelang/fast_hadamard_transform) are absent. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…tch + spec wiring
glm5_bridge.provider_bridge adds dsa_attention_backend ({megatron-bridge,slime}, default megatron-bridge); CrossLayerDSAttention.forward routes the sparse-attention call through a dispatch (_sparse_attention/_compute_layer_forward) with the default path unchanged (Milestone A: bit-identical); both GLM-5.1 and 5.2 specs point the MLA self-attention module at SlimeMLASelfAttention.
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
The fused slime backend's indexer top-k bypassed rollout-routing replay: the vendored lighting_indexer recomputed the top-k, and the megatron-core DSAIndexer only self-registers the indexer-replay stream in DeepSeek-V4 mode (dsv4_mode), not for GLM. So the train-side DSA selection could diverge from the rollout's, the gap R3 (arxiv 2510.11370) exists to close. - fused/indexer.py: route lighting_indexer's top-k through indexer_replay_manager.get_topk_fn when present+enabled (mirrors slime's indexer.py); guarded so the vendored kernel still runs standalone. - slime_mla.py: SlimeMLASelfAttention.__init__ registers a per-layer indexer_replay stream (mirrors slime glm5.py; the fork does not for GLM), gated on the slime backend; _slime_topk selects that stream before the kernel. No-op unless --use-indexer-replay enabled the manager before build; default backend untouched. Verified e2e (record->mangle stream->replay) on GLM-5.1-6layer and GLM-5.2-7layer (incl. cross-layer skip-layer holder reuse): replay overrides the recompute and the kernel consumes exactly the replayed selection. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
The _IndexerRopeHalfSwapMapping rewrites the DSA indexer wq_b/wk/k_norm assuming RoPE occupies the first qk_pos_emb_head_dim of every dsa_indexer_head_dim-sized head. That assumption only holds for the GlmMoeDsa HF layout; if a future GLM-5.x variant changes the indexer layout the previous code silently mis-permuted weights. Add two conversion-time guards that fail loudly on a changed layout: - _swap(): require the param leading dim to tile head_dim exactly (this is the precondition the following reshape already needs, just with a clearer message). - _dims(): require 0 < qk_pos_emb_head_dim < dsa_indexer_head_dim. _dims() reads the dims from megatron_module.config, which the conversion driver back-fills on the owning rank (model_bridge.py ~1596). On the EXPORT path under pipeline_model_parallel_size > 1, params a PP rank does not own get a broadcast-only fill task with megatron_module=None (model_bridge.py ~1626); megatron_to_hf() still runs the swap there on the PP-broadcast tensor, so when config is unreadable we fall back to the GlmMoeDsa indexer split (128/64) instead of raising -- this keeps PP>1 export working (matching pre-guard behavior) while the owning rank still validates the real dims. No behavior change for GLM-5/5.1/5.2 (dims are set and 64 < 128); other GLM families never reach this mapping. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…out/inference claims glm_moe_dsa/__init__.py had no module docstring. Add one documenting the two DSA sparse-MLA kernel backends selectable via config.dsa_attention_backend (set from the miles --dsa-attention-backend arg under bridge mode): "megatron-bridge" (default, unfused megatron-core kernels) and "slime" (vendored fused TileLang SparseMLA + lighting_indexer). Both support GLM-5.1 and GLM-5.2 (cross-layer index sharing), full or LoRA. Correct three inaccurate claims found while writing the matrix: - The only optional dep is tilelang; fast_hadamard_transform was never imported (the slime indexer path does no Hadamard rotation). Dropped it from the __init__/fused docstrings and the slime_mla ImportError message. - The default (unfused) backend does not "require bshd": it supports both bshd and thd (thd is the recompute-safe preferred carrier; bshd + activation recompute is rejected at forward). - Excluding the indexer from LoRA is a no-op only on slime (the fused indexer gets no gradient); the unfused backend gives it a tiny aux-loss gradient, so excluding it there is a deliberate choice, not a consequence. Docs only; no behavior change. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…er sequence-parallel
The fused thd indexer / SparseMLA path crashed under the canonical TP + sequence-parallel
layout ("CuSeqLenKS shape[0] expected 512, got 128") because _slime_index_qkw / _slime_topk
omitted the SP/CP token-dim reconciliation native slime (miles_plugins glm5.py) performs:
SP all-gather of index_q / index_k / head_weights, CP all-gather of index_k, and CP-scatter of
the per-query cu_seqlens (starts/ends). Add all four, guarded by config.sequence_parallel; the
CP ops are no-ops at CP=1. Mirrors glm5.py:218-219,538-555. The SparseMLA query is already
SP-full via the column-parallel linear_q_up_proj, so it matches the now-SP-full top-k.
Validated: GLM-5.2_5layer LoRA e2e at TP=4 + sequence-parallel + CP=1 runs 10 steps clean
(rollout -> fused fwd -> SparseMLA bwd -> optimizer -> save). CP>1 remains blocked upstream by
megatron-core (transformer_config asserts context_parallel_size==1 for the 'dsa' attention
variant); the CP branches here are no-op-at-CP=1 forward-compat mirroring native slime.
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
… sequences _original_topk called torch.topk(logits, index_topk=2048), which raises 'selected index k out of range' when the packed sequence is shorter than index_topk (e.g. gsm8k toy rollouts ~256 tokens, where the indexer degenerates to dense). Cap k at logits.shape[-1] and pad the result back to the fixed index_topk width with -1, matching the rollout-captured top-k shape and what SparseMLA expects. The long-sequence path (k == index_topk) is unchanged. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…nd split
Rename models/glm_moe_dsa/ to models/glm5/ and reorganize the DSA sparse-MLA
kernels by backend so the two paths are structurally explicit:
glm5/
cross_layer_dsa_dispatch.py # CrossLayerDSAttention dispatcher (was cross_layer_dsa.py)
fused/ # slime fused TileLang kernels (slime_mla.py moved here)
unfused/ # megatron-core DSA entry point (re-exports DSAttention etc.)
The dispatch selects the kernel via config.dsa_attention_backend
("slime" fused vs "megatron-bridge" unfused, default). Both backends support
GLM-5.1 and GLM-5.2 (DSA cross-layer index sharing), full or LoRA.
Pure restructure: only module paths and the new unfused/__init__.py entry
point change. HF source id (model_type="glm_moe_dsa", GlmMoeDsaForCausalLM)
and all kernel logic are unchanged.
Verified e2e (colocate LoRA RL, gsm8k, rollout->train->save) on both backends:
GLM-5.2_5layer fused+unfused, GLM-5.1-6layer fused+unfused.
…s gradient SlimeMLASelfAttention._kv_up_proj_weight_and_norm folded the LoRA delta into the absorbed kv_up weight using a non-differentiable torch.distributed.all_gather on linear_in (LoRA-A), detaching A from the autograd graph -> LoRA-A never received gradients on the fused (slime) backend (LoRA-B trained, A frozen). Use torch.distributed.nn.functional.all_gather (autograd-aware; reduce-scatter backward, correct since the gathered A is used in every TP rank's local delta). Verified by a TP micro-test: linear_in.grad None -> nonzero. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…moe_dsa import shim
- glm5_bridge.py: the qk_rope_head_dim mis-parse override (transformers>=5.12 reports
the derived head_dim 192 instead of the config's 64, mis-sizing linear_kv_down_proj)
only fired when _name_or_path was a local directory with config.json on disk. For a
repo-id load (AutoBridge.from_hf_pretrained("zai-org/GLM-5.2")) the path check failed
and the override was skipped, leaving the wrong dim -> weight-load shape mismatch.
Read the RAW qk_rope_head_dim from config.json resolved for BOTH load styles (local
dir, else HF-cache via cached_file). NOTE: must read raw config.json, NOT
hf_config.qk_rope_head_dim -- transformers overwrites that in-memory attribute with
the mis-parsed 192, so trusting it would skip the correction (GPU-verified: doing so
crashed the GLM-5.2_5layer rope-half swap with qk_pos_emb_head_dim=192 > 128).
- glm5_bridge.py: name the config-less PP-broadcast rope-half-swap fallback dims as a
single-source constant and warn when that export path is taken (a future non-128/64
indexer layout would otherwise diverge from the owning rank silently).
- add megatron/bridge/models/glm_moe_dsa.py deprecation shim re-exporting GLM5Bridge so
the pre-restructure import path keeps working (with a DeprecationWarning).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
yushengsu-thu
force-pushed
the
bridge-dev-glm-merging
branch
from
July 2, 2026 01:38
b64c312 to
d532606
Compare
…ntifiers dsa_attention_backend value rename (pairs with the miles-side flag rename): megatron-bridge -> megatron-bridge-native (portable unfused megatron-core kernels) slime -> glm-native (fused TileLang kernels vendored from slime) 'megatron-bridge' as a value collided with Megatron-Bridge the package (the option is only reachable when Megatron-Bridge IS the training backend), and 'slime' named the kernels' upstream origin rather than the implementation. Identifiers follow the backend name: fused/slime_mla.py -> fused/glm_native_mla.py SlimeMLASelfAttention -> GlmNativeMLASelfAttention _slime_forward/_slime_topk/ _slime_index_qkw/... -> _glm_native_* Provenance references to the slime project (vendored kernel origin, reference implementations, THUDM/slime pointers) are intentionally kept. No kernel or logic change; the layout assert message and dispatch/docstrings now name the new values. Validated on GLM-5.2_5layer via miles: TOY RL smoke on megatron-bridge-native (abs_diff 0.0105 / kl 1.17e-4) and an attention-only toy run on glm-native (thd) through rollout->train step 1, both clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HFoyYjC55SRsNRXUzWMWKH Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
yushengsu-thu
changed the title
…values glm5/fused/ -> glm5/glm_native/ and glm5/unfused/ -> glm5/megatron_bridge_native/ so the package layout matches the backend values (megatron-bridge-native / glm-native; Python package names use underscores). Update import paths and the stale docstring references the value-rename left behind (``slime`` / ``megatron-bridge`` as backend names, the old cross_layer_dsa.py file name, glm_moe_dsa/ paths in the vendored-kernel headers). Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
This was referenced Jul 2, 2026
Collaborator
Author
|
Split into two stacked PRs for review: #16 (1/2 — DSA option: |
Per miles PR review (clearer naming): dsa_attention_backend values
megatron-bridge-native -> megatron and glm-native -> tilelang, with the
package layout, module and identifiers renamed to match:
models/glm5/megatron_bridge_native/ -> models/glm5/megatron/
models/glm5/glm_native/ -> models/glm5/tilelang/
glm_native_mla.py -> tilelang_mla.py
GlmNativeMLASelfAttention -> TileLangMLASelfAttention
_glm_native_{forward,index_qkw,topk} -> _tilelang_*
Validated e2e: GLM-5.2_5layer colocate LoRA RL --ci-test on the default
tilelang backend (weight check clean, step-1 train_rollout_kl ~1e-4).
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
Collaborator
Author
|
Naming update (73a0d3f, per miles PR review): |
…r field GLM5ModelProvider becomes a real MLAModelProvider subclass (was an import alias) declaring dsa_attention_backend: str = "megatron" as a dataclass field, and the bridge registers/returns it. Replaces the ad-hoc attribute set in provider_bridge; since provide() hands the provider object itself to the model as its config, the field reaches every module's config with no caller-side propagation (the miles-side post-build force loop is now deleted). Validated e2e on GLM-5.2_5layer LoRA RL --ci-test (tilelang default; weight check clean, step-1 train_rollout_kl ~1e-4). Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…ped-expert LoRA) Per miles PR review: per-expert becomes the default grouped-expert LoRA layout; the shared-outer layout is opted into solely via the experts_shared_outer_loras field (which takes priority in the transform dispatch and does not depend on share_expert_adapters), so callers select the layout through a single, consistently named knob. The legacy implicit shared layout is no longer a silent default; adapters saved under it are not shape-compatible with per-expert ones. Validated e2e on GLM-5.2_5layer LoRA RL --ci-test in BOTH layouts (shared-outer via the single flag, per-expert via the new default; weight check clean, step-1 train_rollout_kl ~1e-4 each). Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
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
GLM-5 / GLM-5.1 / GLM-5.2 (
glm_moe_dsa) DSA + LoRA enablement for the bridge path. Restructuresmodels/glm_moe_dsa/→models/glm5/with a per-backend package split (glm5/megatron/unfused megatron-core kernels,glm5/tilelang/fused TileLang kernels), fixes two HF↔Megatron conversion bugs, adds GLM-5.2 cross-layer DSA index sharing, and makes GLM-5.x LoRA-trainable end-to-end.For review, this branch is split into two stacked PRs (together they reproduce this tree byte-for-byte): #16 (
megatronbackend + restructure + conversion fixes + cross-layer sharing) and #17 (tilelangfused backend). This branch stays as the integration reference.Key points
GLM5ModelProvider.dsa_attention_backenddataclass field, values"megatron"(default; portable unfused megatron-core kernels, bshd) /"tilelang"(fused TileLang kernels, thd, optionaltilelangdep, training/forward-only — rollout is always served by sglang). Set from the miles--dsa-attention-backendarg; since the provider object is the model's config, the field reaches every module with no caller-side propagation._IndexerRopeHalfSwapMapping; pre-fix HF↔bridge indexer-score pearson ~0.48 → 0.98–1.0), and GLM-5.2 MLA rope dims re-read fromconfig.json(transformers≥5.12 mis-parse).CrossLayerDSAttention): anchor layers publish top-k, skip layers drop their indexer and reuse it; build-time PP-stage assert + forward-time bshd+recompute rejection. Old-core feature detection back-fills the"dsa"spec on megatron-core versions without it.share_expert_adaptersdefaults toFalse(per-expert); the shared-outer layout is opted into solely via theexperts_shared_outer_lorasfield (same name end-to-end from the miles CLI). The two layouts' checkpoints are not interchangeable.Validation (8×H200)
miles PR radixark/miles#1559 CI combination matrix — {shared-outer + virtual-experts, per-expert + no-virtual-experts} × {tilelang, megatron} — is 8/8 green on GLM-5.2_5layer and GLM-5.1-6layer against this head; full 744B GLM-5.2 runs 50+ steps at train↔rollout KL ~1e-4 with healthy reward. R3 indexer replay, sequence-parallel, and old-core (0.16.0rc0) paths validated e2e.
Naming history
dsa_attention_backendvalues wentmegatron-bridge/slime→megatron-bridge-native/glm-native→megatron/tilelang(current; dirsglm5/megatron/,glm5/tilelang/, classTileLangMLASelfAttention). Earlier review comments may use the old names.🤖 Generated with Claude Code