ZeroLock 阅读笔记:用模块化解耦打破流水线并行训练的更新锁定

笔记日期: 2026-08-13 作者: Zhongzhu Zhou 阅读论文: ZeroLock: Concurrent Memory-Efficient LLM Training via Modular Update Decoupling 论文作者: Wentao Dai, Xuanran Li, Yuxiang Zhang, Ming Tang, Chao Huang(南方科技大学;蒙特克莱尔州立大学) arXiv: 2608.07974 发表状态: 预印本(cs.LG),2026年8月

1. 流水线并行微调为什么还是有”气泡”问题

如果你曾经试过在几块边缘 GPU 上微调一个几十亿参数的大模型,或者更极端一点,在一堆安卓手机组成的集群上做这件事,你一定撞上过同一堵墙:模型放不下一张卡,所以必须切分;而几乎所有保留标准反向传播(BP)完整链路的切分方式,最终都会留下要么是空闲算力(“气泡”),要么是被浪费的显存(激活值不得不比你想要的时间保留得更久)。ZeroLock 是一篇系统 + 理论并重的论文,它选择在算法层面而不是调度层面解决这个问题:与其想办法更聪明地调度 BP 的前向/反向流水,不如直接问一句——能不能在整条链路完好无损之前,就让每一段独立更新?

这和大多数流水线并行论文选择的杠杆完全不同。GPipe、1F1B、PipeDream、Zero-Bubble,以及大量以调度为核心的后续工作,本质上都保留了同一张计算图——一长串前向操作接一长串反向操作,靠链式法则粘在一起——然后想办法把这条链的片段在多设备间交错、重排或重叠,从而压缩气泡、减少陈旧度。ZeroLock 则直接打断这条链本身:每个流水线阶段拥有自己的本地损失函数,只用上游邻居的前向输出就能独立更新自己的参数。没有任何梯度会跨越阶段边界。正是这一个设计选择,同时消除了吞吐瓶颈(不需要等下游反向跑完)和显存瓶颈(不需要等整条链的反向都结束才能释放激活),并且让论文真正搭出了一个能在 NVIDIA GPU 和未经改造的安卓手机上都跑起来的系统。

在深入 ZeroLock 具体怎么做之前,值得先把背景铺垫清楚:**“更新锁定”**到底在机制层面意味着什么、为什么此前的 BP-free 方法一直没能在系统层面被用到大模型微调上、以及本地目标训练的理论此前到底缺了什么,这篇论文才需要去补上。

前置知识:读正文之前需要了解的背景

流水线并行与”阶段”。 当模型太大放不进一张卡时,一种办法是流水线并行:把模型的层切成 KK 个连续的块,每个块分配给一台设备(称为一个”阶段”),微批次(micro-batch)像流水线一样依次流过这条阶段链。设备 kk 跑完块 kk 的前向,把输出(“隐藏状态”)交给设备 k+1k+1,以此类推直到最后一个阶段算出损失;然后梯度沿着同一条链反向流动。这在原理上很高效——每台设备只需要保存模型很小的一片——但如果朴素地一次处理一个微批次,大多数设备大部分时间都是空闲的,因为设备 kk 必须等到微批次 ii 已经流过所有下游设备并收到对应的反向信号后,才能开始处理微批次 i+1i+1

GPipe、1F1B、PipeDream:三种不打破 BP 前提下对抗气泡的办法。 GPipe 把一个逻辑批次内所有的前向都跑完,再统一跑反向,所以气泡很大,但代码简单、完全同步。1F1B(one-forward-one-backward)在一个批次内交错前向和反向计算,设备一旦某个微批次的反向可用就立刻执行,不用等所有微批次前向都跑完——这缩小了”流水线填充”气泡,但仍然要在共享优化器步骤之前”冲刷”流水线(清空所有在途的微批次),所以是用周期性的空闲时间换取更新的一致性。PipeDream 建立在 1F1B 之上,用”权重暂存”(weight stashing,即同时保留多个版本的参数)去掉了这次冲刷,让不同更新窗口之间的前向和反向得以重叠,同时保证每次反向看到的参数版本与它对应的前向完全一致。这进一步提升了利用率,代价是要存储多份权重快照——但更重要的是,就本文的论证而言,无论 GPipe、1F1B 还是 PipeDream,都没有改变论文所说的更新锁定(update locking):不管调度怎么变,每个阶段的参数更新依然依赖于完整跑完每个下游阶段的前向和反向计算,因为只要链路末端只有一个共享损失,链式法则就要求如此。

更新锁定,具体到公式层面。 如果模型计算 hk=fk(hk1;Wk)h_k = f_k(h_{k-1}; W_k)h0h_0 是输入嵌入,hKh_K 送入最终损失),那么更新 WkW_k 需要 L/Wk\partial \mathcal{L}/\partial W_k,根据链式法则又需要 L/hk\partial \mathcal{L}/\partial h_k,而这又要求从 hkh_khKh_K整条下游链路都必须先完成前向、再完成反向。这带来两个具体后果:

  • 吞吐瓶颈: 上游阶段(比如 S0S_0)必须等每个下游阶段(S1S_1S2S_2……)都跑完自己的反向,S0S_0 的反向——进而它的更新——才能进行。调度技巧可以缩小由此产生的气泡,但没法让它彻底消失,因为瓶颈来自依赖关系本身,而不只是调度顺序
  • 显存瓶颈: 上游阶段必须把自己前向算出的激活值(中间张量)一直留在显存里,直到自己的反向跑完——而根据上一条,这要等到所有下游阶段的反向都完成之后才可能发生。这意味着上游阶段占用显存的时间,比它”本地实际需要”的时间长得多。论文的 Fig. 1 把这一点直观地画了出来:在某个时刻,BP 基线下的 S0S_0 可能要同时保留三个微批次的激活值(A1A_1A3A_3),纯粹是因为它在等下游;而在更新解耦的方案里,它只需要保留正在实际使用的那一份激活。

BP-free 训练:两条既有路线,以及为什么它们此前都没能在大模型流水线微调的系统层面真正落地。 论文把此前的 BP-free 工作归为两类。反向梯度估计类方法(direct feedback alignment、以 MeZO 为代表的零阶优化)用不依赖链式法则的方式估计梯度——要么把目标误差直接传播到每一层,要么通过扰动参数、用损失差作为梯度代理——但这两种做法要么在复杂任务上明显损失精度,要么带来难以承受的算力开销(零阶方法尤其需要每次更新做大量前向评估才能得到可用的梯度估计)。目标重构类方法则给每个块配一个本地目标,让它的更新完全不需要来自下游的任何信号——NoProp 把每一层看作独立的去噪单元,预测编码在”最小化局部预测误差”和”用局部误差更新权重”之间交替进行,而深度渐进单调学习(depth-progressive monotonic learning,本文最直接的前身)给分类器的每一层配一个独立的本地损失。这一类方法对大模型微调更有前景,因为它往往能较好地保住精度——但论文关键的定位主张是:此前没有人(a)把它适配到大模型特有的要素上(比如 LoRA、自回归的逐 token 损失),(b)在任意数量的分块下证明收敛性(此前一篇 LoPT 只分析了两块的情形),或者(c)真正围绕它搭建出一个可部署的多 GPU 或多手机系统,配上真实流水线训练所需的运行时机制(缓冲、故障恢复、无需 RPC 的隐藏状态交换)。ZeroLock 的贡献正是精准地填上这三个空白——论文把它们表述为三个研究问题:Q1(怎样用本地目标微调大模型)、Q2(模块化解耦是否损害收敛性)、Q3(怎样搭建真实系统)——论文余下的部分正是围绕逐一回答这三个问题展开的。

LoRA 简述,因为 ZeroLock 建立在它之上。 低秩适配(LoRA)冻结第 ll 层预训练权重矩阵 wlw_l,只学习一个低秩更新 Δwl=BlAl\Delta w_l = B_l A_l,其中 BlRd3×d2B_l \in \mathbb{R}^{d_3 \times d_2}AlRd1×d3A_l \in \mathbb{R}^{d_1 \times d_3},秩 d3min(d1,d2)d_3 \ll \min(d_1, d_2)。只有 AlA_lBlB_l 参与训练;基础权重 wlw_l 始终不变。这一点对 ZeroLock 特别重要,因为它让每个流水线阶段的可训练参数体量保持很小——这是让微调在受限的边缘设备(包括手机)上可行的前提条件。

Figure 1(对应论文 Fig.1):(a) BP 基线与 (b) ZeroLock 的吞吐时序图,(c) BP 与 (d) ZeroLock 的显存占用图。在 BP 下,S0 必须同时保留 A1-A3 三份激活,因为要等下游反向完成;在 ZeroLock 下,S0 只保留正在实际使用的那一份激活。

