Out of the Memory Barrier: A Highly Memory-Efficient Training System for LLMs with Million-Token Contexts
Authors: Wenhao Li
Daohai Yu
Gen Luo
Yuxin Zhang
Yifan Wu
Jiaxin Liu
Ziyang Gong
Zimu Liao
Fei Chao
Rongrong Ji
Venue: ICLR 2026 Poster (2026)
Keywords: LLM
NLP
Long-Context LLM
Memory Efficient Training
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: We built a system called OOMB that uses chunk-based processing and smart CPU offloading to train LLMs on million-token contexts using a single GPU, a task that previously required a large computer cluster.

Abstract: Training Large Language Models (LLMs) on long contexts is severely constrained by prohibitive GPU memory overhead, not training time. The primary culprits are the activations, whose memory footprints scale linearly with sequence length. We introduce OOMB, a highly memory-efficient training system that directly confronts this barrier. Our approach employs a chunk-recurrent training framework with on-the-fly activation recomputation, which maintains a constant activation memory footprint ($\mathcal{O}(1)$) and shifts the primary bottleneck to the growing KV cache. To manage the KV cache, OOMB integrates a suite of synergistic optimizations: a paged memory manager for both the KV cache and its gradients to eliminate fragmentation, asynchronous CPU offloading to hide data transfer latency, and page-level sparse attention to reduce both computational complexity and communication overhead. The synergy of these techniques yields exceptional efficiency. Our empirical results show that for every additional 10K tokens of context, the end-to-end training memory overhead increases by a mere 10MB for Qwen2.5-7B. This allows training Qwen2.5-7B with a 4M-token context on a single H200 GPU, a feat that would otherwise require a large cluster using context parallelism. This work represents a substantial advance in resource efficiency for long-context LLM training. The source code is available for review at https://anonymous.4open.science/r/oomb/README.md.

