ZO2：面向极大规模语言模型的可扩展零阶细调，GPU 内存受限
机器学习
2025-03-18 v1 性能
摘要
微调大型预训练语言模型（LLM）通常需要大量 GPU 内存。随着模型规模扩大，传统的一阶优化器如 SGD 在存储前向和反向过程中的激活值和梯度方面会面临巨大的内存需求。或者，可以使用零阶（ZO）技术仅通过前向操作计算梯度，从而无需存储激活值。此外，利用 CPU 能力，是一种可行的方法，可提高单个 GPU 可用的内存和计算资源。我们提出了一种新框架，ZO2（Zeroth-Order Offloading），用于在仅受限 GPU 内存的情况下高效零阶细调 LLM。该框架动态地在 CPU 和 GPU 之间移动模型参数，以满足需求，优化计算流程并通过最小化停机时间最大化 GPU 使用率。这一零阶双前向操作与参数调整的集成减少了不必要的数据移动，从而提高了细调效果。此外，本框架支持一种创新的低位精度方法在 AMP 模式下简化 CPU 和 GPU 之间的数据交换。通过这种方法，我们能够在仅 18GB GPU 上对参数超过 1750 亿的极大模型（如 OPT-175B）进行微调——这超出了传统方法的 reach。此外，本框架几乎没有额外的时间开销，同时在准确率上几乎没有损失，与标准零阶方法相比表现一致。ZO2 的代码已开源于 https://github.com/liangyuwang/zo2。
引用
@article{arxiv.2503.12668,
title = {ZO2: Scalable Zeroth-Order Fine-Tuning for Extremely Large Language Models with Limited GPU Memory},
author = {Liangyu Wang and Jie Ren and Hang Xu and Junxiao Wang and Huanyi Xie and David E. Keyes and Di Wang},
journal= {arXiv preprint arXiv:2503.12668},
year = {2025}
}
备注
14 pages, 7 figures