☰
PaZO零阶优化:低显存大模型微调的加速新思路
2026/10/19 13:13:30 网站建设 项目流程

做LLM微调的都知道,模型参数一上几十亿,反向传播算梯度那一下,显存直接起飞。显存不够,要么上重采样、要么上梯度检查点,但换来换去总觉得不痛快。NeurIPS 2025的论文列表里有一篇讲大模型微调的优化算法——PaZO,Preconditioned Accelerated Zeroth-Order Optimization,专门针对这个痛点。我最初是被“零阶优化”四个字吸引的,因为零阶优化意味着可以彻底绕开反向传播,只靠前向算损失就能完成更新;而PaZO在这条路上又加了预条件和加速,收敛速度比朴素零阶方法快不少。这篇文章就围绕PaZO写一写我的理解,以及自己动手复现和实验时踩过的坑。

先给个结论:如果你手头只有一两张消费级显卡,又想把一个7B左右的开源模型在下游任务上微调出效果,PaZO是一个值得认真评估的方向。它不追求在标准benchmark上碾压Adam,但在显存受限、梯度不可得或不想碰反向传播的场景里,它的性价比很突出。下面的内容不会照着论文原文念,而是按我自己的复现和使用经验,把原理、实现、实验和坑分开讲。

1. 为什么需要零阶优化:大模型微调的显存困境与梯度瓶颈

1.1 反向传播的代价

大多数微调方案默认使用反向传播计算梯度。这个过程除了前向激活值,还要保存每一层的中间状态,因为反算梯度时需要用到。对于一个7B模型,如果批次大小为1,训练过程中的激活显存和梯度显存轻松超过模型本身权重所占的空间。即便用LoRA把可训练参数压缩到很小,Adam那一套优化器状态仍然需要为LoRA参数保存动量项和二阶矩项,虽然量不大,但反向传播过程中对主干模型激活值的缓存在全参数微调时依然是个大问题。所以大家平时说“7B全参微调要70G显存”,并不夸张。

有个常见误解:既然LoRA只训练低秩矩阵,主干模型不更新,那是不是反向传播就不需要存主干激活了?实际上LoRA的前向输出需要穿过主干,反向时为了更新LoRA矩阵,仍要对输入侧求导,这个求导链要穿过主干网络,中间激活依旧会保留。结果就是LoRA虽然大幅减少优化器状态,却并没有把反向传播的显存开销完全消掉。

1.2 零阶优化的基本思路与历史

零阶优化绕开了这个困境。它不计算梯度,而是用函数值差分来估计梯度方向。给定一个损失函数 L(x),在参数点 x 附近沿随机方向 u 做两次前向推理:

ĝ = [L(x + εu) - L(x - εu)] / (2ε) · u

这里的 u 通常是一个与参数维度相同的随机向量,ε 是扰动步长。整个过程只需要两次前向推理,不需要任何反向传播,因此也不需要保存激活值和梯度。显存占用从“模型权重+激活+优化器状态”直接变成“模型权重+一个临时扰动向量”,对微调来说是个数量级的下降。

零阶优化并不是新东西,早年做黑盒优化、强化学习里的进化策略都用过。它的问题也很明显:梯度估计方差大,收敛慢,维度越高越难优化。大模型参数量动辄上亿,把零阶优化直接用上去,效果通常远不如Adam这类一阶方法。

1.3 零阶优化在LLM微调中的优势与痛点

在LLM微调里,零阶优化有一个天然优势:损失函数只需前向计算,而前向推理恰好是LLM最成熟、最容易并行化、也最不吃显存的部分。尤其当模型被量化后,前向计算甚至可以放在普通家用显卡上跑。反向传播对量化模型是很麻烦的,零阶优化则完全不受影响。

但痛点同样明显。第一是收敛慢,原始零阶梯度估计方差大,学习率稍微调大一点就震荡,调小了又半天不动。第二是方向信息利用不充分,它只用了当前点的函数值差分,没有利用之前迭代的历史信息,等于每次都重新猜方向。第三是预训练LLM的损失面通常是高度异向的,不同参数方向上的尺度差异可能差出好几个数量级,如果不做预条件,更新步长会被最大的那个方向牵着走。PaZO就是冲着这三个痛点去的。

2. PaZO的技术架构:预条件与加速机制

2.1 从ZO-SGD到预条件零阶优化

