SkyRL extends GLM-5.3-Flash training context to 1 million tokens
Through four types of system optimizations, SkyRL achieves a single-sequence training length of 1 million tokens for GLM-5.3-Flash on B200, with a peak of about 144.5 GiB, roughly 64 times that of the out-of-the-box Megatron baseline. The optimizations include running DSA on sparse kernels, chunking by query and sharding the indexer by TP, recomputing full-layer activations and CPU offloading layer inputs, and removing redundant fp32 upcasting. On real book data, the 512k peak is about 122 GiB, and 1M (CP=4) is about 109 GiB; the mean absolute difference in logprob between the trainer and vLLM is about 0.04–0.07.
Why read this
AI score 58
The original text breaks down the four types of system optimizations and the memory and step-time data item by item, showing exactly where the long-context training bottleneck is stuck.
Topics
View original postSuper 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!
View original
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!
Body · AI translation
Pulling GLM-5.3-Flash's training context to 1 million tokens: SkyRL achieves a 64x improvement over the Megatron baseline
On B200 (about 178 GiB), SkyRL uses four types of system optimizations to push GLM-5.3-Flash's single-sequence training length to 1 million tokens, with a peak of about 144.5 GiB, roughly 64x the out-of-the-box Megatron baseline. The key is not just stacking parallelism, but also: making sparse attention actually run on sparse kernels, chunking and sharding the quadratic-complexity indexer, activation recomputation and CPU offload, and removing unnecessary fp32 upcasting. https://t.co/Z3E8qjKsZ8
Model and bottleneck · GLM-5.3-Flash is hybrid attention: 34 of its 45 layers use Kimi Delta Attention (linear attention), and 11 layers use DeepSeek Sparse Attention (DSA). Only DSA's cost grows approximately quadratically with sequence length, which determines the context limit. · DSA itself reads only about 2048 key keys per query (attention is approximately linear), but the indexer must score all "one pool per 4 tokens" for every query, so the quadratic cost of sparse attention is shifted onto the indexer.
Four optimizations 1. DSA truly sparse · Problem: Megatron's fused sparse kernel only recognizes DeepSeek-V3.2's tensor shapes; the Flash version of DSA has no positional encoding (NoPE), narrower dimensions, and selects 2051 rather than 2048 keys. After the kernel silently rejects it, it falls back to full dense attention. With the other optimizations enabled, 16k already requires about 114 GiB/GPU, and 32k goes straight to OOM. · Approach: pad tensors to the kernel's expected shape; choose TileLang (saves about 3x memory versus cuDNN/FlashMLA at 128k). 16k drops to about 56 GiB. · Indexer scores in chunks per query: at 32k, full scoring peaks at about 64 GiB, and after chunking about 2.3 GiB. · Indexer sharded by TP: previously 8 GPUs each redundantly scored the entire O(S²) segment; instead, each GPU computes only its own share of queries, then exchanges the smaller selected pool indices (int32). On 256k / 64×B200, the full step goes from about 218 s to about 102 s (about 2.1x), the largest speedup among DSA-related ones.
2. Activation recomputation and offload · The model uses mHC (4-way parallel residual), so activations are about 4x a normal hidden; selective recomputation only reaches about 32k. Switching to full recompute of entire layers reaches about 256k (about +6% step time versus no recomputation). · Full recompute exposed MoE router probability leakage: after fixing it, 256k saves about 41 GiB/GPU. · Offloading layer inputs to CPU: 256k peak goes from about 142 to 98 GiB, with step time about +21%; above 512k it is almost mandatory. Using pinned memory is about 22% faster than pageable. The cost is host memory: at 1M, offloading only inputs, before CP, is about 180 GB/GPU, over 1 TB for an 8-GPU node.
3. Context parallelism (CP) · KDA: Megatron previously had no CP. Split by head (Ulysses style), each GPU runs its own share of heads over the full sequence, avoiding transfer of recurrent state. Also change communication from "projected Q/K/V/gate" to first exchanging the smaller hidden, then projecting locally; at CP2, inbound bytes per layer drop about 5x, and the 1M full step goes from about 335 to 307 s (−8%). · DSA: fix the misalignment between k-pool gate scores and keys under zigzag sharding. · Enable CP only when a single sequence does not fit: at 512k without CP, aggregate throughput is about 18.2k tok/s; at 1M + CP2, about 13.7k.
4. Removing redundant fp32 · mHC projection: do not store an fp32 copy of the input; recompute the cast during backward; saves about 8 GiB at 512k. · MoE combine: with 2 expert-TP ranks, use bf16 reduce-scatter (NCCL internally still does one addition in fp32), saving half the communication and copies; with more ranks / NVLS, still use fp32.
Real data · Memory measured with random tokens underestimates: real text has uneven MoE routing, and the busiest GPU can reach about 4x the average load; when fine-tuning on ProLong-512K books, 512k fits under random settings but OOMs on real data. Computing MoE in chunks by sequence fixes this. · Real books: 512k peaks at about 122 GiB; 1M (CP=4) at about 109 GiB. · Fine-tuning on the human genome hg38 to demonstrate learning: train on chr1–2, hold out chr22; at 512k, held-out loss drops about 4.8% over 80 steps. Bucketing by position shows later tokens improve more, and the authors take this as evidence of long-context learning. · Correctness: with real weights, the mean absolute difference in logprob between the trainer and vLLM is about 0.04–0.07 (128k–1M); repeatedly fine-tuning four 512k books for 20 steps moves loss from about 1.65 to 0.17.