arXiv:2503.12668cs.LGcs.PF2025-03被引 6

用零阶优化+动态参数迁移,在18GB显存上搞定1750亿参数大模型微调。

ZO2: Scalable Zeroth-Order Fine-Tuning for Extremely Large Language Models with Limited GPU Memory

  • 通过双前向计算和参数在CPU/GPU间动态切换,避免存储梯度与激活值。
  • 在18GB显存下成功微调OPT-175B模型,精度不降、速度几乎无损。
  • 适合资源受限却想训练超大规模语言模型的研究者或开发者。

微调大型预训练语言模型通常需要大量显存。传统一阶优化器如SGD因模型规模扩大,正向和反向传播中需存储大量激活值和梯度,显存需求剧增。零阶(ZO)方法仅通过前向操作即可估算梯度,无需保存中间激活值。结合利用CPU能力,可显著提升单个GPU的内存与算力。本文提出全新框架ZO2(Zeroth-Order Offloading),实现仅用有限显存高效进行零阶微调。该框架按需动态在CPU与GPU间迁移模型参数,优化计算流程,最大限度减少等待时间。将参数调整与零阶双前向操作结合,减少不必要的数据传输,提升微调效率。此外,框架支持在AMP模式下使用创新的低比特精度策略,进一步优化CPU与GPU间的数据交换。该方法使我们能在仅18GB显存的设备上完成如OPT-175B(超过1750亿参数)这类超大规模模型的微调,成果远超传统方法。同时,该方案几乎无额外时间开销,且与标准零阶方法相比精度完全无损失。ZO2代码已开源:https://github.com/liangyuwang/zo2。

原文摘要 · Abstract (English)

Fine-tuning large pre-trained LLMs generally demands extensive GPU memory. Traditional first-order optimizers like SGD encounter substantial difficulties due to increased memory requirements from storing activations and gradients during both the forward and backward phases as the model size expands. Alternatively, zeroth-order (ZO) techniques can compute gradients using just forward operations, eliminating the need to store activations. Furthermore, by leveraging CPU capabilities, it's feasible to enhance both the memory and processing power available to a single GPU. We propose a novel framework, ZO2 (Zeroth-Order Offloading), for efficient zeroth-order fine-tuning of LLMs with only limited GPU memory. Our framework dynamically shifts model parameters between the CPU and GPU as required, optimizing computation flow and maximizing GPU usage by minimizing downtime. This integration of parameter adjustments with ZO's double forward operations reduces unnecessary data movement, enhancing the fine-tuning efficacy. Additionally, our framework supports an innovative low-bit precision approach in AMP mode to streamline data exchanges between the CPU and GPU. Employing this approach allows us to fine-tune extraordinarily large models, such as the OPT-175B with more than 175 billion parameters, on a mere 18GB GPU--achievements beyond the reach of traditional methods. Moreover, our framework achieves these results with almost no additional time overhead and absolutely no accuracy loss compared to standard zeroth-order methods. ZO2's code has been open-sourced in https://github.com/liangyuwang/zo2.

大模型微调零阶优化显存优化分布式训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。