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

📄 arXiv: 2507.03211v1 📥 PDF

作者: Liangyu Wang, Huanyi Xie, Di Wang

分类: cs.LG, cs.PF

发布日期: 2025-07-03

🔗 代码/项目: GITHUB

💡 一句话要点

DistZO2：通过分布式并行计算实现高吞吐、内存高效的Zeroth-Order微调LLM

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 Zeroth-Order优化 分布式训练 并行计算 模型微调 内存效率 OPT-175B

📋 核心要点

  1. 现有LLM微调方法资源消耗巨大，Zeroth-Order优化虽然内存高效，但受限于单GPU的计算能力和内存。
  2. DistZO2通过扰动并行、分布式数据并行和统一的2D并行策略，实现了LLM在多GPU上的高效Zeroth-Order微调。
  3. 实验表明，DistZO2在OPT-175B模型上实现了3倍于ZO2的加速，显著提升了训练吞吐量。

📝 摘要（中文）

由于大型语言模型（LLMs）的规模庞大，微调仍然是资源密集型的。Zeroth-order（ZO）优化通过消除反向传播提供了一种内存高效的替代方案，但其在数千亿参数模型上的应用受到GPU内存和计算吞吐量的限制。ZO2框架通过将模型参数卸载到CPU内存，并在单个GPU上重叠transformer块传输与双重前向计算，从而解决了内存瓶颈。然而，ZO2仍然受到单设备执行的限制，并且吞吐量适中。本文提出了DistZO2，这是一个用于LLM分布式Zeroth-Order微调的高吞吐量、内存高效的框架。DistZO2引入了三种并行策略：（1）扰动并行（PertP），它跨设备并行化两个扰动的前向传播；（2）分布式数据并行（DDP），适用于ZO训练的标量梯度特性；（3）统一的2D并行设计，结合了PertP和DDP。为了进一步缓解参数卸载带来的通信瓶颈，我们提出了一种硬件感知的通信策略，该策略对参数块进行切片，并通过高速互连（如NVLink）在GPU之间重新分配它们。DistZO2将Zeroth-Order微调扩展到现代多GPU系统，在保持ZO2内存效率的同时，显著提高了训练吞吐量。在OPT-175B上的实验表明，DistZO2通过分布式计算实现了比ZO2快3倍的速度。DistZO2的代码已在https://github.com/liangyuwang/zo2上开源。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）微调过程中资源消耗过大的问题，特别是当模型参数规模达到数千亿级别时，传统的微调方法需要大量的GPU内存和计算资源。Zeroth-Order优化虽然可以降低内存需求，但现有方法（如ZO2）在单GPU上的执行效率较低，无法充分利用现代多GPU系统的计算能力。

核心思路：DistZO2的核心思路是通过分布式并行计算来加速Zeroth-Order优化过程，同时保持其内存效率。具体来说，它将计算任务分解到多个GPU上，并设计了多种并行策略来最大化GPU的利用率和减少通信开销。通过硬件感知的通信策略，优化了参数在GPU之间的传输，进一步提升了整体性能。

技术框架：DistZO2的整体框架包括以下几个主要部分：1) 扰动并行（PertP）：将两个扰动的前向传播过程分配到不同的GPU上并行执行。2) 分布式数据并行（DDP）：针对Zeroth-Order优化的特点，改进了传统的DDP算法，使其更适合标量梯度的计算。3) 2D并行：结合了PertP和DDP，进一步提升并行度。4) 硬件感知通信：根据GPU之间的互连拓扑（如NVLink），优化参数的传输方式。

关键创新：DistZO2的关键创新在于其并行策略和硬件感知通信策略的结合。传统的并行方法可能无法充分利用Zeroth-Order优化的特点，而DistZO2通过专门设计的并行策略，实现了更高的GPU利用率。此外，硬件感知通信策略能够根据GPU之间的互连拓扑，优化参数的传输，减少通信开销，从而进一步提升整体性能。与现有方法的本质区别在于，DistZO2不仅考虑了计算的并行化，还考虑了通信的优化，从而实现了更高的效率。

关键设计：DistZO2的关键设计包括：1) 扰动并行策略：如何有效地将两个扰动的前向传播分配到不同的GPU上，并保证计算结果的正确性。2) 分布式数据并行策略：如何针对Zeroth-Order优化的特点，改进传统的DDP算法，使其更适合标量梯度的计算。3) 硬件感知通信策略：如何根据GPU之间的互连拓扑，优化参数的传输方式，减少通信开销。具体的参数设置、损失函数和网络结构等技术细节与ZO2保持一致，DistZO2主要关注的是如何通过并行计算来加速训练过程。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DistZO2在OPT-175B模型上的实验结果表明，通过分布式计算，其训练速度比ZO2提高了3倍。这一显著的加速效果证明了DistZO2在提高Zeroth-Order微调LLM吞吐量方面的有效性。此外，DistZO2保持了ZO2的内存效率，使得在资源受限的环境下微调大型模型成为可能。

🎯 应用场景

DistZO2可应用于各种需要微调大型语言模型的场景，例如自然语言处理、文本生成、机器翻译等。它降低了微调LLM的硬件门槛，使得更多研究者和开发者能够在资源有限的条件下进行LLM的定制化训练。未来，DistZO2可以进一步扩展到其他类型的模型和任务，并与其他优化技术相结合，以实现更高的性能。

📄 摘要（原文）

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.