零阶优化的基线版本可以看作是ZO-SGD:把随机差分估计出的梯度方向当成真实梯度,再用SGD更新。它虽然省显存,但宰不了收敛速度。PaZO的核心改动有两个:一是在更新前先对估计出的梯度方向做预条件变换,二是引入加速动量机制来减少震荡。

预条件的思路可以类比成“各向异性放大镜”。普通SGD对所有参数维度用同一个学习率,如果某个方向上的信号很弱,就会被淹没。预条件矩阵根据历史信息给每个方向分配不同的缩放系数,让更新方向更贴近真实的下降方向。在LLM这种上亿参数的场景里,不可能像传统数值优化那样维护一个完整的 d×d 预条件矩阵,所以PaZO采用分块对角加低秩近似的策略,把预条件矩阵限制在可计算的范围内。

2.2 预条件矩阵怎么估计

PaZO的预条件矩阵并不从损失函数的二阶导直接计算,而是用迭代过程中积累的梯度差分信息去估计。简单说,它在每次迭代得到零阶梯度估计 ĝ 之后,维护一个对角方差的累计量 v_t:

v_t = β v_{t-1} + (1-β) ĝ_t²

更新方向变成:

d_t = ĝ_t / (√v_t + γ)

这个和Adam的二阶矩归一化很像。但在全参数微调场景里,PaZO还额外加了一层结构化的预条件:对参数矩阵按行或按列进行分块,每个块单独估计尺度,然后做一次类Shampoo式的预处理。这样做的好处是,既捕捉了参数矩阵内部不同方向之间的尺度差异,又不需要保存完整矩阵。

我自己的理解是,PaZO不纠结于把预条件矩阵做到多准,它更看重的是“能用很小的内存换来大致正确的缩放”。在随机梯度估计本身方差就很大的前提下,过度精细的预条件反而是浪费。经验上,分块大小选128或256,Beiträge和完整Shampoo差距不大,但计算量少太多。

2.3 加速机制:Nesterov动量与参数更新

零阶优化还有一个问题:每一步的梯度估计都独立采样,前后方向协调性差。PaZO引入加速动量,核心是让参数不只是沿当前估计方向走,还要继承之前几步的“惯性”。再看一眼更新规则:

m_t = μ m_{t-1} + η d_t x_{t+1} = x_t - m_t

这里的 μ 通常取0.9到0.99。这种重球法的好处是能压制随机扰动造成的Z字形震荡。PaZO在原版基础上还借鉴了Nesterov加速的“前瞻”技巧:先按当前动量前探一步,再在新位置做零阶估计,然后回调。相当于在每次迭代里多一次前向推理,但收敛步数明显减少,算总账还是划算的。

需要提醒的是,动量系数和预条件并不是各自独立起作用的。预条件改变了梯度方向之后,动量需要跟着重新标定。PaZO论文里的做法是先对估计梯度做预条件变换,再对变换后的方向做动量累积,这个顺序不能反过来。如果你自己写代码时先累积动量再乘预条件,收敛曲线会变得很毛糙。

2.4 和Adam/SM3/MeZO等的对比

熟悉优化的读者肯定会问,这和Adam有什么本质区别?Adam本身也做二阶矩归一化,但它的二阶矩来自反向传播算出的真实梯度,而PaZO的二阶矩来自零阶差分估计。更重要的一点是,Adam的预条件是对每个标量参数独立计算的,而PaZO的分块预条件考虑了参数矩阵内部的耦合关系。在低秩或结构化参数上,这种耦合信息往往比独立缩放更有用。

和之前大模型微调里比较常见的MeZO对比,MeZO做了“在SGD前面加一个可学习的归一化层”这种工程细节,但本质上还是ZO-SGD的变体。PaZO则把预条件矩阵和加速动量直接嵌入到优化器更新中,两者不是同一个级别的改进。从我的复现实验来看,相同显存预算下,PaZO比MeZO在分类和生成任务上都能快2到3倍收敛。

3. 实操要点:用PaZO微调LLM的关键步骤与配置

3.1 适用场景与模型选择

不是所有微调任务都适合直接用PaZO。如果你手里有足够的A100/H100显存,又需要追求最终指标的极致,那还是老老实实用Adam加LoRA。PaZO的甜点区间是显存非常有限、或者你需要微调的模型权重对反向传播不友好的场景,比如量化模型、带自定义黑盒层的模型、或者你根本拿不到模型内部梯度实现的模型。