Published as a conference paper at ICLR 2026
OUT OF THE MEMORY BARRIER: A HIGHLY MEM- ORY EFFICIENT TRAINING SYSTEM FOR LLMS WITH MILLION-TOKE CONTEXTS
Wenhao Li; Daohai Yu; Gen Luo, Yuxin Zhang, Fei Chao Rongrong Ji! Key Laboratory of Multimedia Trusted Perception and Efficient Computing Ministry of Education of China, Xiamen University
Yifan Wu Peking University
Jiaxin Liu University of Illinois Urbana-Champaign
Ziyang Gong
Zimu Liao
Shanghai Jiao Tong University
Shanghai AI Laboratory
ABSTRACT
Training Large Language Models (LLMs) on long contexts is severely constrained by prohibitive GPU memory overhead, not training time. The primary culprits are the activations, whose memory footprints scale linearly with sequence length. We introduce OOMB, a highly memory-efficiem training system that directly confronts this barrier. Our approach employs chunk-recurrent training framework with on-the-fly activation recomputation, which maintains constant activation memory footprint (O) and shifts the primary bottleneck the growing KV cache. To manage the KV cache, OOMB integrates a suite of synergistic optimizations: a paged memory manager for both the KV cache and its gradients to eliminate fragmentation, asynchronous CPU offloading to hide data transfer latency, and page- level sparse attention reduce both computational complexity and communication overhead. The synergy of these techniques yields exceptional efficiency. Our empirical results show that for every additional 10K tokens of context, the end-to- end training memory overhead increases by mere 10MB for Qwen2.5-7B. This allows training Qwen2.5-7B with a 4M-token context on a single H200 GPU, a feat that would otherwise require large cluster using context parallelism. This work represents a substantial advance in resource efficiency for long-context LLM training. The source code available https://gihhbb.comm/ewmaaaii///m
INTRODUCTION
Long-context modeling remains a formidable challenge in training Large Language Models (LLMs). The principal obstacle is not only training time but also the prohibitive GPU memory overhead associated with long sequences (Liu et al.) 2025b: Wang et al. 2024). Research indicates that significant long-context capabilities can be achieved with a surprisingly small number of training
steps (Peng et al.) 2024 Ding et al. 2024),, and that fine-tuning
little as 5% of the original
pretraining data may suffice for commercial -gade performance (Gao et al. 2025 Bai et al. (024)..
Despite this modest data requirement, the number of tokens processed in single iteration grows dramatically. This causes GPU memory consumption from activations and the KV cache to scale linearly with context length, rapidly exhausting available resources. For instance, model with 4x GQA (Ainslie et al., 2023) ratio requires 64GB for the KV cache alone when processing a 256K context. This demand overwhelms an A000 GPU before even accounting for other network activations, often making 32K tokens a practical limit for single-GPU training (Grattationi et al. 2024 Yang el al.
Equal contribution Corresponding author

---

Published as a conference paper at ICLR 2026
Chunk-wise Training Mechanism
Sparse Attention
Performance
Input Tokens
Standard Porallel Training
Pages
Step1:: Procesa Chunk #3
Simllorty Search
OOMB (Scorse Attr)
Top Indices
150
Step2: Process Chunk
#2
Pased Mennory
75
KV Pages
Backward Processing Direction
Qwen2.5 Soarse Pattern
Content Length
Primary Training Pipeline
CPU
Sparsel(Cachhe()
Peak Memory Usage (GB) on single H200
# forward propagation
Retrieve Part
Fixed Part
141
Standard Panalilel Tra nnx
inferenco_mode(::
GPT-OSS Sparse Attention
COMB (Dense Attnn
model(nnputs, ky. cache)
COMB (Sparse Attr)
#l backward propagation
Persed Hemory
irps reversod [h
outout mocellips, kv. cache)
fetch
GPT-OSS Sparse Pattern
$7.75
loss outputt [loss"
pacchegre.proccespp.. less.heclowrdd))
10
Local Windtw
Content
An overview of the OOMB training framework. (Left) OOMB processes sequences chunk- by-chunk with activation recomputation, maintaining a constant activation memory footprint (O(1)) and shifting the bottleneck to the KV cache. (Center) The growing KV cache is managed through a combination of paged memory, asynchronous CPU offloading, and page-level sparse attention. (Right) Performance benchmarks show that the integrated system drastically reduces peak GPU memory, enabling efficient training on long contexts with minimal memory overhead.
Figure 1:
2025). This immense memory pressure renders popular techniques like ZeRO3 (Rajbhandari et 2020 and tensor parallelism (Shoeybi et al., 2020) insufficient.
While various training-free methods for context extension exist (Jin et their practical utility is limited. Prior work has established that common evaluation metrics like perplexity and targeted retrieval are unreliable indicators of true long-range reasoning (Gao al.) 2024
2024 Han al. 2024).
Bai et al. 2024). Achieving robust performance necessitates dedicated contiguous pretraining.
To this end, this paper directly confronts the central problem: how to train longer-context models with fewer resources without compromising performance. We introduce Out Of the Memory Barrier (OOMB), a chunk-recurrent training framework that processes sequences segments. During the forward pass, each chunk's activations are computed and then immediately discarded. For the backward pass, they are recomputed on-the-fly. This strategy maintains a constant activation memory footprint, regardless of the total sequence length. This approach, however, shifts the primary bottleneck to the KV cache, which must retained throughout the training step and scales linearly with context length.
To overcome this challenge, OOMB integrates a system of co-designed optimizations for efficient KV cache management:
Paged KV Cache and Gradient Management. We introduce paged memory manager, inspired by (Kwon et al. 2023). for both the KV cache and its gradients. This design eliminates expensive memory operations and mitigates fragmentation when appending
key-value pairs.
Specialized Kernels. We implement custom operators that execute all operations the KV cache and its gradients directly, making them opaque to PyTorch's 's autograd system. This avoids storing the KV cache activation and allows for direct gradient accumulation.
Asynchronous KV Cache CPU Offloading. We designed asynchronous mechanism to preemp- tively offload the KV cache and its gradients, whose memory grows with context length, to CPU memory. The resulting data transfer latency is effectively masked by computation.
2

---

Published as a conference paper at ICLR 2026
Table Comparison of Long-Context Training Methods. OOMB achieves (1)) activation memory complexity, where N denotes the context length. This key advantage allows OOMB, when combined with other techniques, to practically enable single-GPU training on contexts exceeding 4M tokens.
Method
GPUs
Activation
Exact Attn
Max Tokens
SeCO (Li et al.
2025
O(1)
128K
1xH200
LongL.RRA (Chen et
2024
O(N)
128K
ZcRO3 Offload
O(N)
V
128K
Ring Flash Attention
8xH200
((N/8)
256K
Tensor Parallelism
((N/8)
256K
OOMB + Dense Attn
O(1)
4M
x2H000
OOMB Sparse Attn
((1)
X
4M+
Page-level Sparse Attention. Our paged memory architecture natively supports page-level sparse attention (Yuan et al. 2025 et al. 2025). This reduces the computational complexity of attention and minimizes the communication overhead of our offloading mechanism.
The synergy of these techniques yields exceptional GPU memory efficiency. Our results show that for every additional 10K context tokens, the end-to-end training memory overhead increases by mere 10MB for Qwen2.5-7B. This allows training Qwen2.5-771 with a 4-miilion-token context on a single H200 GPU, contrast, context parallelism methods (Liu al. 2024 Lin 2025 Shoeybi et al.) 2020) require large clusters for similar tasks (Xu et al. 2025), highlighting substantial advance in resource efficiency.
RELATED WORK
Efficient Long-Context Extension. core challenge in extending language models long contexts is their failure to generalize to positional encodings not seen during pretraining (Han et al. 2024).
While continued training can resolve this issue, the associated GPU memory costs are often prohibitive.
This has motivated training-free free context extension methods, which typically modify positional
encodings to accommodate longer sequences (Chen al.) 2023 Han al. 2024, JTn al.) 2024] An et al. 2024). However, these approaches introduce significant trade-offs. Strategies based on interpolation or extrapolation of positional encodings can degrade performance on short-text tasks by reducing the resolution of positional information (Gao et al. (0244). Others that reuse position indices achieve high scores on perplexity benchmarks but have been shown to fail on tasks requiring deep contextual understanding, indicating superficial grasp of the extended context (Gao al. 2024).
Given these limitations, state-of of the-art models still rely on fine-tuning to expand their context windows (Grattafiori et al.) 2024: Yang al. 2025 Liu al., (025a). While recent advances in parameter-efficient fine-tuning (PEFT) have successfully minimized the memory footprint of model weights and optimizer states (Zhao al. 2025 Zhang al. 2025 Zhong et al.) 2025), they do alleviate the activation memory bottleneck that scales linearly with sequence length. Consequently. the prevailing trend has shifted from algorithmic efficiency (Chen et al. 2024) toward overcoming
memory barriers with massive computational resources (Liu et al.) 2024 Li al. 2024 Jacobs et al.
2024 Shoeybi et al. 2020). For instance, recent efforts have required 256-GPU cluster to extend
model's 8 context window to four million tokens Xu et al. 2025).
Serial vs. Parallel Training Paradigms. Parallel training processes an entire sequence in a single pass, maximizing GPU utilization and scalability. However, its memory footprint scales linearly with sequence length, creating a significant bottleneck for long-context models. In contrast, serial training is highly memory-efficient as it only activates small part of the network at any given time, but at the cost of reduced parallelism.
For long-context fine-tuning, where dataset sizes are often moderate, memory consumption rather than throughput is the primary constraint (Peng et al. 2024: Gao et al. 2024). This observation motivates hybrid approach that processes sequences in chunks: parallel within each chunk and serial between them. This strategy is standard in modern LLM inference engines (Kwon et al., 2023
Ye et al. 2025) where it incurs negligible latency. Its application training, however, remains
3

---

Published as a conference paper at ICLR 2026
Sparse attention lever intra- locality,
The page-level granularity of sparse attention aligns with underlying paged memory system, enabling and rapid access,
4
exhibit attention petterns, maximiz the efficiency of the retrieval process.
Chunk-wise Training
Paged KV Cache
2
training, growth of the KV cache necessitates paged management to prevent fragmentation and realiocation.
Sparse attention enables the selective offloading thereby signfiiantlly reducing communication overhead.
6
3
Offloading the linearty growing KV cache the key bottleneck chunk-wise training. ylelding a constant O1)) memory
The page level granularity of the paged 6 KV cache enables selective transfer data, making offloading flexible highly officient.
Sparse Attention
KV Cache Offload
Figure 2: The synergistic architecture of OOMB's efficient KV cache management relies four deeply interdependent core components: chunk-wise training, paged key-value cache, sparse attention, and key-value eache offload. Each optimization enables and enhances the others, collectively resolving the memory bottleneck of long-context training.
underdeveloped. An early exploration, SeCO (Li et al., 2025). introduced chunk-wise training but lacked critical optimizations the operator and memory management levels. Consequently, it could only train a 16K context model on a single consumer-grade GPU. Under identical conditions, our fully-optimized system extends this capability by more than tenfold.
Context Parallelism. Context Parallelism (CP) and our proposed method both segment long sequences but diverge significantly in their scalability and overhead profiles. CP, inspired by Ring Attention (Liu et al. 2024) and implemented in frameworks like Megatron-LMM (Shoeybi et al. 2020 and DeepSpeed Ulysses (Jacobs et al., (222), distributes single sequence across multiple GPUs. This design introduces substantial inter-GPU communication overhead that scales with the number of devices. In contrast, the overhead of our serial approach is limited repeated KV cache accesses and the latency from CPU offloading. The trade-off is clear: CP excels with large GPU clusters and specialized bandwdhh networks, whereas our method maintains moderate and predictable overhead, eliminating the need for such extensive infrastructure.
3 PRELIMINARY: CHUNK- WISE TRAINING
The core strategy of OOMB is rooted in the principles of sequential processing, which have long been used make training Recurrent Neural Networks (RNNs) memory efficient. We first revisit this foundation and then describe how it can be adapted to the Transformer architecture, setting the stage for the challenges that OOMB is designed to solve.
Activation Recomputation in Sequential Models. In a standard RNN, the forward pass iteratively computes an output yi and a new hidden state m. at each timestep i:
Vi m, a; FRNN(E; m;-1; +)
where represents the intermediate activations cached for the backward pass. For long sequences, storing all activations (a,,..,,ay) } creates a memory bottleneck that scales linearly with sequence length. Activation recomputation (Sohoni et al.) 2022; Bencheikh et al. 2024) circumvents this by discarding aj during the forward pass and regenerating it on-the-fly fly just before its use the backward pass. This trades modest amount of computation for significant reduction in memory.
Adapting Sequential Processing to Transformers. While Transformers are inherently parallel, they can be trained using similar chunk-wise sequential paradigm. We partition a long input sequence into S chunks, {X1, X2, ..., Xs}, }. and process them serially. The forward pass for chunk attends to the KV caches from all preceding chunks, which act analogously
RNN's hidden state:
Y,, M,, A; Firanstomer (X,, {M1, M2 M;
(2)
Here, M, is the KV cache generated by chunk i, and A, represents its intermediate activations. By applying the same activation recomputation strategy, we discard A; after the forward step and
4

---

Published as a conference paper at ICLR 2026
lnvibble to the AutoGrad system
Autodirad System
Invéible the AutoGrad system.
Mensory Managemant Without Paging
CE GralChe
Lie 3000
Lood From HEM
VDOAee grdienn
booum Back To HEM Tntoet 2080
new imnoushi
Book
Forward
(mstual)
masaggrrupstto Jan,, (en)))
Reellgggesstttniiirg Lraw, VJONA nussage))
Our paged KV cache management system for training. A custom kernel handles both forward and back- ward passes, updating gradients in- in-place to bypass the PyTorch autograd system for greater efficiency.
Figure 3:
Figure 4: A non-paged implemen- tation can result in memory usage nearly 3x higher than the theoretical requirement.
recompute it for the backward pass:
dO, (dM,...., dM,-1} backward(dY,, Ai)
(3)
This chunk-recurren approach ensures that the activation memory footprint remains constant, de- termined only by the size of single chunk. Consequently, this paradigm fundamentally shifts the memory bottleneck: the challenge is longer the activations, but the management of the KV cache {M,,..., } and its gradients, which still grow linearly with the sequence length.
4 THE OOMB TRAINING SYSTEM
OOMB builds upon the chunk-wise training paradigm by integrating three synergistic components: paged KV cache management, sparse attention, and asynchronous KV cache offloading. As illustrated in Figure 2 these components are highly interdependent, and their combined effect is critical to the system's overall efficiency. This section provides detailed description of each component.
4.1 PAGED MEMORY MANAGEMENT FOR KV CACHE AND GRADIENTS
Chunk-wise training requires the KV cache to grow incrementally, process analogous autore- gressive decoding. Standard tensor concatenation for this purpose inefficient, leading frequent memory reallocations, copy operations, and severe memory fragmentation (Kwon et al. (023). As shown in Figure this inefficiency causes memory consumption to grow disproportionately with context length, limiting scalability.
To resolve this, OOMB incorporates a paged memory manager for both the KV cache and its gradients. Since existing paged attention solutions designed for inference and lack backpropagation support, we developed custom Triton kernel that handles both forward and backward passes (Figure 3)
Our implementation is designed to bypass the PyTorch autograd engine. This approach provides two key benefits. First, custom CUDA kernels accumulate gradients in-place using at omi c_add, which eliminates intermediate buffers and reduces memory I/O. Second, decoupling the KV cache from the autograd graph prevents Py Torch from storing it as an activation. This design simplifies memory management and facilitates efficient offloading by maintaining a leaner computation graph.
4.2 SPARSE ATTENTION
The quadratic complexity of attention is primary bottleneck in long-context training. OOMB addresses this by integrating page-level sparse attention, which reduces computational cost and naturally supported by paged KV cache architecture (Chen al., 2024 Li et al. 2025 We implement two variants of sparse attention tailored to different model architectures.
Approximating Dense Attention for Qwen2.5. For models with dense attention like Qwen2.5 Yang et al. 2025), we employ a Top-K page retrieval method to approximate the full attention pattern (Tang The selected page indices from the forward pass are cached for reuse
al. 2024 Lu
2025).
5

---

Published as a conference paper at ICLR 2026
K Page Index
Context - 32K
Context - 64K
O
Layer
Layer
Layer
Layer 3
Layer 0
Layer
Layer
Layer
Context - 128K
Context 256K
Layer
Layer1
Layer
Layer
Layer
Layer
Figure 5: Visualization the page-level sparse attention patterns for the Qwen2.5-7B model. Each subplot displays the retrieved key pages (X-axis) for each query page (y-axis) across the initial four layers at context lengths scaling from 32K to 256K.
during backpropagation. The retrieval process begins by computing a single representative vector Kavg RTx for each of the n key pages by averaging. Given queries Q E (mxxP)) from the current chunk, where m is the number of query pages in the chunk and P is the page size (tokens page), we compute similarity scores between each query token and page-representative vectors. voting mechanism then aggregates these scores for each of the m query pages to identify the most relevant key pages:
Score(i. < Pii K [k), R
Score < softmax(Score), MxPx
Score +- sum Score )
m. X
mx Pxn
This retrieval process, visualized in Figure incurs negligible computational overhead with page sizes (P) of 128 on H200 GPUs.
Native Sparse Attention for GPT-OSS. For natively sparse models like GPT-OSS (Agarwal et al. 2025). which alternate between dense and local attention layers, our implementation is straightforward. The dense layers use the same Top-K retrieval method described above, while the local attention layers retrieve only the most recent KV pages.
4.3 ASYNCHRONOUS KV CACHE OFFLOADING
The KV cache is the only component in OOMB whose memory footprint scales with sequence length. To enable training on extremely long contexts, introduce asynchronous offloading mechanism that transfers the KV cache between GPU and CPU memory. The offloading strategy is adapted to the attention type maximize the overlap between data transfer and computation.
Offloading for Dense and Local Attention. For these attention patterns, the required KV pages for given layer are known in advance. We employ pre-fetching strategy. During the forward pass for layer: - 1, the KV cache for layer is asynchronously fetched from the CPU. A symmetric process occurs during the backward pass. This approach effectively hides most of the data transfer latency. resulting in an end-to-end overhead lower than 5%.
satrivaal-basen sparse attention, the required KV pages are known until the query vectors for the current layer have been computed. Consequently, during the forward pass, we initiate page retrieval and the corresponding asynchronous data transfer immediately after the query projection. The small volume of data required by sparse attention allows this transfer to be overlapped with the subsequent key and value projection computations. For the backward pass, the required page indices already cached from the forward pass, allowing to the same pre-fetching strategy as in dense attention.
Offloading for Sparse Attention.
Both offloading schemes, illustrated in Figure 6 are implemented using pinned CPU memory and dedicated CUDA streams to ensure true asynchronicity. Data transfers are handled by the GPU's Direct Memory Access (DMA) engine minimize latency.

