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

Key points

• Problem: GPU memory is the main bottleneck for training large models. Mixed-precision training of a 7B model requires ~84GB just for parameters, gradients, and Adam optimizer states (~12x parameter count), exceeding even an H200's 141GB HBM for larger models.
• Core idea: Break the constraint of GPU memory by leveraging the memory hierarchy. MegaTrain stores all persistent states (parameters, gradients, optimizer states) in host RAM and treats the GPU as a temporary compute engine.

Design highlights

1. Parameters and optimizer states in host memory: Parameters are streamed layer-by-layer to the GPU, gradients streamed back, and optimizer steps (e.g., Adam) run entirely on CPU. GPU memory demand drops from ~12x parameter size to roughly a single layer's worth. 2. Streaming, pipelined execution: Layer-by-layer parameter prefetch and gradient offload, with double buffering (ping-pong) so the next layer's parameters arrive while the current layer computes, overlapping PCIe I/O with GPU computation across multiple CUDA streams. 3. Aggressive activation recomputation: Activations are recomputed every few layers instead of being stored, trading extra compute for memory; segmented checkpointing (every K layers) balances the trade-off. 4. Stateless layer templates: Instead of PyTorch's global autograd graph, each layer uses a lightweight template with dynamically bound parameters, so GPU memory never exceeds one layer and gradients can stream back to the host easily.

Reported performance

• Scale: Reliably trained a 120B-parameter model on a single H200 with 1.5TB host memory; existing offloading systems typically fail around 30B due to exploding host memory needs. Host memory grows linearly with parameter count thanks to flat tensor layouts and an authoritative host storage design.
• Throughput: 1.84x DeepSpeed ZeRO-3 (CPU offload) on a 14B model; sustained over 250 TFLOPS on a single GH200 at 32B, where ZeRO-3 cannot run due to OOM.
• Long context: Trained a 7B model with 512K-token context on a single GH200 — a setup that would require 64 GPUs with FSDP+CP.

Positioning vs. existing frameworks

• DeepSpeed ZeRO-Offload/Infinity: GPU-centric with selective offloading; parameters still mostly reside on GPUs, causing scaling walls for very large models. MegaTrain moves parameters to host fundamentally.
• Megatron-LM: Multi-GPU tensor/pipeline parallelism requiring large clusters; pursues maximum MFU. MegaTrain targets single-GPU scenarios via storage restructuring rather than parallel scaling.
• PyTorch FSDP2: Shards state across multiple GPUs but cannot handle single-GPU huge models; MegaTrain is like FSDP extended to offload-to-host on a single card, with a custom streaming engine.
• Colossal-AI: Ease-of-use distributed framework that still requires multi-GPU environments; MegaTrain optimizes a single GPU to the extreme.

Takeaway

MegaTrain demonstrates that when hardware is scarce, system-level redesign of the memory hierarchy can substitute for hardware stacking. It complements rather than replaces cluster-scale frameworks: single-machine training/fine-tuning with MegaTrain, cluster scaling with tools like Megatron-LM. For memory-constrained practitioners, it makes single-card training of frontier-scale models feasible — albeit with trade-offs in communication-bound workloads and extra recomputation overhead.