Efficient Tensor Offloading for Large Deep-Learning Model Training based on Compute Express Link
Dong Xu, Yuan Feng, Kwang-Sik Shin, Dae Woo Kim, Hyeran Jeon, Dong Li
Abstract
The deep learning models (DL) are becoming bigger, easily beyond the memory capacity of a single accelerator. The recent progress in large DL training utilizes CPU memory as an extension of accelerator memory and offloads tensors to CPU memory to save accelerator memory. This solution transfers tensors between the two memories, creating a major performance bottleneck. We identify two problems during tensor transfers: (1) the coarse-grained tensor transfer creating difficulty in hiding transfer overhead, and (2) the redundant transfer that unnecessarily migrates value-unchanged bytes from CPU to accelerator. We introduce a cache coherence interconnect based on Compute Express Link (CXL) to build a cache coherence domain between CPU memory and accelerator memory. By slightly extending CXL to support an update cache-coherence protocol and avoiding unnecessary data transfers, we reduce training time by $33.7 \%$ (up to $55.4 \%$) without changing model convergence and accuracy, compared with the state-of-the-art work in DeepSpeed [62].