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

The paper introduces MegaTrain, a technique that enables training large language models with over 100 billion parameters on a single GPU by streaming parameters from CPU memory and overlapping computation and data transfer.

Details

The core problem addressed is that training large language models (LLMs) with over 100 billion parameters typically requires distributing the model across dozens or hundreds of GPUs, as the parameters, gradients, and optimizer states do not fit in the VRAM of a single GPU. MegaTrain tackles this by using a memory-time tradeoff - instead of keeping all parameters active in VRAM, it streams the parameters from CPU memory or storage to the GPU as needed for the forward and backward passes, and discards them afterwards. This is done at the granularity of individual layers, with intelligent prefetching to ensure the GPU never waits. Crucially, MegaTrain maintains full precision (FP32 or BF16) for the active parameters, unlike techniques like LoRA that reduce precision to fit in memory. The optimizer states are also stored on the CPU and synchronized as needed. The aggressive overlapping of data transfer and computation is a key aspect of the approach.