Fig. 1 值得多看一会儿,因为它是让整篇论文的主张变得直观易懂的那张图。在 BP 时序图 (a) 中,注意阶段 S0S_0 的前向 F1,F2,F3F_1, F_2, F_3 都很早就完成了,但它的反向 B1,B2,B3B_1, B_2, B_3 却被往后推——每一个都要等 S1S_1S2S_2 上对应的反向先完成,图中画出的对角依赖箭头正是在描绘这条链式法则依赖关系。在 ZeroLock 时序图 (b) 中,每个阶段的反向紧跟在自己的前向之后,完全不依赖下游阶段——调度变成了彻底本地化的。显存面板 (c)、(d) 展示了随之而来的机制性后果:在 BP 下,S0S_0 要同时保留三个微批次的激活(图中重叠的彩色条),因为最早的那份要等到它自己的反向最终跑完才能释放;而在 ZeroLock 下,每一份激活在被自己的本地反向消耗完的那一刻就立刻释放。

2. ZeroLock 算法:给每个块配上自己的本地目标

2.1 分块与”读出头”

考虑一个由嵌入算子 E()E(\cdot) 加一叠 Transformer 层 L\mathcal{L} 组成的大模型。ZeroLock 把 L\mathcal{L} 切成 KK 个连续的块,第 kk 块拥有一部分层 Lk\mathcal{L}_k,包含冻结的基础权重 WkW_k 和可训练的 LoRA 参数 ΔWk=(BlAllLk)\Delta W_k = (B_l A_l \mid l \in \mathcal{L}_k)。设 xx 为输入,SS 为序列长度。跨块的前向计算非常简单:

h0=E(x),hk=fk(hk1;Wk+ΔWk),k=1,,K.(1)h_0 = E(x), \qquad h_k = f_k(h_{k-1}; W_k + \Delta W_k), \quad k = 1, \dots, K. \tag{1}

真正让本地训练成为可能的关键结构设计是读出头(readout head):在每一个kk 之后(不仅仅是最后一个块),ZeroLock 都挂上一个(冻结、共享的)读出头,直接把该块的隐藏状态映射到词表 logits:

zk=Norm(hk)Wlm,pk,i=softmax(zk,i).(2)z_k = \mathrm{Norm}(h_k) W_{lm}^{\top}, \qquad p_{k,i} = \mathrm{softmax}(z_{k,i}). \tag{2}

这里 WlmRV×dW_{lm} \in \mathbb{R}^{V \times d} 是所有块共享的(冻结的)语言模型头,zkRS×Vz_k \in \mathbb{R}^{S \times V} 是把第 kk 块的中间隐藏状态当作最终隐藏状态来计算得到的逐 token logits,pk,iR1×Vp_{k,i} \in \mathbb{R}^{1 \times V} 是第 kk 块对第 ii 个 token 位置在词表上的预测分布。正是这个设计把”某个中间块的输出”变成了一个可以直接算损失的东西,不需要等网络的其余部分——每个块都可以用它目前为止看到的信息,独立地对下一个 token 的分布做出自己的猜测。

2.2 本地损失:任务项 + 一致性项

给定第 kk 块的读出分布 pkp_k,ZeroLock 定义了一个由两部分组成的本地损失:

Lk(pk)=αLTask(pk)+(1α)LConsiskk1(pk),α(0,1].(3)L_k(p_k) = \alpha L_{\text{Task}}(p_k) + (1-\alpha) L^{k \to k-1}_{\text{Consis}}(p_k), \qquad \alpha \in (0, 1]. \tag{3}

(一)任务相关项。 这是全局交叉熵目标的一个本地化版本——它把第 kk 块的预测拉向真实的下一个 token:

LTask(pk)=1ΩiΩDψ ⁣(pk,i,py),(4)L_{\text{Task}}(p_k) = -\frac{1}{|\Omega|} \sum_{i \in \Omega} D_{\psi}\!\left(p_{k,i}, p_y\right), \tag{4}

其中 Ω\Omega 是有效 token 位置的下标集合(排除 prompt/padding),pyp_y 是真实标签的 one-hot 向量,DψD_\psi 是由一个严格凸、可微的势函数 ψ\psi 诱导出的 Bregman 散度Dψ(u,v)=ψ(u)ψ(v)ψ(v),uvD_\psi(u, v) = \psi(u) - \psi(v) - \langle \nabla\psi(v), u - v \rangle。取负熵作为 ψ\psi,就能得到 KL 散度这个特例;由于最小化预测分布与 one-hot 目标之间的 KL 散度等价于最小化交叉熵,在实现层面 Eq. (4) 其实就是普通的逐 token 交叉熵——用 Bregman 散度来表述,是为了让后面的收敛性证明保持一般性,而不是因为实现时真的需要什么特别玄妙的东西。

(二)一致性项。 这一项负责防止各个块彼此”飘远”——它把第 kk 块的预测拉向第 k1k-1 块的预测(通过停梯度算子 sg()\mathrm{sg}(\cdot) 把后者当作固定目标,所以这一项的梯度永远不会流回第 k1k-1 块):

LConsiskk1(pk)=1ΩiΩDψ ⁣(pk,i,sg(pk1,i)).(5)L^{k \to k-1}_{\text{Consis}}(p_k) = \frac{1}{|\Omega|} \sum_{i \in \Omega} D_\psi\!\left(p_{k,i}, \mathrm{sg}(p_{k-1,i})\right). \tag{5}

为什么两项都要,去掉任何一项会怎样? 如果完全去掉一致性项(α=1\alpha = 1),每个块就纯粹靠自己(不完整的)输入视角去拟合真实标签——早期的块只看过网络的前几层处理,却被要求自己单独解决整个语言建模任务,这比”作为深层网络最终预测的一部分贡献”要难得多;论文的实验(Fig. 3 中的 α=1\alpha=1 曲线)确实显示,这种设定下负对数似然明显比 α=0.5\alpha=0.5 更差。如果完全去掉任务项(α0\alpha \to 0),各块就只顾彼此保持一致,完全收不到与真实标签相连的信号,整条流水线可能收敛到一个自洽但错误的答案上。一致性项其实是在替代普通 BP 中链式法则本来隐式在做的事情:把”下一个块应该给出什么样的好预测”这一”形状”信息反向传播出去——只不过这里是靠匹配一个分布,而不是通过它反传一个梯度张量来实现。

微调循环。 有了本地损失,第 tt 轮迭代中第 kk 块内第 ll 层的实际参数更新就是一次完全普通的 SGD 步,只不过是本地计算的:

AlAlηLkAl,BlBlηLkBl.(6)A_l \leftarrow A_l - \eta \frac{\partial L_k}{\partial A_l}, \qquad B_l \leftarrow B_l - \eta \frac{\partial L_k}{\partial B_l}. \tag{6}

Remark 1(这在机制层面到底带来了什么)。 同一个块内部的层仍然用普通反向传播更新——块内部的 BP 完全没问题,因为一个块整体活在一台设备上,不跨阶段边界。变化的是不同块之间被解耦了:第 kk 块的反向只需要 hk1h_{k-1}(已经从上游收到,现成可用)和标签——它不需要来自第 k+1k+1 块的任何梯度信号。这正是同时消除吞吐气泡(Fig. 1(b))和显存膨胀(Fig. 1(d))的原因,而且是同一个原因:两者都不再是调度层面的产物,所以没有什么可以留给更聪明的调度器去优化掉了。

2.3 算法 1:ZeroLock 本地更新逐步展开

把上面的内容展开成针对单个块 kk 处理一个微批次的明确编号流程:

算法 1:ZeroLock 本地分块更新(每个块 k 针对每个微批次独立执行)
--------------------------------------------------------------------------------
输入 :上游隐藏状态 h_{k-1}(来自第 k-1 块的发送缓冲区,若 k=1 则为 E(x))
        当前微批次的标签/目标 token p_y,注意力掩码,位置编号
        冻结的基础权重 W_k,第 k 块内各层的可训练 LoRA 参数 (A_l, B_l)
        冻结的共享读出头 W_lm,混合权重 alpha,学习率 eta
输出 :更新后的 (A_l, B_l)(l 属于第 k 块);发往第 k+1 块的隐藏状态 h_k

1:  h_k <- f_k(h_{k-1}; W_k + Delta W_k)                     // 第 k 块内各层的前向计算
2:  将 h_k 从自动求导图中分离(DETACH);异步发送分离后的 h_k 到第 k+1 块的接收缓冲区
3:  z_k <- Norm(h_k) @ W_lm^T                                 // 读出头投影,对应 Eq. (2)
4:  p_k <- softmax(z_k)                                       // 逐 token 预测分布
5:  L_Task     <- 交叉熵(p_k[Omega], p_y[Omega])                // Eq. (4),Bregman->KL->交叉熵
6:  L_Consis   <- 散度(p_k[Omega], 停梯度(p_{k-1}[Omega]))       // Eq. (5)
7:  L_k <- alpha * L_Task + (1 - alpha) * L_Consis             // Eq. (3)
8:  仅对读出头(冻结,不更新)和第 k 块的 LoRA 参数反向传播 L_k
9:  在整个更新窗口内累积梯度 dL_k/dA_l、dL_k/dB_l(l 属于第 k 块)
10: 如果 当前微批次是本更新窗口内最后一个:
11:     A_l <- A_l - eta * (累积梯度 dL_k/dA_l),对第 k 块内每个 l              // 对应 Eq. (6)
12:     B_l <- B_l - eta * (累积梯度 dL_k/dB_l)
13:     在本地为该块保存检查点 (A_l, B_l, 优化器状态)
14: 返回 h_k(已在第 2 步发出,早于第 3-13 步完成)

