通过多级多路径卸载,突破大模型训练的GPU显存瓶颈
MLP-Offload: Multi-Level, Multi-Path Offloading for LLM Pre-training to Break the GPU Memory Wall
- 将优化器状态分层卸载,避免并发冲突
- 280B参数模型下迭代速度提升2.5倍
- 适合资源受限环境下大模型训练
由于大模型规模增长快于GPU显存容量,训练超出多张GPU总显存的模型日益必要。现有基于主机内存或磁盘的多级异步卸载技术虽先进,但仍存在关键路径上的显著I/O开销,导致迭代变慢。本文提出MLP-Offload,一种专为资源受限环境设计的多级、多路径卸载引擎,旨在缓解训练中反向传播与参数更新阶段的I/O瓶颈。关键观察包括:更新阶段的I/O开销主导迭代时间;三级远程存储的带宽未被充分利用;并发卸载引发的竞争加剧了瓶颈。基于此,MLP-Offload采用缓存高效且并发可控的方式,在多层级间卸载优化器状态。在高达280B参数的模型上评估显示,相比当前最优训练框架,迭代速度提升2.5倍。
原文摘要 · Abstract (English)
Training LLMs larger than the aggregated memory of multiple GPUs is increasingly necessary due to the faster growth of LLM sizes compared to GPU memory. To this end, multi-tier host memory or disk offloading techniques are proposed by state of art. Despite advanced asynchronous multi-tier read/write strategies, such offloading strategies result in significant I/O overheads in the critical path of training, resulting in slower iterations. To this end, we propose MLP-Offload, a novel multi-level, multi-path offloading engine specifically designed for optimizing LLM training on resource-constrained setups by mitigating I/O bottlenecks. We make several key observations that drive the design of MLP-Offload, such as I/O overheads during the update dominate the iteration time; I/O bandwidth of the third-level remote storage tier remains unutilized; and, contention due to concurrent offloading amplifies I/O bottlenecks. Driven by these insights, we design and implement MLP-Offload to offload the optimizer states across multiple tiers in a cache-efficient and concurrency-controlled fashion to mitigate I/O bottlenecks during the backward and update phases. Evaluations on models up to 280B parameters shows that MLP-Offload achieves 2.5$\times$ faster iterations compared to the state-of-the-art LLM training runtimes.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。