---

Published as a conference paper at ICLR 2026
Forward Propagation
Forward Propagation
Attma
Self
MLP
Pood X
Frod
Rop
RoPE
Lnn
MLP
launch
sync
launch
sync
Onload KV Cache
Onload Sparse KV
Memory
DMA Bus
Memory
Onload Grad
Onload Sparse Gradi
launch
sync
launch
Backward Propagation
Backward Propagation
Figure 6: Asynchronous KV cache offloading in OOMB. (Left) For dense and local attention, the KV cache for the next layer pre-fetched during computation of the current layer. (Right) For sparse attention, the forward pass overlaps data transfer with key-value projections, while the backward pass reverts to the pre-fetching strategy.
EXPERIMENTS
5.1 EXPERIMENTAL SETUP
Unless otherwise specified, all experiments adhere to the following configuration.
Model and Dataset. We use the Qwen2.5-7B model (Yang al. 2025). powerful foundation model, for all efficiency and scalability benchmarks. The training data is sourced from the arxiv dataset (Soboleva al. 2023). To construct sequences of the required lengths for our long-context experiments, we concatenate samples from the dataset.
Baseline Methods. For efficiency evaluation, we compare our method against two baselines: (1) standard parallel training setup using FlashAttention (Dao (022)) and layer-wist gradient check- pointing; and (2) Ring Flash Attention (RFA) (Lin 2025 state-of-the-art implementation of Ring Attention (Liu et al. 2024).
Additional details are provided in Appendix B
5.2 MEMORY, TIME, AND SCALABILITY ANALYSIS
We evaluate OOMB's efficiency against two paradigms: standard parallel training and context parallelism.
Comparison with Parallel Training. Table presents comprehensive efficiency comparison between OOMB and a parallel training baseline using FlashAttention. Our analysis highlights several key findings:
At 32K context length, OOMB with dense attention already outperforms the baseline, and this performance advantage widens as the context length increases.
The asynchronous offloading mechanism imposes negligible latency overhead on dense attention, demonstrating its high efficiency.
Sparse attention provides substantial acceleration longer sequences, yielding speedups of approximately 3.7x x at 64K and up to 13.5x at 256K.
For sparse attention with offloading, a larger chunk size mitigates latency; a 12K chunk size, for example, adds only 12% overhead.
Our offloading system dramatically improves memory scalability for both dense and sparse attention, keeping peak memory usage nearly constant irrespective of context length.
Even without CPU offloading, OOMB reduces the memory growth rate by more than fivefold compared to the parallel training baseline.
7

---

Published as a conference paper at ICLR 2026
Table 2: Performance and Ablation of OOMB for Qwen2.5-7B on an H200 GPU. Per-iteration latency and peak memory for various OOMB configurations (ablating chunk size, token budget, gradient checkpointing, and CPU offload) and a FlashAttention (Dao 2024) baseline. Red intensity indicates:
(I) for latency, the slowdown vs. the Ckpt-only counterpart; (2) for memory, the growth rate as context length increases. Single-GPU offload results are reported under full 8-GPU system load.
Method
Budget
Coaftguration
Per Per-iteration Latency
Peak Memory Usage (MB)
Chyp
Office
6K
16K
32K
64K
128
256K
16K
32K
64K
126K
256K
Baseline
2.4
5.7
20.4
76.1
00M
0OM
30298
47005
69985
1000516
0OM
DOM
Chunk Site
4K
00MB Donse Are
2.6
2.4
6.0
6.0
19.3
18.6
63.4
68.9
260.5
233.9
1015.9
899.6
57369
34310
57529
36326
57743
39804
53175
44390
59065
55143
60736
76650
300
20.7
70.7
260.4
985.8
33686
33848
34076
34504
35404
37129
2.7
4.4
8.8
18.0
36.6
75.2
34317
36335
39826
64409
55157
76719
4K+512
3.0
4.5
9.5
19.6
40.0
8333
57346
57372
57377
57380
574.0
57446
2.8
5.2
10.8
46.5
96.9
33669
33694
33990
33712
33732
35769
2.7
4.9
10.6
21.9
44.9
93.2
34117
36335
39827
64811
55198
76725
Sparee
00MB-
4K+2018
2.9
3.1
6.2
5.5
11.5
13.4
28.6
24.5
50.9
127.5
107.2
57346
33669
57401
35732
57413
33735
33778
57456
57479
33800
57515
33337
2.6
5.9
14.3
31.5
66.2
138.1
34117
36335
30827
44413
55183
76735
4K+ 8192
2.4
6.7
16.3
57.0
K1.4
17K.6
57348
57421
S7487
57614
57647
$7752
2.9
7.4
18.3
42.8
94.0
201.4
33669
33745
33809
33857
34083
34063
2.5
6.0
19.1
61.1
148.
924.9
34117
36335
99626
44816
55196
76768
4K+ 32768
2.4
6.4
20.6
67.1
154.7
192.1
57348
57421
57520
57716
57716
58238
3.0
7.2
22.7
79.8
187.7
427.4
33669
33743
33842
34036
34436
34503
Clunk Neer AK
Desse Atie
0UMB +
3.0
2.3
2.8
5.7
5.5
176
19.0
19.7
68.4
64.3
70.6
237
259,0
264.3
1021.4
1014.5
922.2
31628
30374
71025
85705
36293
40795
85808
43483
38485
86382
45850
35870
87051
59613
39639
86589
H1119
41177
2.5
5.1
11.4
24.1
49.6
101.9
31836
40796
45488
45872
59639
81181
SK4512
2.2
5.5
12.1
25.3
53.2
107.8
70979
8542
85634
85650
85604
85726
2.5
13.4
26.7
59.9
123.7
3038
38193
38215
38253
38254
36270
2.5
5.4
12.4
26.3
54.4
111.6
31826
40797
45489
45873
59640
81183
Sparse
OOMB-
8K4-248
2.6
2.4
6.00
5.4
12.6
14.7
27.9
33.9
67.3
58.2
142.1
120.7
78979
3030
85610
38201
85635
38226
85716
36262
85732
38317
85757
38364
2.7
5.8
16.0
36.6
78.1
162.6
31526
40797
43489
43874
59646
81392
SK 8192
2.4
5.7
16.3
36.7
85.2
183.4
78979
85613
85645
85813
85952
85250
3.0
0.4
17.9
43.7
94.9
204.7
3038
38204
35286
38319
38427
38A08
2.8
5.7
18.9
63.1
159.8
345.6
31526
40797
43488
45877
59658
H1227
HK 32768
2.4
5.7
18.9
64.88
165.0
355.2
70979
85513
85713
85917
86122
H6112
3.0
0.4
20.8
70.8
177.7
408.5
30380
38204
38304
35508
38731
38556
Cluss Nice 12K
Dosse Alla
00MB +
2.8
2.8
2.3
6.1
6.1
5.7
19.6
18.9
19.6
68.2
70.9
70.9
25H.7
262.5
262.5
1024.0
1024.0
1017.2
71366
31575
31627
113827
44599
42705
115169
46645
42679
114398
S3654
44081
115199
62735
45858
117984
84913
45558
2.7
4.7
12.5
28.5
60.5
125.3
31626
44599
25689
S2864
62760
X4073
12K4512
2.4
4.7
12.3
25.6
60.3
126.3
70979
113753
113765
113810
113881
113490
2.9
5.4
14.0
32.8
73.8
145
30579
42639
42681
43698
42737
42774
2.8
4.9
13.5
31.3
66.7
138.1
31626
44599
25609
S2865
62761
84975
0OMB- Spare: Alm
12K+2048
2.4
3.0
5.7
5.1
13.9
15.2
32.8
36.4
09.8
78.9
146.7
163.5
70979
30379
113753
42509
113817
42763
133830
42756
113923
42809
113963
42545
2.8
5.7
17.6
43.4
90.6
188.7
31626
44593
25609
S2067
62765
88985
12K+ 8192
x
2.4
5.8
18.1
43.6
90.5
197.9
70979
113753
113830
1139/1
112017
114138
3.0
6.6
21.5
45.2
105.1
224.1
30379
42639
42716
42850
42968
43018
2.8
5.7
18.9
65.4
170.0
373.9
31626
44589
45609
52070
62778
85019
12K+ 32768
2.4
5.7
18.5
67.2
169.0
385.3
70979
113753
113809
114028
114284
114558
311
65
20.6
74.0
185.5
421.2
30379
42639
42715
42964
43204
43443
Table 3: Per-device training throughput (tokens/sec) comparison against context-paralle methods. 4K chunk size with checkpointing and CPU offload; the sparse
All OOMB configurations attention version adds an 8192 retrieval budget. TP: Tensor Parallel.
Ring Atin (report)
RFA
OOMB (Dense Alin)
OOMB (Sparse Atin)
Context
8xA100
16xTPUv4
4x H200
1xH200
4xH200 (TP)
1xH200
4xH200 (TP)
64K
872.51
936.22
933.84
1560.38
1542.85
128K
463.37
504.12
500.26
1394.16
1371.41
256K
50
49
218.41
266.13
261.47
1301.60
1265.59
Comparison with Context Parallelism Methods. We also benchmarked OOMB against Ring Flash Attention (RFA) (Lin 2025). a ofate-of-the-art implementation of Context Parallelism. This paradigm manages memory by splitting sequences across GPUs but introduces significant inter- device communication overhead. As shown in Table OOMB achieves higher per-device training throughput, outperforming the heavily optimized RFA.
8

---

Published as a conference paper at ICLR 2026
Training Loss Curves
Gradient Approsimation Error Distribution (12 Norm)
4.0
0.6
Butggr: 6C132766
Budget ai-d11z Budget 46--27668
Busigr
Buskgr
Budger
2.6
0.4
1.2
0.1
Contest a 256X
Contest IM
Contect + 16K
Context 64K
Context 250K
Context - TM
Layer-wise Breakdown Gradient Estimation
1.6
12
PrO)
Pro
Proj
Up Proj
Gate Proj
Down Proj
Figure 7: (Top-left) Training loss curves for Qwen2.5-7B YYang al. 2025) with 256K and IM context lengths, employing sparse attention with retrieval budgets of 4K+8192 and 4K+32768. (Top- right) Distribution of the 1.2 norm of gradient approximation error for sparse attention across various context lengths ( (16K to 1M) and retrieval budgets. (Bottom) Layer-wise L2 norm of the gradient estimation error for different projection matrices, analyzed with 64K context length and a 4K+8 192 retrieval budget. Our work offers a preliminary validation of this sparse training; for more optimal sparse attention strategies, we refer readers to the established literature.
5.3 ACCURACY VALIDATION OF SPARSE ATTENTION
We empirically validated the accuracy of our sparse attention mechanism and its impact on model training. First, analyzed the gradient approximation error for Qwen2.5-7B (Yang et al.) 2025) across 16 configurations, composed of 4 context lengths and 4 retrieval budgets. As shown in Figure (Top-right), within the model's 128K native context limit, larger retrieval budget results in smaller approximation error. Beyond this limit, the error differences between budgets become less significant. The loss values for these settings, presented Table further indicate that sparse attention does not substantially degrade model performance.
Next, we performed training runs at 256K and 1M context lengths using 8K and 32K retrieval budgets. The results in Figure 7 (Top-left) show that at 256K, both budgets yield rapid convergence to loss below 1.2 with highly similar loss curves. At a 1M context, the 32K budget exhibited minor instability late in training, effect we attribute to the learning rate and limitations in positional encoding extrapolation. Nevertheless, both configurations converged successfully to low loss values. As detailed in 4bble this approach is practical, allowing IM-context model to be trained for 256 steps in approximately 10 hours on a H200 system with tensor parallelism.
Collectively, these findings confirm that sparse attention is viable and effecctve method for long- context training, as it preserves model performance while achieving a constant O(1) memory footprint.
Table 4: (a) Training loss for the 16 configurations (4 context lengths sparse budgets) in the top-right subplot of 7 (b) Resource usage during training for the contigurations (2 context lengths sparse budgets) the top-left subplot of Figure These experiments were conducted on four H200 GPUs using tensor parallelism, without gradient checkpointing or CPU offload.
(a) Loss values under different settings
(b) Resource usage (4x H200 tensor parallelism)
Contect
Demse Attention
Spanse Attentive
Metrie
4K191
4K+32768
4K+512
4K+20488
4K451922
4K<43768
256K
IM
256K
IM
16K
1.30030
1.50000
1.36718
1.33091
1.30031
64K
1.223887
147656
1.9281
1.31875
1.21875
Latency/lter (6) Peak Memory (MB)
43.454
176.305 67944
87.340
422.383
256K
IM
4,31250
2.21875
2.34355
6.1250
2228125
4,50000
2.21875
4.71875
2,218755
4,78125
35637
35671
68145