初读论文时很容易漏掉的关键细节在第 2 步:隐藏状态在本地损失甚至还没有计算、更谈不上反向传播之前,就已经被分离并发送给下一个阶段。这正是论文第 3 节描述的”提前转发”(early forwarding)优化,它让第 k+1k+1 块几乎能在第 kk 块前向计算一结束就立刻开始自己的前向——第 k+1k+1 块完全不需要等第 kk 块的本地反向(第 3-13 步)完成,因为这些步骤根本不影响要往下游发送的内容。

Figure 2(对应论文 Fig.2):系统架构——一个协调器(成员管理、调度器、故障恢复)与多个执行器,每个执行器运行一个阶段工作进程(前向、读出头、本地反向更新),各自拥有独立的缓冲区和检查点。

3. 解耦会不会损害收敛性?论文的理论核心

这是 ZeroLock 之所以不只是一个工程技巧的关键所在,也是最容易被读者一带而过的部分——但不应该跳过,因为它回答了一个并不显然的问题:如果把块之间的梯度链条切断,各自优化 KK 个独立的本地目标,跟优化你真正关心的那个全局目标之间到底还有没有关系?论文的回答分两层:先是一个逐块表现界(第 II-B 节),说明在训练的某个固定点上损失如何随着块的推进演化;然后是一个完整的收敛速率分析(第 II-C 节),把整个解耦算法当作一个随机优化过程来处理,给出和普通 SGD 一样形式的收敛界。

3.1 引理 1:本地最优点,用全局目标重新表达

暂时略去 token 下标,一次只考虑一个 token 位置;设 pPp \in \mathcal{P}(词表上的概率单纯形)为某个块的读出分布,设 L:PR\mathcal{L}: \mathcal{P} \to \mathbb{R}全局目标,L(p)LTask(p)\mathcal{L}(p) \triangleq L_{\text{Task}}(p)。设 pkp_k^{\star} 为恰好最小化 Eq. (3) 中本地损失 Lk()L_k(\cdot) 的那个 pkp_k 取值。第一个结构性结论是:

pk=argminpPL(pk1),ppk1+1αDψ(p,pk1).(7)p_k^\star = \arg\min_{p \in \mathcal{P}} \left\langle \nabla \mathcal{L}(p_{k-1}), \, p - p_{k-1} \right\rangle + \frac{1}{\alpha} D_\psi(p, p_{k-1}). \tag{7}

推导展开。 从 Bregman 散度的定义 Dψ(u,v)=ψ(u)ψ(v)ψ(v),uvD_\psi(u,v) = \psi(u) - \psi(v) - \langle \nabla\psi(v), u-v\rangle 出发。分别代入 (u=p,v=pk1)(u=p, v=p_{k-1})(u=p,v=py)(u=p, v=p_y),把前一个代入结果整理出 ψ(p)\psi(p),再代入后一个式子。经过代数化简(论文附录形式的证明中完整给出),会收缩为:

Dψ(p,py)=Dψ(pk1,py)+L(pk1),ppk1+Dψ(p,pk1),(9)D_\psi(p, p_y) = D_\psi(p_{k-1}, p_y) + \langle \nabla \mathcal{L}(p_{k-1}), p - p_{k-1}\rangle + D_\psi(p, p_{k-1}), \tag{9}

这里用到了恒等式 pDψ(p,py)p=pk1=L(pk1)\nabla_p D_\psi(p, p_y)\big|_{p=p_{k-1}} = \nabla \mathcal{L}(p_{k-1}),也就是说”到目标的 Bregman 散度损失”的梯度,就是全局目标在上一块输出处求值得到的梯度。把 (9) 代回本地损失 Lk(p)=αDψ(p,py)+(1α)Dψ(p,pk1)L_k(p) = \alpha D_\psi(p, p_y) + (1-\alpha) D_\psi(p, p_{k-1})(这是 Eq. (3) 用一致性项以 pk1p_{k-1} 为目标重新写出的形式),并去掉不依赖 pp 的加性常数 αDψ(pk1,py)\alpha D_\psi(p_{k-1}, p_y),就恰好得到 Eq. (7)(相差一个 α\alpha 的缩放)。

这为什么重要,从直觉上讲。 Eq. (7) 说的是:本地最优点 pkp_k^\star 是这样一个点——它同时(a)沿着在上一块输出处求值的全局损失 L\mathcal{L} 的最速下降方向移动,(b)因为偏离 pk1p_{k-1} 太远而受到惩罚(通过 1αDψ\frac{1}{\alpha}D_\psi 这一项)。这在结构上与全局目标上的一次**近端梯度步(proximal gradient step)**完全一致,其中 α\alpha 扮演步长的角色,DψD_\psi 扮演近端正则项的角色。这是理解”本地目标构造到底在做什么”的一个很令人满意的视角:每个块并不是在解一个不相关的本地小谜题,而是在只用自己所在位置能获得的信息的前提下,朝全局最优迈出一步近端梯度步。

3.2 命题 1:逐块表现,以及额外误差项从哪里来

假设 L\mathcal{L} 在 Bregman 几何下相对 ψ\psiβ\beta-光滑的(假设 1):对所有 u,vu, vL(u)L(v)+L(v),uv+βDψ(u,v)\mathcal{L}(u) \le \mathcal{L}(v) + \langle \nabla\mathcal{L}(v), u-v\rangle + \beta D_\psi(u,v)。定义块级次优性 δkDψ(pk,pk)\delta_k \triangleq D_\psi(p_k, p_k^\star)——即实际训练出的块输出 pkp_k(受限于有限的模型容量,只能近似达到理论最优 pkp_k^\star)与该理论最优之间的差距。那么,对于 α<1/β\alpha < 1/\beta

L(pK)L(p0)(1αβ)k=1KDψ(pk,pk1)+1αk=1Kδk.(11)\mathcal{L}(p_K) \le \mathcal{L}(p_0) - \left(\frac{1}{\alpha} - \beta\right)\sum_{k=1}^{K} D_\psi(p_k, p_{k-1}) + \frac{1}{\alpha}\sum_{k=1}^{K} \delta_k. \tag{11}

推导要点。u=pk,v=pk1u = p_k, v = p_{k-1} 代入光滑性假设,得到 L(pk)L(pk1),pkpk(i)+L(pk1),pkpk1(ii)+L(pk1)+βDψ(pk,pk1)\mathcal{L}(p_k) \le \underbrace{\langle \nabla\mathcal{L}(p_{k-1}), p_k - p_k^\star\rangle}_{(i)} + \underbrace{\langle \nabla\mathcal{L}(p_{k-1}), p_k^\star - p_{k-1}\rangle}_{(ii)} + \mathcal{L}(p_{k-1}) + \beta D_\psi(p_k, p_{k-1})(式 12)。项 (i) 用本地最优点 (7) 的 KKT 稳定性条件加上 Bregman 三点恒等式来界定,化简后得到 αL(pk1),pkpk=δk+Dψ(pk,pk1)Dψ(pk,pk1)\alpha \langle \nabla \mathcal{L}(p_{k-1}), p_k - p_k^\star\rangle = \delta_k + D_\psi(p_k^\star, p_{k-1}) - D_\psi(p_k, p_{k-1})(式 13)。项 (ii) 用 pkp_k^\star 作为最小化点这一定义来界定,得到 L(pk1),pkpk1+1αDψ(pk,pk1)0\langle \nabla \mathcal{L}(p_{k-1}), p_k^\star - p_{k-1}\rangle + \frac{1}{\alpha}D_\psi(p_k^\star, p_{k-1}) \le 0。把两个界代回 (12),再对 k=1,,Kk = 1, \dots, K 求和做望远镜相消(相邻 kk 之间的 Dψ(pk,pk1)D_\psi(p_k^\star, p_{k-1}) 项相消),就得到了 Eq. (11)。

这到底告诉了我们什么。 随着分块数 KK 增大,最后一块的损失 L(pK)\mathcal{L}(p_K) 被界定为初始损失 L(p0)\mathcal{L}(p_0) 减去一个严格为负的收缩项(只要 α<1/β\alpha < 1/\beta,这是一个温和的步长条件),加上一个累积误差 1αkδk\frac{1}{\alpha}\sum_k \delta_k——它随分块数增长,反映了每个块的实际训练结果偏离自己本地最优点的程度。这是一个很直观的担忧的理论化版本:把模型切得越碎,理应在极限情况下开始变差,因为你累积的”本地不完美”项越来越多——但只要每个 δk\delta_k 都保持很小(也就是说每个块都有足够的容量,能相当接近自己的本地最优),这个界就仍然是有利的。这是一个确实有用的合理性界,不过值得指出的是(详见第 8 节批判性分析),它完全没有说明在真实的 Transformer 块上 δk\delta_k 实际有多大——这是一个经验问题,论文只用实验来回答,而没有给出进一步的理论界。

