SkyRL 将 GLM-5.3-Flash 训练上下文提升至 100 万 token
SkyRL 通过四类系统优化，在 B200（约 178 GiB）上把 GLM-5.3-Flash 的单序列训练长度做到 100 万 token，峰值约 144.5 GiB，约为开箱 Megatron 基线的 64 倍。
meng shao· @shao__meng · X·2026-10-23 10:24· 1 小时前AI 评分47
AI 导读
SkyRL 通过四类系统优化，在 B200（约 178 GiB）上把 GLM-5.3-Flash 的单序列训练长度做到 100 万 token，峰值约 144.5 GiB，约为开箱 Megatron 基线的 64 倍。
正文 · 原文
把 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。
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://blog.skyrl.ai/blog/1m-context-glm-5-3-flash! 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!在 X 查看被引用的帖子
来源：meng shao · x.com
来源：RadarAI · aihot.news