GLM-5.2 131k context parallelism: dsv4-next stack + GLM fused-DSA contiguous-CP adapter

[JackRao123]

Purpose

One diff against trainers-main showing every Megatron-LM change needed for GLM-5.2 131k CP. Retargeted from trainers-main-next-dsv4 per Jack — the diff therefore includes that whole line (not just my commits), because 131k needs it:

• Inherited via trainers-main-next-dsv4 (the newer NVIDIA-dev base): upstream PR [dev] [DeepSeek-v4] Context Parallel support NVIDIA/Megatron-LM#5087 (DSv4 context-parallel machinery — csa_cp_utils, q_causal_offsets through the DSA indexer kernels, cp_partition_mode='contiguous'), THD packed sequences ([dev] [DeepSeek-v4] Packed Sequence (THD) support for DSv4 Hybrid Attention NVIDIA/Megatron-LM#5011), the GLM-5.2 fused frozen-indexer commit re-layered, RoPE int64/split-recompute, mHC leak fixes, 131k memory hardening. These are jerry's/upstream's — listed here for visibility, authorship in the commits.
• Mine (jackrao/glm-dsa-cp, the last 4 commits):
  • glm_dsa_fused.py _forward_thd_cp: autograd CP all-gather of the compressed latent KV (bwd reduce-scatters dKV), non-autograd gather of the frozen indexer's K rows, per-rank top-k against global KV via compute_cp_indexer_topk (ratio=1) + THD flat-index lowering; FlashMLA/cuDNN sparse attention unchanged; IndexShare holders store (topk, layout).
  • rope_utils.py: contiguous branch in _apply_rotary_pos_emb_thd (positions from GLOBAL cu_seqlens over each rank's contiguous packed-row interval); zigzag untouched; absorbed_mla max_seqlen plumb.
  • transformer_config.py: allow dsa alongside dsv4_hybrid for contiguous CP validation.
  • CP1 fix needed even without CP: the re-layered GLM commit still called the pre-[dev] [DeepSeek-v4] Context Parallel support NVIDIA/Megatron-LM#5087 build_flat_topk_idxs(seqlen_kv=...) signature — GLM died at warmup on this line before this fix. Plus indexer rope-table sizing by padded length (OOB gather device-assert otherwise).

Validation

Debug GLM: CP1-vs-CP2 loss parity 0.03%, per-rank peak shards 2.33 → 1.70 GiB. Full GLM-5.2-FP8, 4×8 B200, TP1/PP1/EP32/CP32: 131 072-token LoRA-SFT fwd-bwd at 98.2 GiB peak-alloc / 138.2 GiB hottest used, ~56 s/fb, memory fit R²=0.999 over 6 points.

Caveats: fb/memory validated; real-data parity + convergence pending; CE-only v1; ~1e-3 forward jitter at CP2 (nondeterministic kernel in the gather/FlashMLA path).

Companions: basetenlabs/trainers#592, basetenlabs/Megatron-Bridge#12.

🤖 Generated with Claude Code

[JackRao123]

Closing — superseded by the upstream fused DSA+CP in trainers-main-20260907 (NVIDIA NVIDIA#5099 / NVIDIA#5246). Of this PR's 3 commits on top of the dev sync: (1) contiguous-CP GLM frozen-indexer is replaced by upstream's zigzag fused DSA CP (the trainer is being pivoted to that); (2) the mHC backward-leak fixes touch hyper_connection.py — DSv4-hybrid only, not GLM-5.2; (3) the 131k memory hardening is mostly superseded (its MoE dispatcher-clear is already on the new base as 922ebcb; RoPE-int64 / split-recompute / int32-stride already upstream). The only pieces not clearly upstream are two pure-memory optimizations (chunked Triton indexer scorer for the SBHD non-CP path, GLU linear-offset skip) — parked as a separate memory/perf follow-up to re-port onto upstream's reworked kernels IF CP32 131k OOMs on the new stack. Branch left intact.