Hybrid GPU CPU method speeds up large model training efficiently
Fast and Memory Efficient Offload Training Framework with Hybrid XPU Computation
Distributed, Parallel, and Cluster Computing
Summary
Training very large deep learning models often runs out of GPU memory, slowing down or stopping training. The authors of this paper created a new way to split work between GPU and CPU memories using a technology called Direct Host Access, which lets certain GPU parts use CPU memory directly. Their system, MemFerry, smartly decides which parts of the model go to GPU or CPU to speed up training and reduce memory use. Tests showed that MemFerry can train bigger models faster than previous methods and works well even when using multiple GPUs or special hardware setups.
What this means in practice
- •For machine learning engineers: Train larger deep learning models faster on existing GPU hardware by using hybrid GPU-CPU memory management with MemFerry.
- •For cloud infrastructure operators: Reduce training iteration times on scale-up nodes with multiple NPUs by using ScaleUp-MemFerry to better utilize interconnect bandwidth.
Authors
Zhiyi Yao, Zuning Liang, Yuedong Xu, Jin Zhao, Jessie Hui Wang, Tong Li
Abstract
With the ever-growing size of deep learning models, GPU memory is prone to being insufficient during training. A prominent approach is ZeRO-Offload, which moves the optimizer states to CPU memory and performs parameter update using CPU. However, the deficiencies of ZeRO-Offload include low GPU utilization, imperfect overlapping of communication and computation, and inflexible offloading. In this paper, we leverage Direct Host Access (DHA) on the GPU that can compute data in CPU memory, forming a novel hybrid on-GPU and DHA. We design and implement MemFerry consisting of an execution scheduler and a shadow model. The scheduler strategically chooses layers of parameters for DHA computation and transmits the remaining parameters to GPU memory simultaneously to shorten forward propagation time, and further loads DHA parameters to GPU memory to reduce backward propagation time. The shadow model presents a unified memory abstraction for the parameter partitions stored separately in GPU and CPU memories. To further reduce GPU memory usage, we present MemFerry along with its dynamic programming algorithm that offloads gradients to CPU memory via DHA. We further extend MemFerry to emerging scale-up domains with ScaleUp-MemFerry, which exploits otherwise underutilized accelerator interconnect bandwidth to assist host-to-accelerator data movement through adaptive multi-path transfer. Our experiments show that \system trains up to $1.68\times$ faster and MemFerry can train $1.52\times$ larger model compared to ZeRO-Offload on a single GPU, and increase training speed by at least $28.1\%$ when scaling to data parallelism on 8 GPUs. We further extend the design to a Huawei CloudMatrix384 scale-Up node with up to 8 NPUs, and our ScaleUp-MemFerry reduces the end-to-end iteration time by up to $20.7\%$ over DeepSpeed.