3.3 引理 2-3:把本地更新等价地变换成一次全局更新

理论上更难的挑战是把算法实际的更新规则(各个块独立在本地损失上做 SGD)和一个可以识别为对所有参数 ω\omega 联合的全局函数做下降联系起来。论文分三步完成这件事。

第一步(引理 2,目标等价性)。 把具体选择 KL 散度作为 DψD_\psi 时的本地损失 Lkt(p;α)L_k^t(p; \alpha) 展开,代数化简后可证明它等于 DKL(ppk,t(α))logZkt(α)D_{KL}(p \,\|\, p_k^{*,t}(\alpha)) - \log Z_k^t(\alpha)Zkt(α)Z_k^t(\alpha) 是一个不依赖 pp 的常数),其中目标分布具体为:

pk,t(l;α)=py(l)αpk1t(l)1αlpy(l)αpk1t(l)1α.(15)p_k^{*,t}(l; \alpha) = \frac{p_y(l)^{\alpha}\, p_{k-1}^t(l)^{1-\alpha}}{\sum_{l'} p_y(l')^{\alpha}\, p_{k-1}^t(l')^{1-\alpha}}. \tag{15}

这句话的意思是:最小化本地损失完全等价于最小化到某个特定目标分布 pk,t(α)p_k^{*,t}(\alpha) 的 KL 散度,而这个目标分布本身就是真实标签 pyp_y(权重 α\alpha)和上一块输出 pk1tp_{k-1}^t(权重 1α1-\alpha)之间的一个几何插值(加权乘积再归一化)。这是一个直观、可解释的对象:它其实就是”真实标签怎么说”和”上一块已经相信什么”之间 KL 测地线上的一点,按 α\alpha 向权重更大的那一方偏移。

一个具体的数值例子。 假设词表只有三个符号 {A,B,C}\{A, B, C\},真实标签 py=(1,0,0)p_y = (1, 0, 0)(token 是 AA),上一块的预测 pk1t=(0.5,0.3,0.2)p_{k-1}^t = (0.5, 0.3, 0.2),取 α=0.5\alpha = 0.5。按 Eq. (15) 逐项计算未归一化的权重:

wA=py(A)0.5pk1t(A)0.5=10.5×0.50.50.707,wB=py(B)0.5pk1t(B)0.5=00.5×0.30.5=0,wC=py(C)0.5pk1t(C)0.5=00.5×0.20.5=0.\begin{aligned} w_A &= p_y(A)^{0.5}\, p_{k-1}^t(A)^{0.5} = 1^{0.5} \times 0.5^{0.5} \approx 0.707,\\ w_B &= p_y(B)^{0.5}\, p_{k-1}^t(B)^{0.5} = 0^{0.5} \times 0.3^{0.5} = 0,\\ w_C &= p_y(C)^{0.5}\, p_{k-1}^t(C)^{0.5} = 0^{0.5} \times 0.2^{0.5} = 0. \end{aligned}

归一化后 pk,t(0.5)=(1,0,0)p_k^{*,t}(0.5) = (1, 0, 0)——因为真实标签 pyp_y 在非 AA 位置上取值严格为 0,任何非零的 α\alpha 都会把这些位置的权重直接乘成 0,目标分布退化为跟纯 one-hot 目标一样。换一种更有信息量的情形:假设 pyp_y 不是硬标签,而是标签平滑后的软目标 py=(0.9,0.06,0.04)p_y = (0.9, 0.06, 0.04),其余不变,取 α=0.5\alpha=0.5

wA=0.90.5×0.50.50.949×0.7070.671,wB=0.060.5×0.30.50.245×0.5480.134,wC=0.040.5×0.20.50.200×0.4470.089.\begin{aligned} w_A &= 0.9^{0.5} \times 0.5^{0.5} \approx 0.949 \times 0.707 \approx 0.671,\\ w_B &= 0.06^{0.5} \times 0.3^{0.5} \approx 0.245 \times 0.548 \approx 0.134,\\ w_C &= 0.04^{0.5} \times 0.2^{0.5} \approx 0.200 \times 0.447 \approx 0.089. \end{aligned}

归一化(除以 0.671+0.134+0.089=0.8940.671+0.134+0.089=0.894)得到 pk,t(0.5)(0.751,0.150,0.100)p_k^{*,t}(0.5) \approx (0.751, 0.150, 0.100)——介于纯真实标签 (0.9,0.06,0.04)(0.9, 0.06, 0.04) 与上一块输出 (0.5,0.3,0.2)(0.5, 0.3, 0.2) 之间,且比简单算术平均 (0.7,0.18,0.12)(0.7, 0.18, 0.12) 更靠近真实标签一些(因为几何插值在两个分布形状差异较大的位置上,惩罚更严格)。若改用 α=0.8\alpha=0.8(更偏向任务项),重复同样的计算会发现 AA 位置的权重进一步升高、目标分布进一步逼近真实标签,这与直觉完全吻合:α\alpha 越大,一致性项的拉力越弱,本地目标越接近纯监督学习。

第二步(引理 3,全局更新等价性)。 定义读出映射 pkt=f^k(pk1t;ωkt)p_k^t = \hat f_k(p_{k-1}^t; \omega_k^t)(从第 kk 块的输入分布到输出分布,在当前参数 ωkt\omega_k^t 下),及其雅可比矩阵 Jωk,ktf^k(pk1t;ωkt)/ωktJ_{\omega_k, k}^t \triangleq \partial \hat f_k(p_{k-1}^t;\omega_k^t)/\partial \omega_k^t。把这些矩阵沿所有 KK 个块块对角地堆叠成 JtJ^t,并定义 ekt(α)=logpktlogpk,t(α)e_k^t(\alpha) = \log p_k^t - \log p_k^{*,t}(\alpha)——这是当前输出与引理 2 中理想目标之间在对数空间的差距。对复合损失 Rt(ω;α)=E(x,y)[kDKL(f^k(pk1t;ωk)pk,t(α))]R^t(\omega;\alpha) = \mathbb{E}_{(x,y)}\left[\sum_k D_{KL}(\hat f_k(p_{k-1}^t;\omega_k) \,\|\, p_k^{*,t}(\alpha))\right](式 19)用链式法则,论文证明了:

ωt+1ωtηtωRt(ωt;α),ωRt(ω;α)=(Jt)et.(20)\omega^{t+1} \leftarrow \omega^t - \eta_t \, \nabla_\omega R^t(\omega^t; \alpha), \qquad \nabla_\omega R^t(\omega;\alpha) = (J^t)^{\top} e^t. \tag{20}

结论是:同时对所有 KK 个块独立计算出的本地 SGD 更新集合,在数学上完全等价于对某个定义良好的全局函数 Rt(ω;α)R^t(\omega;\alpha)(所有参数一起)做的一次梯度下降。这是整个理论贡献的核心所在——它把”一堆在不同设备上并行进行的解耦本地更新”转化成了”一个普通(虽然形式上有点特别)的随机梯度下降实例”,而后者的收敛性工具已经现成存在,几乎可以直接拿来复用。

第三步(引理 4,有界漂移)。 因为作为第 kk目标使用的 pk1tp_{k-1}^t 本身也在随着第 k1k-1 块参数的更新而变化,所以 Rt()R^t(\cdot) 严格来说是一个移动目标——第 tt 轮被优化的函数和第 t+1t+1 轮并不完全相同。论文对这种漂移给出了界:DT(α)t=0T1E[δtref(α)](Cfo(α)+βR(α)/2)ρ2M(α)t=0T1ηt2D_T(\alpha) \triangleq \sum_{t=0}^{T-1}\mathbb{E}[\delta_t^{\text{ref}}(\alpha)] \le (C_{fo}(\alpha) + \beta_R(\alpha)/2)\rho^2 M(\alpha)\sum_{t=0}^{T-1}\eta_t^2,利用了对读出映射的光滑性(假设 4)和 Lipschitz 连续性(假设 3)假设。直观来说:只要学习率 ηt\eta_t 按合适的方式衰减(下面定理 1 中确实如此),“目标在脚下漂移”这件事的累积影响就是有界的,不会让分析失控。

3.4 定理 1:收敛速率,以及为什么代价只是一个对数因子

有了上面的机制——本地更新等价于一次全局 SGD 步(引理 3),且由此产生的非平稳性是有界的(引理 4)——论文引入标准的非凸 SGD 收敛性假设(有界梯度、Lipschitz 连续、光滑性、稳定性、无偏有界方差采样;假设 2-6),对步长 ηt=η0/t+γ1/βR(α)\eta_t = \eta_0/\sqrt{t+\gamma} \le 1/\beta_R(\alpha) 推导出:

1Tt=0T1EωRt(ω;α)22(R0Ropt+DT(α))Tη0/T+γ+βR(α)σ(α)2t=0T1ηt2/BTη0/T+γ=O~ ⁣(1T).(23)\frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}\|\nabla_\omega R^t(\omega;\alpha)\|^2 \le \frac{2(R_0 - R_{\text{opt}} + D_T(\alpha))}{T\eta_0/\sqrt{T}+\gamma} + \frac{\beta_R(\alpha)\sigma(\alpha)^2 \sum_{t=0}^{T-1}\eta_t^2/B}{T\eta_0/\sqrt{T}+\gamma} = \tilde{\mathcal{O}}\!\left(\frac{1}{\sqrt{T}}\right). \tag{23}

证明思路。 对一次采样 SGD 更新 ωt+1ωtηtR^t(ωt)\omega^{t+1} \leftarrow \omega^t - \eta_t \nabla \hat R^t(\omega^t)(使用采样梯度)套用光滑性假设,用有界方差采样假设取条件期望,代入引理 4 中的漂移项 δtref(α)\delta_t^{\text{ref}}(\alpha),把跨越移动目标的 Rt(ωt+1)R^t(\omega^{t+1})Rt+1(ωt+1)R^{t+1}(\omega^{t+1}) 联系起来,整理出 ERt(ωt)2\mathbb{E}\|\nabla R^t(\omega^t)\|^2,再对 t=0,,T1t=0,\ldots,T-1 望远镜求和。这是一个完全标准的非凸 SGD 收敛证明模板——本文真正新颖、专属于这篇论文的部分,只有引理 3 这一个等价性,正是它让这个模板可以应用到一个解耦的、多块的、本地训练的系统上。

用大白话讲的核心结论。 普通基于反向传播的 SGD 在非凸目标上以速率 O(1/T)\mathcal{O}(1/\sqrt{T}) 收敛(衡量平均平方梯度范数)。ZeroLock 尽管彻底解耦了各块的更新、从来没有一个梯度跨越阶段边界,却仍然以速率 O~(1/T)\tilde{\mathcal{O}}(1/\sqrt{T}) 收敛——同样的多项式速率,只相差一个隐藏在波浪号里的多重对数因子(对某个常数 cc 而言是 logcT\log^c T)。换句话说:理论上讲,你几乎是”免费”拿掉了更新锁定这个约束,至少在渐近意义下、且在所列的光滑性/有界性假设成立的前提下如此。这是一个相当有分量的结论,也是本文最重要的科学贡献——在我看来,比系统工程部分更重要,因为这是同类结果中第一个针对任意分块数量给出的证明(据论文相关工作部分所述,此前的分析仅限于两块的情形)。

4. 从算法到系统:让它既快又稳的四项技术

一个能解耦更新的算法,对于打造一个快速、健壮的分布式系统来说是必要条件,但并不充分——你仍然要在设备之间搬运张量、处理掉队者和崩溃,还要避免通过”后门”重新引入同步开销。论文第三节给出了一个协调器/执行器架构(Fig. 2,见前文),以及四项具体技术,论文很小心地把它们分别归因到”吞吐”或”健壮性”这两个目标之一。

(一)提前转发(Early Forwarding)。 正如算法 1 第 2 步所强调的,一个块前向计算产生的隐藏状态,在这个块自己的本地反向跑之前,就已经被分离并发往下游。为什么选这个设计,而不是显然的替代方案? 显然的替代方案——先前向、再反向、再发送——会重新引入一个串行依赖(下游要等上游完整走完本地那一步),即使本地目标构造并不要求这种顺序;提前转发是一个纯收益、没有正确性代价的选择,因为本地反向只消耗前向阶段已经算出的信息(隐藏状态 hk1h_{k-1}hkh_k),不需要任何来自下游块反向计算的东西。这个设计唯一会失效的地方,正如论文在第 III-C 节坦率指出的那样,是安卓上的 ExecuTorch 后端,它默认把前向和反向捆绑成一个原子方法——这需要下面会讲到的定制”流水线标记算子”来绕过。

(二)独立执行与检查点。 每个执行器只保存自己的可训练参数、优化器状态和检查点。曾考虑过的替代方案: 一个全局同步的检查点(大多数标准分布式训练检查点系统的做法,确保所有阶段在同一个逻辑步检查点)会更容易推理,但会强迫各阶段之间为了检查点一致性而做部分同步——这恰恰违背了解耦的初衷。所选设计的代价是一个可复现性上的微妙之处:由于每个阶段独立、按自己的节奏检查点,从检查点恢复训练需要小心保证所有阶段的检查点对应到一个”足够一致”的训练时刻,这需要下面的故障恢复机制主动通过重放来管理,而不是靠共享时钟免费获得。

(三)仅交换状态的跨阶段通信。 执行器之间只传输前向隐藏状态——从不传输梯度,也不传输优化器状态。这在哪里带来了可衡量的收益: 后面会讨论的 Fig. 6(b) 显示,ZeroLock 的吞吐优势之所以随网络链路变慢而扩大,正是因为依赖 BP 的基线最终必须在同一条链路上传输和前向隐藏状态体量相当的反向梯度张量,大致把通信负载翻倍,而 ZeroLock 只有前向流量。

(四)缓冲辅助的状态交换。 每个执行器保留一个有界的缓冲区,存放上游阶段最近若干次的隐藏状态,专门用来在故障后重放微批次,而不需要重新跑上游的计算。设计张力: 更大的缓冲区能提供更多重放余地(更快、更局部化的故障恢复),代价是每个阶段要多占用一些显存——这恰好和算法在其他方面想要优化的目标相反。论文没有报告如何调这个缓冲区大小,这是下文第 10 节指出的可复现性缺口之一。

4.1 阶段内执行,展开为伪代码(算法 2)

论文用散文形式(S1/S2/S3 列表)描述了运行时每个微批次的执行流程;把它展开成一个明确的编号流程,同时展示跨阶段的缓冲机制:

算法 2:ZeroLock 跨阶段运行时(协调器 + 一对执行器,稳态运行)
--------------------------------------------------------------------------------
状态 :hidden_state_buffer[k]      // 第 k 个执行器的有界 FIFO 缓冲区,由第 k-1 个执行器填充
        in_flight_depth_limit       // 一对阶段之间允许的最大未确认条目数
        update_window                // 协调器设定的一段连续微批次编号区间
输入 :分配到当前 update_window 的微批次流,编号为 1..M

1:  协调器 定义 update_window = [i_start, i_end] 及微批次调度顺序
2:  对 update_window 内的每个微批次 i:
3:      执行器 k 等待 hidden_state_buffer[k].entry(i) 的就绪事件     // S1:输入
4:          (若 k == 1:entry(i) 是 E(x_i),无需等待)
5:      执行器 k 从 hidden_state_buffer[k] 取出 h_{k-1,i},以及掩码/位置编号/标签(i)
6:      执行器 k 执行 算法1 的第 1-2 步(前向 + 分离 + 异步发送)   // S2
7:      执行器 k 通过发送缓冲区把 h_{k,i} 提交进 hidden_state_buffer[k+1]
8:          -- 此时,执行器 k+1 可以立刻开始为微批次 i 执行自己的第 3 步 --
9:      执行器 k 执行 算法1 的第 3-9 步(读出、本地损失、反向、累积梯度)  // S3
10:     如果 i == i_end(本 update_window 内最后一个微批次):
11:         执行器 k 执行 算法1 的第 11-13 步(优化器步 + 检查点)
12: 协调器 推进 update_window;清理超出重放窗口的 hidden_state_buffer 旧条目
13: 关于预登记接收(仅限 GPU):执行器 k+1 可以在执行器 k 还没产出 h_{k,i+1} 之前,
        提前发出 receive_entry(i+1) 请求,把通信建立与计算过程重叠起来

有两个实现细节容易被忽略,但对正确性很重要:在途深度限制in_flight_depth_limit)防止快上游和慢下游之间堆积无限多的未确认隐藏状态条目——没有它,一个快速的上游执行器可能会无限领先,耗尽下游缓冲区的显存;预登记接收(第 13 行)是一项仅限 GPU 的优化,让下游阶段在数据就绪之前就先登记好接收请求,这样计算流只会等待它真正需要的那个特定就绪事件,而不会被无关的通信流量阻塞。