---

Published as a conference paper at ICLR 2026
CONCLUSION AND LIMITATIONS
We introduced OOMB, a highly memory-efficient training system that overcomes the primary memory barrier in long-context LLM training. By integrating chunk-recurrent framework with activation recomputation, a paged KV cache manager, asynchronous CPU offloading, and page-level sparse attention, OOMB maintains a nearly constant memory footprint regardless of sequence length. Our results demonstrate that OOMB enables training a Qwen2.5-77 7B model with a 4-million-tokee tokknn context on a single H200 GPU, a task that would otherwise require a large-scale cluster. This work represents significant advance in making long-context model training more accessible and resource-efficient
Despite efficiency, OOMB has several limitations. The chunk- wise serial processing introduces latency overhead compared to fully parallel training paradigms, which may be consideration for throughput-sensitiv training scenarios. Additionally, our use sparse attention is an approximation of the full attention mechanism, and its impact model performance for tasks requiring dense global context requires further investigation. The effectiveness of the asynchronous offloading is also dependent on the available CPU-GPU interconnect bandwidth.
10

---

Published as a conference paper at ICLR 2026
ETHICS STATEMENT
The primary contribution of this work is a system designed to improve the computational efficiency of training large language models on long contexts. By significantly reducing the hardware requirements for such tasks, our research aims to democratize access to long-context model development for researchers and institutions with limited computational resources. A direct positive implication of this efficiency is the potential reduction in the overall energy consumption and carbon footprint associated with training these models, promoting more sustainable Al research practices.
We acknowledge the potential for dual use, is common with foundational Al research. The tools we develop could be applied to train models for malicious purposes, such as generating misinformation. While this work focuses on the system's technical aspects, we advocate for the responsible development and deployment of models trained using our framework. Our experiments were conducted the publicly available arXiv dataset, which consists of scientific literature and does not contain personally identifiable or sensitive information.
REPRODUCIBILITY STATEMENT
To facilitate the reproducibility of our findings, we have made the complete source code for the OOMB training system available in the anonymous repository cited in the abstract. A detailed description of our experimental setup, including model specifications, datasets, hyperparameters, and evaluation protocols, provided in Section 5.1 and Appendix B The core methodology and its components thoroughly explained in Section 4 We believe these resources offer a sufficient basis for researchers to verify our results and extend our work.
REFERENCES
Sandhini Agarwal, Lama Ahmad, Jason Ai, Sam Altman, Andy Applebaum, Edwin Arbus, et al. gpt-oss- 120bbsss-bbbb gpt-oss-20b model card. arXiv preprint arXiv:2508. 2025.
Joshua Ainslie, James Lee- Tho-Th Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, et al. GQA: Training generalized multi-query transformer models from multi-head checkpoints. In EMNLP, 2023.
Chenxin An, Fei Huang, Jun Zhang, Shansan Gong, Xipeng Qiu, et al. Training-free free long-context scaling of large language models. In ICMI., 2024.
Yushi Bai, Xin Lv, Jiajie Zhang. Yuze He, Ji Qi. Lei Hou, et al. LongAlign: A recipe for long context alignment of large language models. In EMNLP, 2024.
Wadjih Bencheikh, Jan Finkbeiner, and Emre Neftci. Optimal gradient checkpointing for sparse and recurrent architectures using off-chip memory. arXiv preprint arXiv:24/2./18/00 2024.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.1553 2023.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, et al. Longlora: Efficient fine-tuning of long-context large language models. In ICLR, 2024.
Tri Dao. FlashAttention-22 Faster attention with better parallelism and work partitioning. In ICLR, 2024.
Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu. Ning Shang, et al. Longrope: extending Ilm context window beyond 2 million tokens. In ICML,, 2024.
Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. How to train long-context language models (effectively). arXiv preprint arXiv:2410.02666 2024.
Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. How to train long-context language models (effectively). In ICI.R, 2025.
11

