MegaTrain: Efficient Large Language Model Training on a Single GPU

MegaTrain revolutionizes the training of large language models by leveraging a memory-centric approach, allowing the training of models with over 100 billion parameters on a single GPU. Unlike traditional systems that rely heavily on GPU memory, MegaTrain stores parameters and optimizer states in host memory, treating GPUs as transient compute engines. This design minimizes persistent device state by streaming parameters in and computing gradients out for each layer, effectively decoupling model scale from GPU memory capacity.

To tackle the CPU-GPU bandwidth bottleneck, MegaTrain employs two key optimizations. First, a pipelined double-buffered execution engine overlaps parameter prefetching, computation, and gradient offloading across multiple CUDA streams, enabling continuous GPU execution. Second, it replaces persistent autograd graphs with stateless layer templates, dynamically binding weights as they stream in, which eliminates the need for storing large graph metadata and persistent intermediate tensors.

MegaTrain's innovative approach allows it to achieve 1.84 times the training throughput of existing systems like DeepSpeed ZeRO-3 when training 14 billion parameter models. It also supports ultra-long context training up to 512,000 tokens on a single GPU, a feat unattainable by current offloading-based systems. This system's design suggests that the future of large model training lies in efficient memory and compute organization rather than sheer GPU capacity.

The potential for scaling MegaTrain to multiple GPUs with tensor or expert parallelism is a promising avenue for future research. Incorporating tiered storage with SSDs could further extend its capabilities, possibly making trillion-parameter training feasible on everyday systems.