4.2 故障恢复:算法 3

由于每个阶段独立检查点、隐藏状态又被缓冲而不是立即丢弃,ZeroLock 能够在单个阶段故障后恢复,而不需要整条流水线的全局回滚——这是技术(二)和(四)的直接结构性后果,而不是事后另外补上的机制。

算法 3:ZeroLock 故障恢复(当执行器 k 被检测为故障时触发)
--------------------------------------------------------------------------------
输入 :执行器 k 最近一次提交的检查点(A_l, B_l, 优化器状态, 进度元数据)
        故障前从上游存活下来的 hidden_state_buffer[k] 内容(依靠技术四)
        故障发生前已完全处理完毕、被标记为"已提交"的更新窗口集合

1:  协调器 检测到执行器 k 故障(通过成员管理/执行监控发现心跳缺失)
2:  协调器 将执行器 k 在途的更新窗口标记为"故障,待重放"
3:  执行器 k 从最近一次本地检查点恢复 (A_l, B_l, 优化器状态)          // 被动式恢复
4:  执行器 k 找出差距:(检查点记录的进度)与(其他阶段已到达的流水线前沿)之间的差距
5:  对差距中的每个微批次 i(已在上游缓冲,尚未在本地重新处理过):
6:      执行器 k 从 hidden_state_buffer[k] 重新取出 h_{k-1,i}    // 无需上游重新计算
7:      执行器 k 重新执行 算法1(第 1-13 步)处理该微批次 i           // "重放"
8:  执行器 k 追上流水线前沿后,向协调器发出信号
9:  协调器 恢复正常调度;始终存活的下游阶段从未被回滚过

