Offload-DMSS: Heterogeneous Training with Limited CPU Memory Bandwidth by Direct Model State Swap

In recent years, large language models (LLMs) have demonstrated extraordinary performance and have become the brightest stars on the artificial intelligence stage. However, training these models requires a significant amount of GPU memory, which discourages researchers and impedes progress in deep learning. Fortunately, heterogeneous training provides a promising way. The current state-of-the-art (SOTA) updates parameters on CPU side, which can be a serious bottleneck, especially in scenarios with limited CPU memory bandwidth. To address this issue, in this paper, we introduce Offload-DMSS, a novel heterogeneous training strategy based on Direct Model State Swap (DMSS). In Offload-DMSS, we overlap the swapping of model states with backpropagation; consequently, it achieves higher training performance than the current SOTA while remaining insensitive to CPU-memory bandwidth. In our experiments, its training throughput surpasses the current SOTA by 1.4 \(\times \) on average and by up to 1.6 \(\times \) in the best case.