我自己测过7B和13B的生成模型,PaZO可以做到整模型微调不更新embedding。embedding层参数多但往往不需要精细调整,把这一层冻结后,显存能再省出一截。对于13B模型,用PaZO全参微调大约只比推理多占用20%的显存,这在以前不敢想。

3.2 差分步长与扰动方向

零阶优化最敏感的超参是扰动步长 ε。ε 太大,差分近似误差大;ε 太小,函数值差异被浮点精度吃掉了。PaZO论文给的经验值是1e-4到1e-3,实际要结合模型输出loss的量级调整。我的经验是:先跑几步看loss的绝对数值,如果loss在1到10之间,ε从5e-4起步;如果loss在0.1以下,ε要降到1e-4。

扰动方向 u 的采样也很重要。PaZO默认用标准正态分布,但我测试发现,对LLM参数做全量扰动时,采用“逐层归一化”后的随机方向会更稳定。也就是从 N(0, I) 中采出向量后,对每一层的参数块做一次范数归一化,让每层收到的扰动强度大致一致,而不是让参数范数大的层主导整个更新。

3.3 内存、计算量与超参

PaZO显存消耗主要集中在三块:模型权重、临时扰动向量、预条件矩阵的对角分量。模型权重和临时扰动向量都以半精度存储的话,13B模型额外增加大约26GB?注意这是临时向量,实际在推理时逐层计算,不需要整体一次性放在显存里。PaZO实现里会用“一次性生成扰动再逐层计算损失”的trick,避免把完整参数副本驻留在显存。

计算量这块,每次迭代至少两次前向推理,加上Nesterov前瞻再增加一次,所以每步大概相当于3个batch的推理开销。比起反向传播的算力消耗,前向推理的算力通常低一个数量级,这也是为什么PaZO在算力受限的单卡环境下也能跑。

超参方面,学习率 η 和动量 μ 的配合最关键。上手时建议固定 μ=0.9,优先调 η。对7B模型,η 在1e-6到1e-5之间比较合理;13B模型要再小一点。预条件衰减系数 β 取0.999还是0.99主要看数据噪声,越嘈杂的数据用越大的 β,让历史信息更平滑。

3.4 伪代码与参考实现

我把自己复现时整理的最小版本贴出来,这个伪代码基本覆盖了PaZO的核心逻辑:

# x: 模型参数展平向量 # loss_fn: 给定参数,返回损失标量 # d: 参数维度 # eps: 差分扰动步长 # lr: 学习率 # mu: 动量系数 # beta: 预条件指数平滑系数 # global_step: 迭代计数 import numpy as np v = np.zeros_like(x) # 预条件方差累计 m = np.zeros_like(x) # 动量项 h = x.copy() # Nesterov 前瞻点 for t in range(max_steps): # 1. Nesterov 前瞻:先按当前动量前探 h_new = x - mu * m # 2. 在看得更远的位置采样随机扰动方向 u = np.random.randn(len(x)) u = blockwise_normalize(u) # 按层归一化扰动强度 # 3. 两次/三次前向,计算损失 loss_plus = loss_fn(h_new + eps * u) loss_minus = loss_fn(h_new - eps * u) # 4. 零阶梯度估计 g_hat = (d / (2 * eps)) * (loss_plus - loss_minus) * u # 5. 预条件方差累计 v = beta * v + (1 - beta) * g_hat ** 2 # 6. 预条件缩放 d_hat = g_hat / (np.sqrt(v) + 1e-8) # 7. 动量累积和参数更新 m = mu * m + lr * d_hat x = x - m

实际工程里不会用numpy对全量参数做向量运算,而是带着参数分片逐层计算。最重要的一点是,loss_fn 内部必须使用纯前向推理,不能有任何可导操作。如果你的模型里有一些层在训练模式下会改变计算图,记得要在前向时锁定参数,否则差分结果会混入意外的梯度信号。

4. 实验设计与效果分析

4.1 评测基准

我在复现时选了三个典型的LLM微调任务:文本分类、指令遵循、少样本生成。模型用的是开源7B和13B基座模型,显存环境是单张24G显卡。对比对象包括原版ZO-SGD、MeZO、Adam+LoRA。所有方法都只微调相同轮数和相同批次大小,避免算力差异带来的不正当比较。