---

Published as a conference paper at ICLR 2026
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, et al. The llama herd of models. arXiv preprint arXiv:2407.21783 2024.
Chi Han, Qifan Wang, Hao Peng, Wenhan Xiong, Yu Chen, et al. Lm-infinite: Zero-shot extreme length generalization for large language models. In NAACL, 2024.
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Reza Yazdani Aminadabi, et al. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. In PODC, 2024.
Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, et al. Lim maybe longlm: Selfextend Ilm context window without tuning. In ICML, 2024.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In SOSP, 2023.
Dacheng Li, Rulin Shao, Anze Xie, Eric P. Xing, Xuezhe Ma. et al. DISTFLASHATIN: Distributed memory-efficient attention for long-context LLMs training. CoLM, 2024.
Wenhao Li, Yuxin Zhang, Gen Luo, Daohai Yu, and Rongrong Ji. Training long-context Ilms efficiently via chunk-wise optimization. In ACL, 2025.
Zhuzi Lin. Ring flash attention: Ring attention implementation with flash attention. Github, 2025.
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, al. Deepseek-vä v3 technical report. arXiv preprint 21Xv:22412.19437, 2025a.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ringattention with blockwise transformers for near- infinite context. In ICLR, 2024.
Jiaheng Liu, Dawei Zhu, Zhiqi Bai, Yancheng He, Huanxuan Liao, Haoran Que, Zekun Wang, Chenchen Zhang, Ge Zhang, Jiebin Zhang, et al. A comprehensive survey on long context language modeling, arXiv preprint arXiv:2503.17407, 2025b.
Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, al. Moba: Mixture of block attention for long-context llms. arXiv preprint arXiv:2502. 13189, 2025.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. YaRN; Efficient context window extension of large language models. ICLR, 2024.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: memory optimizations toward training trillion parameter models. In SC, 2020.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, et al. Megatron- Im: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2020.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama. HuggingFace, 2023.
Nimit Sohoni, Christopher R. Aberger, Megan Leszezynski, Jian Zhang, and Christopher Ré. Low-memory neural network training: A technical report. arXiv preprint arXiv: / 10631, 2022.
Jiaming Tang. Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest: Query-aware sparsity for efficient long-context Ilm inference. arXiv preprint arXiv:2406. 10774, 2024.
Xindi Wang, Mahsa Salmani, Parsa Omidi, Xiangyu Ren. Mehdi Rezagholizadeh, al. Beyond the limits: a survey of techniques to extend the context length in large language models. In IJCAI, 2024.
Chejian Xu, Wei Ping, Peng Xu, Zihan Liu, Boxin Wang, et al. From 128k to 4m: Efficient training of ultra-long context large language models. arXiv preprint arXiv:2504.067 062 2025.
12

