DistZO2: High-Throughput and Memory-Efficient ZerothOrder Fine-tuning LLMs with Distributed Parallel

文章主要内容和创新点

主要内容

本文提出了DistZO2，一个高吞吐量、内存高效的分布式零阶（zeroth-order, ZO）微调框架，用于大规模语言模型（LLMs）的微调。零阶优化通过仅使用前向计算估计梯度（无需反向传播），大幅降低内存消耗，但现有ZO2框架受限于单设备执行，吞吐量较低。DistZO2在ZO2的内存高效设计基础上，引入分布式并行策略和通信优化，实现了多GPU系统上的高效零阶微调，在OPT-175B模型上实现了相较于ZO2的3倍加速，同时保持低内存占用（单GPU内存可低于20GB）。

创新点

1. 三种并行策略：

   • 扰动并行（Perturbation Parallelism, PertP）：将零阶优化所需的两个扰动前向传播（+ϵz和-ϵz）分配到不同设备并行执行，利用两者的独立性提升吞吐量。
   • 分布式数据并行（Distributed Data Parallelism, DDP）：适配零阶训练的标量梯度特性，通过同步标量投影梯度而非完整张量，实现高效模型复制和数据并行。
   • 2D并行：融合PertP和DDP，形成“扰动方向-数据批次”的二维并行框架，支持更多GPU的高效扩展。

2. 硬件感知