arXiv:2507.03211cs.LGcs.PF2025-07ACL被引 2

让百亿参数大模型零阶微调更快更省显存,支持多卡并行。

DistZO2: High-Throughput and Memory-Efficient Zeroth-Order Fine-tuning LLMs with Distributed Parallel Computing

  • 用多设备并行处理扰动前向计算,提升效率。
  • 在OPT-175B上实现比ZO2快3倍的训练速度。
  • 适合资源受限下大规模模型微调的研究者使用。

大语言模型微调因规模庞大而资源消耗巨大。零阶(ZO)优化通过省去反向传播实现内存高效,但其在数百亿参数模型上的应用受限于显存与计算吞吐。现有框架ZO2通过将模型参数移至CPU内存并重叠Transformer块传输与双前向计算缓解内存瓶颈,但仍受限于单设备执行,吞吐量有限。本文提出DistZO2,一种面向分布式零阶微调的高吞吐、内存高效框架。DistZO2引入三种并行策略:(1)扰动并行(PertP),跨设备并行两个扰动前向;(2)分布式数据并行(DDP),适配零阶训练的标量梯度特性;(3)统一2D并行设计,融合PertP与DDP。为缓解参数卸载带来的通信瓶颈,提出基于硬件感知的通信策略,通过高速互连如NVLink对参数块分片并重新分配至多GPU。DistZO2在现代多卡系统上实现零阶微调的可扩展性,在保持ZO2内存效率的同时显著提升训练吞吐。在OPT-175B上的实验表明,相比单机版ZO2,DistZO2实现3倍加速。代码已开源:https://github.com/liangyuwang/zo2。

原文摘要 · Abstract (English)

Fine-tuning large language models (LLMs) remains resource-intensive due to their sheer scale. While zeroth-order (ZO) optimization provides a memory-efficient alternative by eliminating backward passes, its application to multi-hundred-billion-parameter models is constrained by GPU memory and compute throughput. The ZO2 framework addresses the memory bottleneck by offloading model parameters to CPU memory and overlapping transformer block transfer with dual forward computation on a single GPU. However, ZO2 remains limited by its single-device execution and achieves modest throughput. In this work, we present DistZO2, a high-throughput, memory-efficient framework for distributed zeroth-order fine-tuning of LLMs. DistZO2 introduces three parallel strategies: (1) Perturbation Parallelism (PertP), which parallelizes the two perturbed forward passes across devices; (2) Distributed Data Parallelism (DDP), adapted to the scalar-gradient nature of ZO training; and (3) a unified 2D Parallelism design that combines PertP and DDP. To further mitigate communication bottlenecks introduced by parameter offloading, we propose a hardware-aware communication strategy that slices parameter blocks and redistributes them across GPUs via high-speed interconnects such as NVLink. DistZO2 scales zeroth-order fine-tuning to modern multi-GPU systems, preserving ZO2's memory efficiency while substantially improving training throughput. In our experiments on OPT-175B, DistZO2 achieves a 3x speedup over ZO2 with distributed computing. DistZO2'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 官方产品;中文卡片由大模型生成,请以原文为准。