---

Published as a conference paper at ICLR 2026
An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, et al. Qwen2.5 technical report. arXiv preprint arXiv:2412.15115, 2025.
Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, et al. Flashinfer: Efficient and customizable attention engine for Ilm inference serving. arXiv preprint arXiv:2501.01005, 2025.
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. 2025.
Xueyan Zhang, Jinman Zhao, Zhifei Yang, Yibo Zhong, Shuhao Guan, Linbo Cao, and Yining Wang. UORA: Uniform orthogonal reinitialization adaptation in parameter efficient fine-tuning of large models. In ACL, 2025.
Jinman Zhao, Xueyan Zhang, Jiaru Li, Jingcheng Niu, Yulan Hu, Erxue Min, and Gerald Penn. Tiny budgets, big gains: Parameter placement strategy in parameter super-efficient fine-tuning. In EMNLP, 2025.
Yibo Zhong, Jinman Zhao, and Yao Zhou, Low-rank interconnected adaptation across layers. In ACL, July 2025.
A USE OF LARGE LANGUAGE MODELS (LLMS)
During the writing of this paper, we utilized LLM solely for language editing to improve clarity and readability. We critically reviewed and revised all Al-generated suggestions to ensure the final text accurately reflects our original intent. All intellectual contributions, including the research design, methodology, analysis, and conclusions, are our exclusive work, and we take full responsibility for the academic integrity of this publication.
B EXPERIMENTAL DETAILS
All experiments were conducted using full-parae full-parameter fine-tuning in bfloat16 precision. We employed the Adam optimizer with learning rate of 5 x 10-5 and betas of (0.9, 0.98). To isolate the impact of sequence length system performance, the Per-GPU batch size model and training hyperparameters were maintained at their default settings. For OOMB-specific configurations, the KV cache page size was set to 128 for H200 GPUs.
to one, and all other
System performance was evaluated using two primary metrics: peak GPU memory and per-iteration latency. Peak memory per device was recorded using PyTorch's max_ feax__mmmry_aa aaloccaaallltttaaoaaaaeeaeo allocated () function, while latency was measured precisely with CUDA events. To ensure reliability and mitigate measurement noise, all reported results represent the minimum value from three independent runs.
C ADDITIONAL EXPERIMENTS AND ABLATION STUDIES
This section provides supplementary experiments that further validate the performance, scalability, and correctness of the OOMB framework.
PERFORMANCE AND ABLATION STUDIES
Scalability to 4M+ Tokens. A key objective of OOMB is to enable training on sequence lengths that are infeasible with conventional methods on limited hardware. We henchmarked the per-iteration training time and memory usage of OOMB contexts up to 8 million tokens. Figure 8 shows that OOMB maintains exceptional memory efficiency even at this scale. When using sparse attention, the training time scales in near-linear fashion, making ultra-long context training practically achievable.
Ablation on Chunk Size. The chunk size is a key hyperparameter in OOMB that balances com- putational parallelism and activation memory. We evaluated the effect of varying the chunk size from 512 to 4096. As shown in Figure S larger chunk sizes improve computational efficiency by enabling better hardware utilization. However, these performance gains diminish beyond size of
13

