Multi-node LLM training scales down GPUs using CPU memory extension
TRANSIT: Transparent Scale-in for Multi-Node LLM Training
Distributed, Parallel, and Cluster ComputingHardware Architecture
Summary
Training large language models usually needs many graphics cards (GPUs), which can be expensive and hard to manage. The authors present TRANSIT, a system that lets training run efficiently using fewer GPUs by temporarily using regular computer memory (CPU DRAM) as extra space. Importantly, TRANSIT works without changing existing software or hardware setups, and moves data quickly between CPU and GPU. Their tests show that TRANSIT keeps most of the training speed while reducing hardware needs and network traffic.
What this means in practice
- •For ai infrastructure engineers: Enable model training with fewer GPUs by using CPU memory smartly without altering existing cluster setups.
- •For high performance computing teams: Reduce network load and improve GPU throughput for large-scale distributed training jobs across data center networks.
Authors
Hyungyo Kim, Nicholas Satchanov, Hrishi Shah, Gaohan Ye, Jiaqi Lou, Robert Walkup, Shweta Salaria, I-Hsin Chung, Hubertus Franke, Seetharami Seelam, Apoorve Mohan, Nam Sung Kim
Abstract
TRANSIT is a transparent scale-in framework to enable multi-node model training on fewer GPUs while maintaining training efficiency by transparently leveraging CPU DRAM as an extension of GPU memory during distributed training. It achieves this through a user-space interposition layer, requiring no modifications to the application, training framework, cluster scheduler, device driver, or operating system. Furthermore, TRANSIT achieves higher efficiency by leveraging a zero-copy data path for CPU-GPU transfers. We evaluate TRANSIT on dense and MoE models across scales up to 64 NVIDIA H100 GPUs and multiple parallelism configurations over a RoCE network. Our evaluation shows that TRANSIT can: (a) outperform state-of-the-art framework-managed offloading techniques, achieving up to 68%, 59%, and 42% higher per-GPU throughput than TorchTitan, ZeRO-Offload, and ZeRO-Infinity, respectively, (b) enables training with 50% fewer GPUs while maintaining over 90% of baseline per-GPU throughput, (c) lower per-node network traffic by up to 33%, and (d) improve per-GPU throughput by up to 35% in communication-bound settings.