通过动态调度优化器状态,实现大模型训练中GPU与CPU内存的高效协同。
Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading
- 按迭代周期动态在主机与GPU间移动部分优化器状态,利用计算间隙提升内存利用率。
- 在相同硬件下,相比现有方法迭代速度提升2.5倍,有效突破内存瓶颈。
- 适合大规模语言模型训练场景,尤其适用于资源受限但需高吞吐的部署环境。
Transformer与大型语言模型(LLMs)已在各领域快速普及,其参数规模已达数百亿并持续增长。在此背景下,训练过程成本高昂,常遭遇“内存墙”问题:即使采用3D并行(流水线、张量、数据)并聚合多块GPU内存,仍不足以容纳模型参数、优化器状态、梯度和激活等数据结构。为缓解此问题,当前主流方法将部分优化器状态卸载至主机内存,并执行混合CPU-GPU计算。然而,主机与GPU内存的管理往往不够优化,导致数据传输与计算重叠不足,错失同时利用互连带宽与算力的机会。本文观察到前向、反向与更新阶段的交错引发GPU内存使用波动,可借此在每轮迭代中动态调度部分优化器状态在主机与GPU内存之间移动。为此,我们提出Deep Optimizer States,将大模型分组,基于所提出的性能模型,在每轮中决策将更新操作调度至CPU或GPU,权衡数据移动开销、加速效果及共享资源竞争。我们将该方法集成至DeepSpeed,在大量实验中验证其较现有最优方法实现2.5倍的迭代加速。
原文摘要 · Abstract (English)
Transformers and large language models~(LLMs) have seen rapid adoption in all domains. Their sizes have exploded to hundreds of billions of parameters and keep increasing. Under these circumstances, the training of transformers is very expensive and often hits a ``memory wall'', i.e., even when using 3D parallelism (pipeline, tensor, data) and aggregating the memory of many GPUs, it is still not enough to hold the necessary data structures (model parameters, optimizer state, gradients, activations) in GPU memory. To compensate, state-of-the-art approaches offload the optimizer state, at least partially, to the host memory and perform hybrid CPU-GPU computations. However, the management of the combined host-GPU memory is often suboptimal and results in poor overlapping between data movements and computations. This leads to missed opportunities to simultaneously leverage the interconnect bandwidth and computational capabilities of CPUs and GPUs. In this paper, we leverage a key observation that the interleaving of the forward, backward, and update phases generates fluctuations in the GPU memory utilization, which can be exploited to dynamically move a part of the optimizer state between the host and the GPU memory at each iteration. To this end, we design and implement Deep Optimizer States, a novel technique to split the LLM into subgroups, whose update phase is scheduled on either the CPU or the GPU based on our proposed performance model that addresses the trade-off between data movement cost, acceleration on the GPUs vs the CPUs, and competition for shared resources. We integrate our approach with DeepSpeed and demonstrate 2.5$\times$ faster iterations over state-of-the-art approaches using extensive experiments.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。