为什么这样设计,而不是显然的替代方案(全局回滚)? 显然的替代方案——本文自己用来对比的基线,同步 1F1B 采用的方案——是把整条流水线回滚到最近一次全局一致的检查点,然后从那里重放,因为 1F1B 的权重暂存和共享优化器步骤确实会在阶段之间产生真正的状态依赖,使得纯本地回滚是不正确的。ZeroLock 的本地恢复之所以可能做到,正是因为更新从一开始就没有跨越阶段边界——所以”一致性”从来就不是一个全局性质,恢复一个阶段不可能破坏另一个阶段的状态,因为从来没有梯度或共享优化器状态在它们之间流动过。这个设计可能在哪里失效: 如果故障恰好发生在某个阶段完成本地优化器步骤、但还没成功写入检查点的这个短暂窗口内,重放机制就必须正确检测到该次更新未被持久化提交,且不能重复应用它——论文没有明确走一遍这个边界情形,这是下文批判性分析部分指出的具体缺口之一。

4.3 移动端执行后端:一个被意外约束逼出来的设计选择

对于安卓部署,论文使用 ExecuTorch——Meta 官方的 PyTorch 原生端侧运行时,它通过提前编译(AOT)生成静态计算表示,序列化为 .pte 程序,可以在没有 Python 解释器的资源受限终端设备上确定性地执行。具体的工程障碍: ExecuTorch 的训练支持把前向和反向计算捆绑进一个原子方法调用,只有在本地反向梯度计算已经完成之后才会暴露前向隐藏状态——这直接破坏了提前转发(技术一),因为等你终于有东西可以往下游发的时候,你已经付出了本来想要隐藏掉的完整本地反向延迟。解决办法: 作者在编译好的 PTE 图中插入一个定制的”流水线标记”算子,放在分离后的隐藏状态输出之后、参数梯度子图之前。这个算子不做任何张量计算——它唯一的作用是通知运行时捕获隐藏输出,并挂起该方法的执行状态(调用栈、张量、优化器缓冲区)。一个定制的两阶段安卓接口让第一阶段执行到标记处为止(返回隐藏状态供立即传输),第二阶段则从挂起状态恢复,完成本地反向并应用设备端的 AdamW 更新。这是一段相当巧妙的工程实现,正因为它证明了提前转发不仅仅是 GPU 上的调度小花招,而是作者们愿意去和底层运行时”较劲”、在一个约束大得多的平台上也要保住的设计属性——这是一个很好的信号,说明这项技术在实践中确实重要,而不只是在理想化的 GPU 场景下才有意义。

5. 实验:理论碰上真实系统,还站得住吗

论文搭建了两套真实原型——一套多 GPU/CPU 服务器环境,一套安卓手机部署——评估三方面内容:(E1)相对 BP 基线的显存与吞吐对比,(E2)故障恢复,(E3)设备端可行性。模型与数据:TinyLlama,切成三个连续的块(分别放在一张 NVIDIA L40 GPU 上,或一台安卓手机上),用 LoRA(秩 4,缩放因子 16)在 AG News 的一个固定 10,000 条样本子集上微调,序列填充/截断到 128 token,重复三个随机种子。对比基线:GPipe、1F1B、以及带权重暂存的 PipeDream。

5.1 E1:GPU 原型上的显存与吞吐

Figure 3(对应论文 Fig.3):(a) 准确率与 (b) 负对数似然随训练进度的变化,对比 BP(单卡)、GPipe、1F1B、PipeDream,以及 alpha=1 和 alpha=0.5 两种 ZeroLock 配置。

第一个也是最重要的合理性检验是 Fig. 3:解耦各块真的会损害模型质量吗?答案基本是不会——ZeroLock 的准确率曲线和 BP 基线贴得很紧,NLL 也具有竞争力,其中 α=0.5\alpha=0.5(任务项和一致性项权重相等)明显优于 α=1\alpha=1(只有任务项),这正好印证了第 2.2 节讨论过的直觉:对于早期、信息量不足的块来说,没有一致性信号支撑、只靠纯本地任务损失,是一个明显更难的学习问题。

Figure 4(对应论文 Fig.4):GPU 峰值显存按组件(模型+状态、激活、前向/损失/反向、输入/接收、权重缓存、其他)拆解,对比 GPipe、1F1B、PipeDream、ZeroLock 在三个阶段上的表现(b=8, m=4)。

Fig. 4 把显存收益具体来自哪里拆开来看,如果有人问”这是不是一篇披着流水线并行外衣的激活显存论文”,我会指着这张图说:很大程度上是的。ZeroLock 相较 GPipe、1F1B、PipeDream,平均每阶段峰值显存分别降低 47.8%、14.7%、14.6%(对最大阶段峰值来说是 55.3%、26.6%、26.5%),论文明确把大部分降幅归因于激活值的消除——具体来说”激活”一栏分别降低了 75.4%、40.8%、48.8%。注意”模型+状态”和”权重缓存”两栏在各方法间基本没变(ZeroLock 并没有改变 LoRA 参数或基础模型的存储方式)——收益完全集中在不需要为等待整条链的反向而一直保留前向激活值上,这正是论文 Fig. 1 在概念层面展示的机制。文中还藏着一个实践上很关键的数据点:在 m=4m=4 时,GPipe、1F1B、PipeDream 分别在物理批大小 b=10,20,20b = 10, 20, 20 时就发生显存溢出,而 ZeroLock 可以扩展到 b=28b = 28 才溢出——说明显存节省直接转化为一个更大的可用批大小空间,而不仅仅是图表上一个更小的数字。

Figure 5(对应论文 Fig.5):(a)(c) 最大单阶段显存和 (b)(d) 吞吐量随批次几何形状 b/m 的变化,总批大小分别为 B=32 和 B=128。

Fig. 5 增加了批次几何形状这一维度:随着物理批次与微批次数量之比 b/mb/m 增大(也就是每次优化器更新前,微批次调用次数更少、每次更大),GPipe 的显存急剧增长,因为它必须同时保留更多的前向计算图;而 ZeroLock 的显存增长要慢得多——证实了收益在”基线保留激活最痛苦”(即大 b/mb/m)的场景下会更明显。在吞吐方面,默认的 b=8,m=4b=8, m=4 设置下,ZeroLock 相对 GPipe、1F1B、PipeDream 分别提升吞吐 55.8%、62.8%、4.9%——注意相对 PipeDream 的优势明显小得多,这是合理的,因为 PipeDream 本来就是气泡优化最充分的 BP 基线;ZeroLock 相对 PipeDream 的优势几乎完全来自显存侧(支持更大批次),而不是消除 PipeDream 已经消除得差不多的额外调度气泡。

Figure 6(对应论文 Fig.6):吞吐量随 (a) 流水线阶段数/GPU 数量和 (b) 逐渐变慢的发送端网络链路(本地、Wi-Fi、移动、受限)的变化。

Fig. 6(a) 显示 ZeroLock 的吞吐优势随着阶段数增加(2 到 4 块 GPU)而扩大——与”去掉一个随流水线深度增长的依赖关系”这一预期一致,而 BP 基线的气泡开销会随阶段数增多而累积。Fig. 6(b) 是最直接验证第 4 节设计选择(三)的图:随着模拟网络链路变慢(本地 → Wi-Fi → 移动 → 受限,带宽从约 1000 Mbps 降到 50 Mbps,延迟从约 0 升到 30 ms),ZeroLock 相对 1F1B 和 PipeDream 的相对吞吐优势扩大,正是因为那些基线必须在同一条链路上除了前向隐藏状态之外还要传输反向梯度张量,而 ZeroLock 只需要传输前向隐藏状态。

Figure 7(对应论文 Fig.7):一秒稳态窗口内 CUDA 活动轨迹(CUDA 核函数 vs. 空闲),(a) GPipe、(b) 1F1B、(c) PipeDream、(d) ZeroLock,跨阶段 S0-S2。

