SkyRL 将 GLM-5.3-Flash 训练上下文扩至 100 万 token
把 GLM-5.3-Flash 的训练上下文拉到 100 万 token：SkyRL 相对 Megatron 基线提升 64 倍
SkyRL 通过四类系统优化，在 B200 上把 GLM-5.3-Flash 的单序列训练长度做到 100 万 token，峰值约 144.5 GiB，约为开箱 Megatron 基线的 64 倍。优化包括让 DSA 跑在稀疏核上、按 query 分块并按 TP 分片 indexer、整层激活重算与层输入 CPU offload、去掉多余 fp32 升精度。真实书籍数据下 512k 峰值约 122 GiB，1M（CP=4）约 109 GiB，trainer 与 vLLM 的 logprob 平均绝对差约 0.04–0.07。
查看原帖非常激动地分享我在 SkyRL 上为 GLM 5.3 Flash 支持 100 万 token 上下文训练所做的一些工作，比开箱即用的 Megatron 基线高出 64 倍！🚀 这项工作是与 @LilaSciences 的优秀伙伴们合作完成的！🤩 实现这一目标需要在四个主要方面进行关键的算法和系统优化： 1) DeepSeek 稀疏注意力：在稀疏核上运行 DSA；对其二次索引器进行分块和分片。 2) 激活重算与卸载：在反向传播中重算激活，并将保存的输入卸载到 CPU。 3) 上下文并行：将每个序列拆分到多个 GPU 上，为 KDA 新增 CP 支持，并修复 DSA。 4) 避免不必要的精度升位：让 FP32 副本不占用内存、不经过网络。 有关我们实现和结果的更多细节，请查看我们的完整技术博客：https://t.co/Dna0CGNeJt！ 这项工作得益于许多人在过程中做出的贡献。衷心感谢来自 @anyscalecompute 的 @pcmoritz、@erictang000、@sumanthrh、@CyrusHakha、@richliaw，以及来自 Lila Sciences 的 Keith Truongcao、Tushar Khot、Vivek Subramanian、Semih Kara 的所有帮助与指导！
查看原文
Super excited to share some of the work I’ve been doing on supporting 1M token context training on SkyRL for GLM 5.3 Flash, 64x higher than the out of the box Megatron baseline! 🚀 This work was done in collaboration with the amazing folks at @LilaSciences! 🤩 Getting there took key algorithmic and systems optimizations in four main areas: 1) DeepSeek Sparse Attention: run DSA on a sparse kernel; chunk and shard its quadratic indexer. 2) Activation recompute and offloading: recompute activations in backward and offload saved inputs to CPU. 3) Context parallelism: split each sequence across GPUs, with new CP support for KDA and a fix for DSA. 4) Avoiding unnecessary precision upcasts: keep FP32 copies out of memory and off the network. For more details on our implementations and our results, check out our full technical blog here: https://t.co/Dna0CGNeJt! This work benefited from the contributions of many people along the way. Big thanks to @pcmoritz, @erictang000, @sumanthrh, @CyrusHakha, @richliaw from @anyscalecompute and Keith Truongcao, Tushar Khot, Vivek Subramanian, Semih Kara from Lila Sciences for all the help and guidance!
正文
把 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://t.co/Z3E8qjKsZ8
模型与瓶颈 · 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。