[model, peft] feat: GLM-5 DSA option: tilelang (2/2)

[yushengsu-thu]

Summary

Part 2/2 of splitting #15 (part 1/2 = #16: the megatron unfused backend + GLM-5.x enablement). Stacked on #16 — this branch contains part 1's commit, so the diff shows both until part 1 merges; review the [model, peft] feat: ... tilelang ... (2/2) commit only. The two parts together reproduce the bridge-dev-glm-merging model tree byte-for-byte (verified: git diff vs its head is empty for src/).

This part adds the second DSA sparse-MLA kernel backend: dsa_attention_backend="tilelang" — vendored fused TileLang kernels (SparseMLA + lighting_indexer fwd/bwd under models/glm5/tilelang/) driven by TileLangMLASelfAttention (tilelang/tilelang_mla.py), which replicates slime's absorbed-latent q/kv + apex fuse_rope numerics so the train-side attention output matches the sglang rollout, including R3 indexer replay. Requires the optional tilelang dependency (imported lazily — the default megatron path stays dependency-free) and the thd packed layout (--qkv-format thd). Training/forward-only: asserts inference_context is None (no KV cache); rollout is always served by sglang.

Bug fixes (fused-backend specific)

• Crashes under TP + sequence-parallel → "CuSeqLenKS shape[0] expected 512, got 128": _tilelang_index_qkw / _tilelang_topk omitted the SP/CP token-dim reconciliation native slime performs → added SP all-gather of index_q/index_k/head_weights, CP all-gather of index_k, and CP-scatter of the per-query starts/ends (CP ops are no-ops at CP=1).
• Fused indexer crashes on short sequences → torch.topk(logits, index_topk=2048) raises "selected index k out of range" when the packed sequence is shorter than index_topk → _original_topk (tilelang/indexer.py) caps k at logits.shape[-1] and pads the selection back to the fixed topk width with the -1 sentinel, preserving the shape SparseMLA and R3 replay expect.
• LoRA-A frozen at TP>1 → TileLangMLASelfAttention._kv_up_proj_weight_and_norm folded the kv_b_proj LoRA delta using a non-differentiable torch.distributed.all_gather on linear_in (LoRA-A), detaching it from the autograd graph, so LoRA-B trained while LoRA-A never received gradients → use torch.distributed.nn.functional.all_gather (autograd-aware; reduce-scatter backward, correct because the gathered A participates in every TP rank's local delta).

New features

• Fused backend dispatch: the fused path is dispatched in TileLangMLASelfAttention.forward, which intercepts at the MLA level (the fused kernel consumes the absorbed-latent q/kv, which only exist upstream of core attention); the block-spec builders (glm5_bridge._build_glm5_dsa_block_spec, cross_layer_dsa_dispatch.get_glm5_crosslayer_dsa_spec) point the MLA module at TileLangMLASelfAttention — on the default megatron backend its forward delegates to the base class → byte-identical. CrossLayerDSAttention also reads the flag so GLM-5.1/5.2 forwards share one dispatch point.
• R3 indexer replay on the fused backend: lighting_indexer routes its top-k through miles' indexer_replay_manager.get_topk_fn when present + enabled, and TileLangMLASelfAttention.__init__ registers the per-layer replay stream (megatron-core's DSAIndexer only self-registers in DeepSeek-V4 mode). No-op unless --use-indexer-replay enabled the manager; import-guarded so the package works without miles.
• LoRA-capable: the kv_b_proj adapter delta is folded into the absorb weight so its adapters genuinely train. LoRA on the DSA indexer is a genuine no-op here (the fused lighting_indexer returns only discrete top-k, no indexer loss → no gradient); the miles launcher excludes the indexer from LoRA targets by default.

Validation

• Restructure + both backends verified e2e (colocate LoRA RL, gsm8k, rollout→train→save) on GLM-5.2_5layer fused+unfused and GLM-5.1-6layer fused+unfused.
• Sequence-parallel fix: GLM-5.2_5layer LoRA e2e at TP=4 + sequence-parallel runs 10 steps clean (rollout → fused fwd → SparseMLA bwd → optimizer → save).
• Differentiable all-gather: with plain dist.all_gather the kv_up lora_A stays frozen at TP>1 (5-angle verification + CPU micro-test: linear_in.grad None → nonzero with the fix).
• R3 replay verified e2e (record → mangle stream → replay) on GLM-5.1-6layer and GLM-5.2-7layer, including skip-layer holder reuse: replay overrides the recompute and the kernel consumes exactly the replayed selection.
• Full GLM-5.2 744B-A40B LoRA RL (fused backend, 64×H200) runs multi-step with healthy on-policy numbers (train↔rollout KL ~1e-4) on the integration branch this splits from.

Notes for reviewers

• CP>1 is still blocked upstream (megatron-core asserts context_parallel_size==1 for the "dsa" variant); the CP branches here are no-op-at-CP=1 forward-compat mirroring native slime.
• CrossLayerDSAttention._sparse_attention currently invokes the unfused kernel on both branches (leftover Step-4 TODO; the fused SparseMLA is never called from there — the real fused dispatch is at the MLA level).
• Supersedes the model-tree half of GLM-5/5.1/5.2 DSA + LoRA enablement (DSA option: megatron-bridge-native and glm-native) #15 together with part 1/2; bridge-dev-glm-merging is kept (not deleted) as the integration branch.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HFoyYjC55SRsNRXUzWMWKH

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.