MegaTrain: Training 100B+ Parameter Models on a Single GPU

Introduction: Breaking GPU Memory Barriers

For most AI researchers, GPU memory is the bottleneck limiting model size. Training a 7B-parameter model with mixed precision (bf16 weights + gradients, fp32 optimizer states) requires roughly 12x the parameter count in memory — about 84GB — for parameters, gradients, and Adam momentum/variance. Even a top-end H200 GPU (141GB HBM) cannot hold much larger models, and 14B, 30B, or 100B+ models are out of reach. Traditional approaches rely on model parallelism or data parallelism across many GPUs, which brings high hardware costs and communication overhead.

The Memory Hierarchy: An Overlooked Opportunity

MegaTrain draws on the classic computer architecture concept of the memory hierarchy:

• On-chip SRAM: a few MB, tens of TB/s bandwidth
• GPU memory (HBM/GDDR): e.g., 141GB HBM on H200, ~4.8TB/s bandwidth
• Host memory (CPU DDR/LPDDR): up to multiple TB, hundreds of GB/s bandwidth
• NVMe SSD: tens of TB, tens of GB/s bandwidth

Traditional training systems keep all model states (parameters, gradients, optimizer states) and activations resident in GPU memory, using only one tier of the hierarchy. MegaTrain's core idea is to move persistent states — parameters and optimizer states — to host memory, treating the GPU as a temporary compute engine or high-level cache.

MegaTrain's Memory-Centric Design

1. Parameters and optimizer states live in host memory

Parameters, gradients, and optimizer states (e.g., Adam momentum and variance) are all stored in host memory. The GPU streams parameters in on demand, writes gradients back, and immediately frees its buffers. Optimizer steps execute entirely on the CPU. GPU memory demand drops from ~12x the parameter count to roughly one layer's worth of parameters.

2. Layer-wise streaming of parameters and gradients

During the forward pass, parameters are prefetched layer by layer and released immediately after use. During the backward pass, gradients are computed layer by layer and streamed back to host memory rather than accumulated on the GPU. A double-buffering (ping-pong) scheme pre-loads the next layer's parameters while the current layer computes, overlapping communication and computation.

3. Aggressive activation recomputation

MegaTrain recomputes activations every few layers instead of storing them, trading extra compute for drastically lower memory use. Segmented checkpointing (e.g., every K layers) balances compute against storage.

4. Pipelined double-buffered execution and stateless layer templates

Parameter prefetch, computation, and gradient offload run as pipeline stages across different CUDA streams. MegaTrain abandons PyTorch's autograd graph in favor of stateless layer templates: each layer binds parameters dynamically at runtime and releases them after computing. This eliminates global graph metadata and intermediate tensors, keeping GPU memory usage bounded by a single layer's size.

Performance and Capabilities

• Very large models: reliably trained a 120B-parameter model on a single H200 with 1.5TB host memory. Existing offloading-based systems typically fail around 30B due to exploding host memory requirements; MegaTrain's flat tensor layout and authoritative host storage keep host memory linear in parameter count.
• High throughput: achieved 1.84x the training throughput of DeepSpeed ZeRO-3 (with CPU Offload) on a 14B model. At 32B parameters, MegaTrain sustains over 250 TFLOPS on a single GH200, where ZeRO-3 cannot run due to OOM.
• Ultra-long context: as a byproduct of its layer-wise memory design, MegaTrain trained a 7B model with 512K-token context on a single GH200. Conventional FSDP+CP setups require at least 64 GPUs for 512K context.

Comparison with Existing Frameworks

• DeepSpeed ZeRO-Offload/Infinity: shards and offloads some states to host memory, but keeps parameters primarily on GPU ("GPU-centric, host-assisted"). DeepSpeed needs proportionally more GPUs as model size grows and incurs frequent CPU-GPU synchronization. MegaTrain fundamentally restructures storage by keeping parameters on the host.
• Megatron-LM: NVIDIA's model-parallel framework (tensor + pipeline parallelism) targeting GPU clusters with hundreds of devices and MFU above 50%. MegaTrain instead targets single-GPU scenarios, democratizing access rather than maximizing cluster performance. The two are complementary.
• PyTorch FSDP2: fully shards states across multiple GPUs but cannot handle single-GPU large models. MegaTrain can be seen as extending the sharding idea to host offloading, with a custom streaming engine that minimizes GPU memory further.
• Colossal-AI: an easy-to-use distributed framework with multiple parallelism strategies, still requiring multi-GPU environments. MegaTrain focuses on extreme single-GPU optimization.

Conclusion

MegaTrain demonstrates that by rethinking the storage hierarchy, a single GPU can train 100B+ parameter models — a significant shift for researchers without large GPU clusters. It is not a panacea: it may underperform multi-GPU parallelism on communication-heavy workloads, and aggressive recomputation costs extra compute when GPUs are plentiful. But for memory-constrained scenarios, MegaTrain opens unprecedented possibilities, and future systems may combine such storage-level optimizations within nodes with Megatron-style parallelism across nodes.