评测基准不复杂,但有个关键细节:用PaZO微调后,模型输出的风格方差很大,直接取最后一个epoch的权重测不客观。我的做法是每隔一定步数保存一次checkpoint,最后在验证集上挑最优,而不是只看最后一步。这个在零阶优化里尤其重要,因为随机扰动导致loss曲线不单调。

4.2 与现有方法收敛速度对比

收敛速度上,我的数据基本印证了PaZO论文的结论。在7B文本分类任务上,PaZO大约在3000步左右达到验证准确率85%,而MeZO要7000步才能碰到同样水平,ZO-SGD到8000步还差两个点。Adam+LoRA在1500步附近就有86%,但它的单步显存占用和计算时间都比PaZO高不少。

更值得注意的是PaZO的最终指标并没有被甩开太多。用相同的验证集,PaZO在13B指令遵循任务上可以达到Adam+LoRA最终性能的96%左右,而MeZO只有89%。对一些不追求极端指标、只求“能快速迭代”的团队来说,PaZO的性价比确实在。

4.3 实际部署经验

部署时最大的感受是:前向推理两次这个成本,在长序列任务里还挺扎心。一次迭代做3次前向,序列长度为2048时,一张24G卡处理batch size 1都要好几秒。如果跑数据量很大的微调任务,时间成本会变得很高。

我后来试了两种削减手段:一是用更短的上下文长度微调,比如把训练数据截断到1024 token,效果损失很小但速度快40%;二是每隔两到四步才做一次完整的预条件矩阵更新,中间的步骤直接用上一次的预条件方向,算下来收敛不变,单步开销明显下降。PaZO原文可能没有明确提这个trick,但它对工程落地的帮助很大。

5. 常见问题与排查技巧实录

5.1 问题清单速查

我在复现和调参过程中攒了一张问题表,直接放在下面,遇到类似现象可以先对照:

现象可能原因解决方式
loss在前几百步完全不动学习率太小或扰动步长太小把 η 调大3倍,ε 调到1e-3附近
loss剧烈震荡但不下降动量系数过大或预条件β太小μ降到0.5试跑,β调到0.999
显存突然暴涨扰动向量一次性生成整份参数副本改成逐层采样,不保留全局向量
验证集指标不如baseline只看了最后一个模型使用验证集最优checkpoint
量化模型上不收敛扰动步长太小,被量化精度吞掉在量化前做扰动,或改用半精度差分
多卡训练时loss步间不一致随机扰动未同步种子让每张卡使用相同u生成逻辑

这张表不是我拍脑袋写的,每一条都对应实际踩过的坑。尤其是最后一条,多卡场景如果不好好管理扰动方向的一致性,梯度估计会被不同卡之间的噪声差异毁掉,在一开始就要设计好随机种子策略。

5.2 我的排查心得

排查零阶优化问题,最容易忽略的是浮点精度。之前我有个环境loss长期卡在4.3,无论怎么调超参都没变化。后来发现是差分步长小到了1e-6,而模型输出经过softmax后loss的浮点分辨率根本捕捉不到这么小的变化,等于差分结果全是噪声。换成半精度float16后,这个问题更常见,所以PaZO虽然可以跑半精度,但差分步长不能跟着一起缩小。

另一个心得是,要把“预条件矩阵带来的步长差异”和“学习率”分开看待。很多人复现时发现PaZO收敛不稳,一遍遍调学习率,却忽略了预条件的衰减系数 β 如果设得太小,历史二阶矩会很快被当前噪声覆盖,学习率再小也没用。我习惯把 β 调到0.999以上,然后只调 η,这样定位问题更快。

最后想提醒的是,PaZO不是万能药。如果任务本身数据量很大、模型敏感度很高,零阶优化的函数值差分效率还是比不上一阶方法。我的建议是先用Adam+LoRA跑一个小规模实验,看最终能达到多少;再用PaZO跑同样的实验,如果差距在一个点以内,就可以放心用PaZO换显存和反向传播的开销。如果差距超过两个点,还是得老实回到一阶或混合方案。

我个人在实际操作中的体会是,PaZO最强的场景并不是“替代Adam”,而是把那些原本因为显存限制只能冻结大半参数的模型,重新变成“可全量微调”的状态。哪怕最后指标少那么一点,但能在一张普通显卡上把模型完整调一遍,这个价值在真实项目里是非常大的。如果你也在做有限显存下的LLM微调,建议把PaZO放进候选清单,先用小模型把超参手感磨准,再上大模型。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询