Fig. 7 是一段很不错的机制性证据,而不只是一个汇总统计量:你可以直接从视觉上看出 GPipe 的各阶段(a)有很长、结构规整的空闲间隙(经典的”先填充后排空”气泡模式),1F1B(b)的间隙更短但依然可见,PipeDream(c)和 ZeroLock(d)看起来都密集得多——但论文报告说,ZeroLock 相较 PipeDream 表现出更均衡的各阶段活动负载,活跃比例高出 13.7%,正文对此很小心地补充说这”是流水线调度的诊断性指标,而不是吞吐量的代理指标”,因为空闲时间还包含主机端调度开销和 Gloo 通信开销,与算法本身无关。我很欣赏这个补充说明——这正是我希望更多系统论文能做到的、对指标能衡量什么/不能衡量什么保持诚实的态度。

5.2 E2:故障恢复

用三个阶段,b=1,m=8b=1, m=8,论文在训练中途的某个窗口注入一次故障(先四个”序幕”窗口,再四个 Stage 1 处的”故障”窗口,再四个”恢复”窗口),整个过程中工作进程和 GPU 上下文都保持存活(这种故障注入设计把恢复协议本身的代价与原始进程重启的代价区分开来)。与同步 1F1B 的全局回滚+完整重放恢复方式相比,ZeroLock 的本地恢复将恢复延迟降低了 368 毫秒(2013.1 ± 6.4 ms vs. 2381.2 ± 66.6 ms),并把传输流量减半(96 MiB vs. 192 MiB)——这两个数字都直接反映了只有故障阶段需要重放,而不需要整条流水线的同步回滚。

5.3 E3:安卓手机上的设备端评估

Figure 8(对应论文 Fig.8):(a) 按阶段拆分的延迟明细(前向、本地反向、优化器、其他本地耗时、排队等待)和 (b) 三部安卓手机原型上各阶段的本地损失随优化器步数收敛的情况。

论文把三个 TinyLlama 训练 PTE(分别覆盖 Transformer 的 [0,6]、[7,13]、[14,21] 层)部署在一台 NX809J、一台 Lenovo L71091 和一台 Pixel 10 Pro XL 上,用 128 条样本、序列长度 128、LoRA 秩 8 进行训练。核心数据:应用峰值 PSS 低于 4000 MiB,电池温度约 37°C,整个微调过程的墙钟时间为 1644.1 秒(吞吐 0.0779 条/秒)——以服务器标准来看并不快,但重点在于可行性而不是速度:据我所知,这是首批(之一)在未经改造的消费级安卓硬件上端到端跑通 BP-free 本地目标流水线训练的演示。Fig. 8(a) 显示各设备的活跃/等待时间相当均衡(S0 等待 11.4 秒对活跃 12.2 秒,与 S1、S2 的比例大致相当),说明这三台异构设备尽管算力差异明显,负载却相对均衡,不过论文没有说明这种均衡是刻意设计(比如通过不均等的层数分配来补偿设备异构性)的结果,还是恰好碰上了这一组具体硬件的巧合。Fig. 8(b) 显示三个阶段的本地损失在约 125 步优化器步骤内都在收敛,其中 S0(最早、信息量最少的块)的起始损失和最终收敛损失都明显高于 S2——这很好地印证了第 2.2 节关于 α=1\alpha=1 的讨论:即使一致性项在起作用,更早的块在结构上依然面临更难的本地预测问题。

6. 重新审视设计选择:为什么是这样,而不是显然的替代方案

除了第 4 节已经讨论过的四项系统技术之外,还有三个算法层面的设计选择值得展开讲原因/替代方案/边界,因为它们在第一遍阅读时很容易被略过。

冻结的共享读出头,而非逐块可训练的头。 ZeroLock 对每一个中间块都复用预训练模型自己最后的归一化层和 LM 头作为读出头(第 III-B 节中的”(a)静态模型头”选项),而不是给每个块训练一个独立的小分类器头(这是此前一些视觉模型本地学习方法的做法)。替代方案——逐块可训练的头——会增加可训练参数,并且要求中间表示能被专门针对那个深度训练出的头所解读,而冻结共享头能让每个块的本地损失在尺度上直接可比,并直接反映”这个中间表示离最终层已经能读懂的东西有多近”这一信息。这个设计会在哪里失效: 论文自己指出,这假设中间隐藏状态能被最终头”解读”,作者承认这”在复杂的生成式任务上可能不成立”——这正是他们同时提供读出适配器方案(b,一个初始化接近恒等映射的小残差 MLP)作为备选方案的原因,用于中间表示和冻结头期望的输入空间偏离太多的情形。值得注意的是,第 IV 节的实验没有一个真正用到方案(b),也没有报告什么情况下才需要它,这使得”能被最终头读懂”这一假设到底能被拉伸多远才会导致质量下降,仍然是一个悬而未决的问题。

用 KL/Bregman 散度做一致性项,而非更简单的 L2 或特征匹配损失。预测分布层面(pk1p_{k-1},即上一块本来会预测出的下一个 token 分布)而不是原始隐藏状态层面(hk1h_{k-1},即隐藏向量之间的 L2 或余弦损失)强制一致性,正是让引理 1 及第 3 节整套收敛性机制得以成立的关键——整个证明链条依赖于能把任务损失和一致性损失都写成同一个概率单纯形空间 P\mathcal{P} 上的 Bregman 散度形式,这样才能把本地最优点(Eq. 7)表达为全局目标梯度的函数。一个隐藏状态空间的一致性损失可以完全绕开每个块都需要读出头这件事,也许计算上更便宜,但没法同样干净地归约为对一个标量目标做近端梯度步,因为隐藏状态空间中的欧氏距离或余弦相似度和模型最终训练所依赖的、基于 KL 散度的全局交叉熵目标之间,没有明显的关联。这是一个”理论上便利的选择”和”实践上有动机的选择”恰好重合的例子,值得强调这是一次真正出色的设计,而不只是运气使然——读出头在这里身兼二职:既是”如何计算本地损失”这件事的具体实现方式,也是让理论可分析的那个关键对象。

固定的、每次运行一个标量 α\alpha,而非可学习或按 token 自适应的混合权重。 Eq. (3) 中的混合权重 α\alpha 被当作单个全局超参数,每次运行只调一次(论文 Fig. 3 只对比了 α=1\alpha=1α=0.5\alpha=0.5 两种取值)。一种自适应的替代方案——比如为更靠后的块(已经看过更多上下文,或许可以合理地更看重任务项)提高 α\alpha,或者在训练早期(此时一致性目标 pk1p_{k-1} 本身还很嘈杂)降低 α\alpha——鉴于论文自己也发现 α=0.5\alpha=0.5 明显优于 α=1\alpha=1,这看起来是一个很自然的改进方向。论文完全没有探索这个设计空间,这算是一个错失的机会:命题 1 的误差界明确地通过 (1/αβ)(1/\alpha - \beta) 收缩项和 δk/α\delta_k/\alpha 误差项依赖于 α\alpha,所以有理论依据认为按块调优或按迭代进度调度的 α\alpha 可能优于单一全局常数,但这在目前的论文中仍未被探索。

6.1 一张总览表:四种流水线训练方案的对比

为了把 GPipe、1F1B、PipeDream 与 ZeroLock 在关键设计维度上的差异放在一起看,下面这张表总结了四种方案在更新机制、跨阶段信息、内存保留策略、以及故障恢复粒度上的本质区别:

方案更新依据跨阶段传递内容激活値保留策略故障恢复粒度
GPipe全局链式法则,先全部前向再全部反向前向隐藏状态 + 反向梯度每个阶段需保留整个逻辑批次的激活需全局回滚重跑
1F1B全局链式法则,交错前/反向但每批次仍需冲刷前向隐藏状态 + 反向梯度比 GPipe 少,但仍需等待下游反向需全局回滚重跑
PipeDream全局链式法则,权重暂存允许跨窗口重叠前向隐藏状态 + 反向梯度 + 多份权重快照最少,但需额外存储权重快照需全局回滚重跑(本文对比基准)
ZeroLock局部目标,无跨阶段链式法则依赖仅前向隐藏状态(不传梯度)每阶段只保留自己正在使用的激活只需单阶段本地重放

这张表能很直接地看出 ZeroLock 为什么能同时拿下吸吐量、显存、以及故障恢复三项指标上的优势:它是表中唯一一个在“更新依据”一栏不需要全局链式法则的方案,而其余三栌的优势均是这一选择的直接推论结果,而不是独立的工程优化。

7. 论文自己承认的局限,以及一些没有明说的部分

论文对自身的几个边界相当坦诚:(a)静态模型头设计假设中间隐藏状态能被最终头解读,对复杂的生成式任务可能不成立;(b)第 II-C 节的理论框架依赖标准的非凸 SGD 假设(有界梯度、光滑性、Lipschitz 连续性、有界方差采样),虽然这些在文献中很常见,但对深层 Transformer 的损失曲面来说其实是相当强的假设,而且论文并没有针对实验里具体用到的 TinyLlama + LoRA 设置去独立验证这些假设是否成立;(c)论文自己在未来工作里提到”进一步引入算子级优化,以进一步提升吞吐、降低显存”——这其实是在隐含承认目前的系统实现还没有在算子层面做优化,所以报告出来的吞吐数字很可能低估了一个更成熟实现能达到的水平(不管是对 ZeroLock 还是基线都一样,因为基线同样没有做算子级优化,但这确实意味着报告的百分比反映的是这套具体实现的成熟度,而不一定是算法本身的上限)。

