ZenFlow通过异步更新关键参数，显著减少大模型训练时的GPU等待时间。
ZenFlow: Enabling Stall-Free Offloading Training via Asynchronous Updates
- 只在GPU上即时更新重要参数，其余参数异步由CPU处理
- 实测提升5倍训练速度，GPU空等时间减少85%以上
- 适合大规模语言模型训练，尤其对显存不足场景优化明显
微调大型语言模型（LLMs）常超出GPU内存容量，现有系统如ZeRO-Offload将所有参数统一处理，全量在CPU上更新，导致严重GPU停顿：高速昂贵的GPU因等待慢速CPU更新和低带宽PCIe传输而闲置。我们提出ZenFlow，一种新型卸载框架，优先处理重要参数，并分离GPU与CPU的更新流程。ZenFlow在GPU上执行重要梯度的就地更新，同时异步将不重要参数卸载至CPU累积，实现CPU工作与GPU计算完全重叠。为跨多GPU扩展，ZenFlow引入轻量级梯度选择方法，利用重要梯度的新颖时空局部性特性，避免耗时的全局同步。实验表明，ZenFlow实现最高5倍端到端加速、PCIe流量降低2倍，且保持精度不变。
原文摘要 · Abstract (English)
Fine-tuning large language models (LLMs) often exceeds GPU memory limits, prompting systems to offload model states to CPU memory. However, existing offloaded training frameworks like ZeRO-Offload treat all parameters equally and update the full model on the CPU, causing severe GPU stalls, where fast, expensive GPUs sit idle waiting for slow CPU updates and limited-bandwidth PCIe transfers. We present ZenFlow, a new offloading framework that prioritizes important parameters and decouples updates between GPU and CPU. ZenFlow performs in-place updates of important gradients on GPU, while asynchronously offloading and accumulating less important ones on CPU, fully overlapping CPU work with GPU computation. To scale across GPUs, ZenFlow introduces a lightweight gradient selection method that exploits a novel spatial and temporal locality property of important gradients, avoiding costly global synchronization. ZenFlow achieves up to 5x end-to-end speedup, 2x lower PCIe traffic, and reduces GPU stalls by over 85 percent, all while preserving accuracy.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。