---

Published as a conference paper at ICLR 2026
Training Overhead for Million-token Contexts
5000
4,802.139
141000
Time / Ittration (sec)
3750
105750
Peak Memory (MB)
2500
70500
2,061.571
1250
913.111
35250
404.827
2M
4M
8M
Context Length
Figure 8: Scalability to 4M+ Tokens. OOMB demonstrates perfect memory efficiency up to 8-million-token context length. With sparse attention, training time scales more linearly compared to dense attention. This experiment was conducted on H200 GPUs using tensor parallelism.
4096. Memory consumption remains stable across all tested sizes. This analysis supports 4096 as well-balanced default value, as it provides strong computational performance without excessive memory usage for activations.
End2End Latency
Poak Memory Usage
A43 det bot
2000
2006
3008
386
155
130K
1258
1008
1154
2026
156
100K
3256
1756
2002
Figure 9: Ablation Chunk Size. Increasing the chunk size improves computational throughput with diminishing returns, while peak memory usage remains stable. A chunk size of 4096 offers strong balance between efficiency and memory. This experiment conducted single A100 80G GPU.
14


Reviews:
[{"review_id": "bQDdrxY93D", "reviewer": "Reviewer_TNui", "summary": "This paper introduces OOMB (Out Of the Memory Barrier) to train LLM efficiently with ultra-long contexts up to 4–8 million tokens using a single GPU. For GPU memory usage, activations and KV caches scale linearly with sequence length. OOMB combines chunk-recurrent training with a series of memory and compute optimizations that make activation memory constant and KV cache growth manageable.\n\nChunk-Recurrent Training: it splits long sequences into smaller chunks processed serially. Activation recomputation keeps activation memory constant to sequence length. Shifts the main bottleneck from activations to the KV cache.\n\nPaged KV cache and gradient management: it proposed a custom paged memory manager for KV caches and gradients, accumulating gradients directly on GPU memory to reduce I/O and fragmentation.\n\nAsynchronous KV cache offloading: Transfers KV caches between GPU and CPU memory asynchronously using pinned memory and DMA\n\nPage-level sparse attention: introduces top-K page retrieval for dense models (e.g., Qwen2.5) and native sparse attention for GPT-OSS", "questions": "It would be great to include a breakdown analysis of time spent on recomputation vs. data transfer vs. forward/backward compute.\n\nClarify what “O(1)” precisely means in practice (e.g., per-chunk constant, independent of total context but dependent on hidden size).", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "originality\nIt showcased the effectiveness of paged KV cache manager that support forward and backward passes, combined with activation recomputation and asynchronous offloading. OOMB innovates on single-GPU training feasibility for million-token contexts. It's quite different to mainstream approaches\n\nQuality\nThe paper has solid benchmarking results across context lengths, chunk sizes, and sparse retrieval budgets. Comparing to FlashAttention and Ring Flash Attention, improvements in runtime and memory are consistent and reproducible. It has clear formulations for chunk-wise training, sparse attention retrieval, and ablation studies up to 8M-token contexts. It also openly discussed latency tradeoffs with dependency on CPU–GPU bandwidth\n\nClarity: Figures 1–3 and 6–9 effectively illustrate the system architecture and memory scaling behavior. The writing style is clear and professional, with minimal ambiguity.\n\nSignificance: The work addresses memory constraints in long-context training. It's critical for long-tail users with limited GPU resource. Achieving 4M-token context training on a single GPU is a concrete result that could democratize access to long-context modeling. The system offers a practical path forward for future foundation model training", "weaknesses": "Whether baselines are state-of-the-art: The authors state (Section 5.3) that “for more optimal sparse attention strategies, we refer readers to the established literature,” which underscores that the sparse attention component remains under-developed. The analysis only reports gradient approximation errors (Fig. 7) without investigating why the sparse retrieval patterns emerge\n\n\nLimited discussion of evaluation results: Section 6 acknowledges that “chunk-wise serial processing introduces latency,” but the magnitude of this penalty relative to GPU compute budget or wall-clock training time is unclear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LnXeoIJ2Ko", "reviewer": "Reviewer_ZwBd", "summary": "Training large language models (LLMs) on long contexts is constrained by exorbitant GPU memory overhead, primarily due to the linear growth of memory usage with sequence length. This paper proposes a memory-efficient training system named OOMB, which adopts a scheme combining a \"block-wise cyclic training framework\" with \"on-the-fly activation recomputation\". To efficiently manage the KV cache, OOMB incorporates a set of collaborative optimization strategies: (1) a paged memory manager for both the KV cache and its gradients to eliminate fragmentation; (2) asynchronous CPU offloading to hide data transfer latency; (3) page-level sparse attention to reduce both computational complexity and communication overhead. The synergy of these techniques yields exceptional efficiency.", "questions": "Refer to Weaknesses", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Paged KV Cache and Gradient Management. This paper introduces a paged memory manager for the key-value (KV) cache and its gradients to eliminate exorbitant memory operation overhead, while alleviating memory fragmentation when appending new KV pairs.  \n2. Specialized Kernels. This paper implement custom operators that execute all operations on the KV cache and its gradients directly, making them opaque to PyTorch’s autograd system. This avoids storing the KV cache as an activation and allows for direct gradient accumulation.\n2. Asynchronous KV Cache CPU Offloading. This paper designed an asynchronous mechanism to preemptively offload the KV cache and its gradients, whose memory grows with context length, to CPU memory. The resulting data transfer latency is effectively masked by computation.\n3. Page-level Sparse Attention. This technique not only reduces the complexity of attention computation but also minimizes the communication overhead introduced by the offloading mechanism to the greatest extent.", "weaknesses": "1. The description of the workflow of the OOMB training system lacks clarity. Could pseudocode be provided for illustration, similar to FlashAttention?\n2. The experiments are only conducted on small model, and validation on larger or more diverse models is anticipated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tXsK1TU3PF", "reviewer": "Reviewer_dZ2Y", "summary": "This paper introduces OOMB, a highly memory-efficient training system designed to overcome the prohibitive GPU memory overhead associated with training Large Language Models (LLMs) on million-token contexts. The core of the system is a chunk-recurrent training framework that employs on-the-fly activation recomputation, achieving a constant ($\\mathcal{O}(1)$) memory footprint for activations. To manage the remaining bottleneck of the linearly growing KV cache, OOMB integrates a suite of synergistic optimizations: a paged memory manager for both the KV cache and its gradients, asynchronous CPU offloading, and page-level sparse attention. The authors demonstrate that their system can train a 7B model with a 4-million-token context on a single H200 GPU, a significant advance in resource efficiency.", "questions": "1.  The KV cache is also a form of activation. The OOMB framework discards other intermediate activations but retains the KV cache to be passed between chunks . Is the primary reason for this special treatment to ensure the attention mechanism has access to the full history for a mathematically correct forward pass, while other activations can be recomputed locally within a chunk without affecting the model's final output?\n2.  In Table 4a, the results for 16K, 64K, and 256K contexts generally show that the loss decreases as the sparse attention budget increases, approaching the performance of dense attention. However, for the 1M context, the loss is higher with larger budgets (4K+8192 and 4K+32768) than with smaller ones. This trend is counter-intuitive. Could the authors provide a potential explanation for this behavior?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- While the core modules used in OOMB (e.g., chunking, paged memory, CPU offload, sparse attention) are existing concepts, the paper does an excellent job of integrating them into a cohesive and interdependent system. \n\n- The paper provides a comprehensive set of experiments to validate the system's efficiency. The performance benchmarks against standard parallel training and state-of-the-art context parallelism methods are convincing. Furthermore, ablation studies on chunk size and scalability tests up to 8 million tokens effectively demonstrate the system's robustness and practicality.\n\n- The OOMB framework is designed not only for single-GPU scenarios but also supports multi-GPU training via Tensor Parallelism. This flexibility significantly increases the practical value of the work for a wider range of researchers and institutions.", "weaknesses": "- Although the authors analyze the performance of sparse attention, I still have concerns about its impact on model quality. A key missing piece of analysis is an ablation study on the page size. The representative vector for each key page is calculated by averaging. It is plausible that a larger page size would average out more details, reducing the distinctiveness between different page-representative vectors and thus degrading the quality of the Top-K retrieval. This seems like a critical hyperparameter that was not explored.\n\n- The training loss curve for the 1M context length in Figure 7 (top-left) shows some instability, especially for the larger budget configuration. The authors briefly attribute this to the learning rate and limitations in positional encoding extrapolation. This explanation is not sufficient to be fully convinced. Could this instability be an artifact of the sparse approximation itself at such an extreme scale? Further analysis is needed to clarify this point.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "csRNn2T0XZ", "reviewer": "Reviewer_qtG7", "summary": "The paper presents OOMB, a system for memory-efficient training of LLMs on long contexts. The approach combines chunk-recurrent training with activation recomputation, paged KV cache management, asynchronous CPU offloading, and sparse attention. The authors claim to enable training of Qwen2.5-7B with 4M-token contexts on a single H200 GPU.", "questions": "Table 3 shows OOMB is only marginally faster than Ring Flash Attention per-device. When is OOMB preferable vs. context parallelism on a cluster?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "1. The paper addresses an Important Problem: GPU memory is indeed a critical bottleneck for long-context LLM training, making this a relevant research direction.\n\n2. The integration of multiple techniques (paging, offloading, sparse attention, custom kernels) shows good systems engineering.\n\n3. Strong Memory Efficiency: The claimed 10MB per 10K tokens overhead is impressive if validated.\n\n 4. Clear Presentation: The paper is generally well-written with good use of figures to illustrate the approach.", "weaknesses": "The paper demonstrates memory efficiency but never validates that the trained models actually work. In particular, the paper did not provide evaluations on long-context benchmarks (LongBench, RULER, BABILong, InfiniteBench), comparisons of final model quality between dense and sparse attention. Meanwhile the training curves stop at ~250 steps with loss ~1.2-4.7, and no evidence of convergence, and only used synthetic data. It would be great if the authors can provide complete training runs with evaluation on established benchmarks, and compare dense vs. sparse attention on major tasks.\n\nThe paper claimed that it can \"training Qwen2.5-7B with a 4M-token context on a single H200 GPU\". However, figure 8 (4M experiment) uses 4×H200 with tensor parallelism, not single GPU, 4M requires sparse attention (Table 2 shows dense attention OOMs at 256K on single GPU), and table 1 lists \"OOMB + Dense Attn\" at 4M—contradicts the data. Can the author clarify w.r.t this matter?\n\nThe paper has limited technical novelty, as each component exists: chunk-wise training: Standard for RNNs, used by SeCO for transformers; paged KV cache: directly from vLLM; CPU offloading: common practice; Sparse attention: well-established.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)