在我看来,还有几处局限论文没有完全展开:

  • 规模。 所有实验都用 TinyLlama(约 11 亿参数),只切成三个块。理论上没有任何东西禁止扩展到更大的模型或更多的块,命题 1 的误差累积界也明确是针对一般的 KK 给出的,但对于分块数超过 3、或基础模型规模超过 11 亿参数时 δk\delta_k(驱动累积误差项的逐块次优性差距)会如何变化,论文完全没有给出经验证据——而这恰恰是流水线并行在实践中真正重要的场景(没有人会仅仅出于生产目的把一个 11 亿参数模型切成三块放到三张 L40 上;这样做通常是为了做原型验证)。
  • 任务多样性。 唯一评估过的微调任务(AG News,一个 128 token 长度、偏文本分类风格的短序列数据集)对于”能被最终头解读”这一假设来说是一个相对容易成立的场景,因为新闻主题信号很可能在前向计算的早期就已经出现。观察 ZeroLock 与 BP 之间的精度差距在需要跨全序列进行长程推理的任务(例如长文档问答或多步算术)上会如何变化,会很有信息量——也是一个自然的后续实验方向,因为在这类任务上,早期块的本地预测大概率是对最终正确答案质量更差的代理。
  • 读出适配器备选方案未经评估。 由于方案(b)的存在专门是为了应对”可解读性假设”失效的情形,而论文没有报告任何一次使用它的实验,我们完全无从判断实践者应该在什么情况下选用它,或者相比默认的静态头它会带来多少开销/质量变化。
  • 没有对额外读出头的前向计算做算力/墙钟核算。 在每个块之后(不只是最后一个)都挂一个读出头,意味着每个微批次要做 KK 次独立的 softmax 加交叉熵计算,而不是一次——对于词表规模达到几万的模型,这不是免费的,论文的显存拆解(Fig. 4)和吞吐数字里大概率已经隐含地包含了这部分开销,但论文从未把它单独拆出来量化,这使得很难判断这项开销会随词表规模或 KK 如何变化。

8. 批判性分析:我会在哪些地方提出质疑

(a)这篇论文本身的具体弱点。 最大的问题在于理论论述的规模(一般 KK、一般模型架构、一般分块方式)和实证验证的规模(一个 11 亿参数模型、三个块、一个短序列分类风格任务)之间的落差。定理 1 给出的收敛速率在 TT 上是渐近的,其明确表述中并不显式依赖 KK,这是一个干净、有吸引力的结果,但隐藏在 O~\tilde{\mathcal{O}} 符号里的常数项(具体来说是 βR(α)\beta_R(\alpha)σ(α)\sigma(\alpha)DT(α)D_T(\alpha))在原则上完全可能随 KK 增长,而论文没有对此做任何分析;实证结果对此也无法提供任何信号,因为 K=3K=3 从未被变化过。一个读者读完这篇论文,很可能(过度)泛化出”ZeroLock 在任意流水线深度下都有相同的理论保证”这样的结论,但证据实际支持的诚实说法更接近于”ZeroLock 的理论在抽象意义上对一般 KK 都成立,在实践中对 K=3K=3 表现良好”。

(b)作者低估或遗漏的局限。 Fig. 4/5/6 中的对比全部是相对基于 BP 的流水线基线(GPipe、1F1B、PipeDream);论文自己的相关工作部分(脚注 5)明确指出,没有和其他的 BP-free 本地目标基线(PPLL、FluidPipe、SCPL、LoPT)做对比,原因是这些工作要么没有开源代码,要么不是为大模型微调设计的,要么没有实现真正的多 GPU 系统级流水线并行。这是一个可以理解的实际原因,但也意味着论文从未用实验证明 ZeroLock 这个具体的本地目标设计(任务项 + 通过共享冻结读出头实现的一致性项)比它在相关工作讨论中明确批评过的其他本地目标设计(比如 SCPL 的逐段监督对比损失,或者 LoPT 局限于两块的方案)更好——对这些方法的批评是概念性/定性的,不是实验性的。另外,故障恢复实验(E2)只在一个受控、人为设置的窗口内注入了一次、针对一个阶段的故障;没有评估并发的多阶段故障,也没有评估恰好发生在检查点写入过程中的故障(也就是第 4.2 节中指出的那个具体边界情形)。

(c)具体、可执行的改进建议。 第一,在参数规模 70 亿到 130 亿的模型上,用 K{4,8,16}K \in \{4, 8, 16\} 个块跑同一套显存/吞吐/故障恢复实验,真正检验”一般分块数量”这一理论框架在实践中通常真正需要流水线并行的模型规模下是否依然成立,而不是仅仅作为研究原型的一个便利选择。第二,实证报告 δk\delta_k(命题 1 中的本地最优次优性)随块数和训练步数的变化——这是可以直接测量的(它就是 Dψ(pk,pk)D_\psi(p_k, p_k^\star),而 pkp_k^\star 可以在一个留出验证批次上数值近似),这能让读者看到理论所警示的误差累积项,在实践中究竟是可以忽略不计,还是对能使用多少块构成了一个真实的限制。第三,至少加入一个长程依赖任务(例如超过 4K token 的大海捞针检索或多跳问答),压力测试支撑静态读出头的”可被最终头解读”这一假设,因为 128 token 长度的 AG News 对这个假设来说接近于最理想的场景。第四,直接和脚注 5 中提到的至少一个 BP-free 本地目标竞争方法做对比(哪怕只是在单机、任务受限的场景下复现 SCPL 或 LoPT 核心更新规则,而不需要完整的系统级流水线并行),这样 Fig. 3-6 中的精度/吞吐对比就不会只针对一个不同的算法家族(BP),也能覆盖到概念上最接近的竞争对手。

9. 可复现性说明

论文提供了一个代码链接(https://anonymous.4open.science/r/unlock_trainer-105B,一个匿名化的投稿仓库),这是一个积极信号,不过作为本次评述的一部分,我并未能独立验证其内容。对 GPU 实验(LoRA 秩 4、缩放因子 16,三个块/三张 L40 GPU,AG News 10,000 条样本子集,序列长度 128,三个随机种子)和移动端实验(LoRA 秩 8、缩放因子 16、学习率 10410^{-4},128 条样本,序列长度 128,b=1b=1)报告的超参数相对完整。从头复现时会缺失或说明不够充分的细节包括:具体的学习率调度和优化器(AdamW 只在移动端后端被明确提及,GPU 实验中并未明确确认);理论步长调度 ηt=η0/t+γ\eta_t = \eta_0/\sqrt{t+\gamma}η0\eta_0γ\gamma 的具体取值,相对于实验中实际使用的实践性调度方案(两者不必相同,但论文没有澄清);运行时用到的在途深度限制和隐藏状态缓冲区大小取值;以及三个块具体切分的层范围——除了安卓部署明确给出的 [0,6]/[7,13]/[14,21] 划分之外,没有确认这是否与 GPU 实验的分块方式一致。论文的”AI 使用披露”部分是一个值得更多系统论文借鉴的透明度实践,明确说明 ChatGPT 被用于环境搭建、基准测试迁移、面向移动部署的 Kotlin 代码生成,以及批量实验脚本编写,作者对最终验证结果负责。

10. 结语

ZeroLock 是一篇难得地把一项真正新颖的理论结果(首次针对任意分块数量给出本地目标构造类 BP-free 训练的收敛性分析,速率与 BP 的 O(1/T)\mathcal{O}(1/\sqrt{T}) 只相差一个多重对数因子)与一个真实可用、跨平台的系统(多 GPU 服务器加未经改造的安卓手机)以及诚实、有机制支撑的实验(展示收益从哪里来,而不只是罗列收益存在)结合在一起的论文。它的核心想法——给每个流水线阶段配上自己的、由共享冻结读出头构造出的本地目标,而不是想办法围绕反向传播的链式法则依赖去做调度——在概念上非常干净,按理论来说,这在渐近意义下几乎是”免费”的。让这篇论文没能成为一场毫无保留的胜利的原因,在于理论的一般性与实证验证的狭窄性之间的落差:三个块,一个 11 亿参数模型,一个短序列分类风格的任务。如果后续工作能压力测试这篇论文留下的正是这几个开放维度——更多的块、更大的模型、更长程的任务,以及和其他 BP-free 本地目标方法的正面对比(目前论文只是概念性地讨论了它们)——ZeroLock 的核心想法看起来确实具备成为内存和通信受限流水线训练场景下一个真正的替代杠杆(与调度类气泡消减手段并行,而非互相取代)的潜力。