We present MegaTrain, a memory-centric system that efficiently trains 100B+ parameter large language models at full precision on a single GPU. Unlike traditional GPU-centric systems, MegaTrain stores parameters and optimizer states in host memory (CPU memory) and treats GPUs as transient compute engines. For each layer, we stream parameters in and compute gradients out, minimizing persistent device state. To battle the CPU-GPU bandwidth bottleneck, we adopt two key optimizations. 1) We introduce a pipelined double-buffered execution engine that overlaps parameter prefetching, computation, and gradient offloading across multiple CUDA streams, enabling continuous GPU execution. 2) We replace persistent autograd graphs with stateless layer templates, binding weights dynamically as they stream in, eliminating persistent graph metadata while providing flexibility in scheduling. On a single H200 GPU with 1.5TB host memory, MegaTrain reliably trains models up to 120B parameters. It also achieves 1.84× the training throughput of DeepSpeed ZeRO-3 with CPU offloading when training 14B models. MegaTrain also enables 7B model training with 512k token context on a single GH200.
Figures & tables
Figure 1. Sustained TFLOPS across model scales. MegaTrain remains efficient while offloading baselines become GPU memory bound. Line plot of sustained training throughput across model scales on GH200 and H200. MegaTrain stays high and stable while the offloading baselines decline and eventually fail at larger model sizes.
Tier
Capacity
Bandwidth
Unit Cost ($/GB)
SRAM (on-chip)
∼ 50–112 MB
∼ 80 TB/s
—
H200 SXM
HBM3e (per GPU)
141 GB
4.8 TB/s
∼ 20
Host DDR5
2–4 TB
∼ 200 GB/s
∼ 5–12
PCIe Gen5 link
—
128 GB/s
—
GH200
Table 1. Memory hierarchy of H200 and GH200.
Figure 2. MegaTrain architecture: CPU serves as the parameter store while GPUs execute transient layer templates via asynchronous parameter streaming and gradient offloading. System diagram showing Host (CPU) memory as the persistent store for parameters and optimizer states, and the GPU as a transient execution engine that streams weights in, computes each layer, and offloads gradients back to host memory.
Figure 3. End-to-end pipelined execution. Weight prefetch ( W ), computation ( F / R / B ), and gradient offload ( G ) overlap across three CUDA streams. Timeline diagram of three CUDA streams. Weight transfers, forward and recompute and backward compute, and gradient offloads overlap in a double buffered schedule so communication stays off the critical path.
Model
Total Params
Layers
Hidden Size
FFN Size
Qwen2.5-7B
7B
28
3584
18944
Qwen2.5-14B
14B
48
5120
13824
Qwen2.5-32B
32B
64
5120
27648
Qwen2.5-72B
72B
80
8192
29568
GPT-OSS-120B (MoE)
120B
36
2880 × 128
2880 × 128
Table 2. Model configurations used in experiments.
Figure 4. Host (CPU) memory footprint versus model scale across training systems. Line plot of host memory usage versus model size for MegaTrain and baseline systems. MegaTrain grows close to linearly while the baselines rise faster and cross practical memory limits at smaller scales.
Metric
Baseline
ZeRO-3 Offload
ZeRO Infinity
PyTorch Native
Ours
7B Acc. (%)
33.47
88.93
88.97
88.91
88.99
14B Acc. (%)
37.58
92.41
92.36
-
92.52
Table 3. Final accuracy comparison across systems at 7B and 14B scales.
Configuration
BS
TFLOPS
Device Mem (GB)
MegaTrain (full)
96
266.3
75.93
w/o Double Buffering
96
182.91
74.11
w/o Gradient Slab Pool
96
257.55
75.93
w/ Checkpoint Interval=1
64 (96 OOM)
240.45
81.34
Table 4. Ablation study of MegaTrain. Removing double buffering significantly degrades throughput, while other components have minor impact.
Layers
Parameters (B)
Device Alloc (GB)
28
7.62
3.83
42
10.88
3.83
56
14.14
3.83
84
20.67
3.83
132
31.85
3.83
180
43.04
3.83
Table 5. Depth study model configurations.
Figure 5. Depth scalability with fixed model width (hidden and FFN) size. Two panel figure for depth scaling at fixed width. The left panel shows MegaTrain throughput staying comparatively stable as layer count increases while FSDP and ZeRO offloading degrade and fail. The right panel shows MegaTrain using less host memory than the baselines across the same models.
Figure 6. Width scalability with fixed model layers Two panel figure for width scaling at fixed depth. Throughput decreases for all methods as width grows, but MegaTrain degrades more slowly and remains feasible at wider models. Host memory rises with width for all methods, with MegaTrain staying below the main offloading baselines.
Width Scale
Hidden Size
FFN Size
Device Alloc (GB)
1.0x
3584
18944
3.83
1.5x
5376
28416
7.01
2.0x
7168
37888
11.07
2.5x
8960
47360
15.99
3.0x
10752
56832
21.78
3.5x
12544
66304
28.44
Table 6. Width study model configurations.
Context
BS
Tokens
Step (s)
TFLOPS
Mem
1K
158
162.7K
27.05
284.7
74.2 GB
8K
25
204.8K
32.36
294.5
86.5 GB
32K
6
196.6K
32.18
316.7
84.0 GB
128K
1
131.1K
26.13
305.3
62.1 GB
256K
1
262.1K
236.1
401.2
88.2 GB
512K ∗
1
524.3K
871.4
407.4
81.9 GB
Table 7. Long-context training performance on GH200. TFLOPS is computed using 6ND+12LHT2 . * denotes results with chunked MLP execution for ultra-long contexts.
A100 PCIe
RTX A6000
RTX 3090
Peak Tensor TFLOPS
312.0
154.8
142.0
Device Memory
80 GB HBM2e
48 GB GDDR6
24 GB GDDR6X
Host Memory
600 GB DDR4
251 GB DDR4
251 GB DDR4
PCIe BW
32 GB/s
32 GB/s
16 GB/s
Host BW
140.8 GB/s
204.8 GB/s
204.8 GB/s
Table 8. Nominal hardware characteristics of the three PCIe-based systems used in this subsection. Peak tensor throughput, PCIe bandwidth, and host memory bandwidth are nominal hardware specifications. PCIe bandwidth is the peak per-direction rate of an × 16 link.
Figure 7. Performance comparison on a single A100 PCIe system. Bar chart comparing achieved TFLOPS on a single A100 PCIe system for 7B, 14B, and 32B models. MegaTrain outperforms Gemini and ZeRO 3 across the supported scales and continues running when the baselines hit memory limits.
Model
Method
Max BS
TFLOPS
GPU Mem
CPU Mem
tok/s
A6000 (48GB VRAM)
Qwen2.5-3B
MegaTrain
15
49.70
46.74 GB
38.3 GB
2153
ZeRO-3 Offload
1
23.89
20.33 GB
–
–
Qwen2.5-7B
MegaTrain
12
55.73
44.74 GB
56.7 GB
1219
ZeRO-3 Offload
1
27.55
20.83 GB
–
–
Qwen2.5-14B
MegaTrain
9
56.82
44.64 GB
104.1 GB
641
Table 9. Training performance comparison between MegaTrain and ZeRO-3 Offloading on A6000 (48GB) and RTX 3090 (24GB). ZeRO-3 results are reported with batch size (BS)=1; BS=2 leads to out-of-memory (OOM).
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Model
TFLOPS
7B
2.03
14B
10.90
32B
10.91
Appendix
Table 10. Ratel reproduction results on GH200 using the official codebase.
Training large language models on limited hardware is increasingly a scheduling problem across GPU compute, host memory, PCIe transfer, and storage bandwidth. Existing offloading systems reduce GPU residency, and MegaTrain shows that a CPU-master layer-streaming executor can train large models on a single GPU, but fixed checkpointing and placement heuristics still leave communication exposed on the critical path. We propose LazyTrain, an optimization layer over a layer-streaming executor. LazyTrain formulates checkpoint selection, activation placement, recomputation, and CPU-GPU-NVMe communication overlap as a mixed-integer scheduling problem, then executes the solved policy during training. It further couples 8-bit optimizer states with fast gradient clipping as a single Hybrid 8-bit operator: state compression reduces optimizer-state memory, while fast clipping counteracts the additional CPU-side update overhead. Across H800 experiments from Qwen2.5-3B to Qwen3.6-27B, LazyTrain improves sustained TFLOPS over matched baselines runs by approximately 1.24×; RTX 3090 experiments likewise increase the maximum feasible batch size by one at each model scale. In the primary Qwen3.6-27B H800 MetaMathQA run, LazyTrain reaches 219.95 TFLOPS and 1361 tokens/s at batch size 72, peaks at 68.84,GB of GPU memory, and obtains 95.42% exact-match accuracy on the full evaluation split. The source code is available at https://github.com/DataArcTech/LazyTrain.
Xiaojun Wu, Cehao Yang, Honghao Liu +5
1IDEA Research · 2The Hong Kong University of Science and Technology (Guangzhou) · 3DataArcTech Ltd.
Fine-tuning large language models (LLMs) often exceeds GPU memory limits, prompting systems to offload model states to CPU memory. However, existing offloaded training frameworks like ZeRO-Offload treat all parameters equally and update the full model on the CPU, causing severe GPU stalls, where fast, expensive GPUs sit idle, waiting for slow CPU updates and limited-bandwidth PCIe transfers. We present ZenFlow, a new offloaded training framework that decouples updates between GPU and CPU by prioritizing important gradients. ZenFlow performs in-place updates of important gradients on GPU while offloading and accumulating less important ones on CPU without blocking training iterations, fully overlapping CPU work with GPU computation. To make this importance-aware design scalable, ZenFlow leverages a novel spatio-temporal locality of important gradients, concentrated in specific positions that remain stable across iterations, to make important gradient identification efficient and lightweight while preserving fidelity. Evaluation shows that ZenFlow achieves up to 5x end-to-end speedup, 2x lower PCIe traffic, 85% less GPU stalls, and up to 72% lower cloud GPU cost, all while preserving accuracy.
Tingfeng Lan, Yusen Wu, Bin Ma +7
University of Virginia · University of California, Merced · Argonne National Laboratory +2
Progress in deep learning is, at scale, more a matter of systems engineering than of modelling: the behaviour of a model in training (its throughput, its memory footprint, and the numerical fidelity of the result) is determined less by the architecture itself than by how that architecture is expressed on the hardware. To achieve absolute control over this hardware expression while abstracting away systems complexity to make modelling seamless and eliminating the need for repetitive orchestration logic, BluTrain was architected from first principles as a robust, lightweight, and architecture-general training framework in standard C++ and the core CUDA programming model. Every layer is implemented natively: a typed tensor module with reverse-mode autograd, a linear-algebra library, a caching allocator, a multi-mode distributed-execution module, and an MLIR-based deep-learning compiler. In formal evaluations training a 124M-parameter GPT-2 baseline in FP32 on an 8-GPU 6000 Ada system, BluTrain outperforms industry-standard baselines in both throughput (sustaining an average of 407K tokens/s versus PyTorch's 395K tokens/s) and memory efficiency (achieving up to a 22% footprint reduction), while strictly preserving numerical fidelity and converging to a marginally lower final validation loss. With every layer explicitly open to native tuning, the performance ceiling is the framework's own to raise.