把 GLM-5.3-Flash 的训练上下文拉到 100 万 token：SkyRL 相对 Megatron 基线提升 64 倍

在 B200（约 178 GiB）上，SkyRL 通过四类系统优化，把 GLM-5.3-Flash 的单序列训练长度做到 100 万 token，峰值约 144.5 GiB，约为开箱 Megatron 基线的 64 倍。关键不只是堆并行，还包括：让稀疏注意力真的跑在稀疏核上、切碎并分片二次复杂度的 indexer、激活重算与 CPU offload、以及去掉不必要的 fp32 升精度。
https://blog.skyrl.ai/blog/1m-context-glm-5-3-flash

模型与瓶颈
· GLM-5.3-Flash 是混合注意力：45 层里 34 层用 Kimi Delta Attention（线性注意力），11 层用 DeepSeek Sparse Attention（DSA）。只有 DSA 的代价随序列长度近似平方增长，决定了上下文上限。
· DSA 本身对每个 query 只读约 2048 个关键 key（注意力近似线性），但 indexer 要对每个 query 给全部「每 4 token 一池」打分，稀疏注意力的二次成本被挪到了 indexer。

四块优化
1. DSA 真稀疏
· 问题：Megatron 融合稀疏核只认 DeepSeek-V3.2 的张量形状；Flash 版 DSA 无位置编码（NoPE）、维度更窄，且选 2051 而非 2048 个 key。内核静默拒掉后回退到全量稠密注意力。其他优化都开着时，16k 就要约 114 GiB/GPU，32k 直接 OOM。
· 做法：把张量 pad 到核期望形状；选用 TileLang（相对 cuDNN/FlashMLA 在 128k 约省 3 倍显存）。16k 降到约 56 GiB。
· Indexer 按 query 分块打分：32k 时全量打分峰值约 64 GiB，分块后约 2.3 GiB。
· Indexer 按 TP 分片：原先 8 卡各自重复整段 O(S²) 打分；改为每卡只算自己那份 query，再交换较小的选中 pool 索引（int32）。256k / 64×B200 上，整步从约 218 s 到约 102 s（约 2.1×），是 DSA 相关里最大的加速。

2. 激活重算与 Offload
· 模型用 mHC（4 路并行残差），激活约为普通 hidden 的 4 倍；选择性重算大约只能到 32k。改成整层 full recompute 可到约 256k（相对无重算约 +6% 步时）。
· Full recompute 暴露 MoE router 概率泄漏：修掉后 256k 约省 41 GiB/GPU。
· 层输入 offload 到 CPU：256k 峰值约 142→98 GiB，步时约 +21%；512k 以上几乎必需。用 pinned memory 可比 pageable 再快约 22%。代价是主机内存：1M 时仅 offload 输入、在 CP 前约 180 GB/GPU，8 卡节点超过 1 TB。

3. 上下文并行（CP）
· KDA：Megatron 原先无 CP。按头拆分（Ulysses 风格），每卡对全序列跑自己那部分头，免传循环状态。并把通信从「投影后的 Q/K/V/gate」改成先交换更小的 hidden，再本地投影；CP2 时每层入站字节约减 5×，1M 整步约 335→307 s（−8%）。
· DSA：修 zigzag 分片下 k-pool gate 分数与 key 对齐的问题。
· 仅在单序列装不下时才开 CP：512k 无 CP 聚合吞吐约 18.2k tok/s，1M + CP2 约 13.7k。

4. 去掉多余 fp32
· mHC 投影：不存 fp32 输入副本，反向时重算 cast；512k 约省 8 GiB。
· MoE combine：在 2 个 expert-TP rank 时用 bf16 reduce-scatter（NCCL 内部仍用 fp32 做一次加），省一半通信与拷贝；更多 rank / NVLS 则仍用 fp32。

真实数据
· 随机 token 测得的显存会低估：真实文本 MoE 路由不均，最忙卡可达约 4× 平均负载；ProLong-512K 书籍微调时，512k 在随机设定能放下、真实数据会 OOM。对 MoE 按序列分块计算可修复。
· 真实书：512k 峰值约 122 GiB；1M（CP=4）约 109 GiB。
· 用人类基因组 hg38 微调演示学习：chr1–2 训练、chr22 留出；512k 上 80 步留出 loss 降约 4.8%。按位置分桶显示后期 token 改善更大，作者据此认为存在长上下文学习。
· 正确性：真实权重下 trainer 与 vLLM 的 logprob 平均绝对差约 0.04–0.07（128k–1M）；四本 512k 书反复微调 20 步，loss 从约 1.65 到 0.17。