ZeroLock：基于模块化更新解耦的并发内存高效LLM训练

# ZeroLock：通过模块化更新解耦实现并发且内存高效的大语言模型训练

来源：https://arxiv.org/html/2608.07974
Wentao Dai∗、Xuanran Li∗、Yuxiang Zhang、Ming Tang、Chao HuangWentao Dai、Yuxiang Zhang 和 Ming Tang 就职于中国南方科技大学计算机科学与工程系。Xuanran Li 就职于中国南方科技大学数学系。Chao Huang 就职于美国新泽西州蒙特克莱尔州立大学计算学院。邮箱：\{12311217,12312110,12410823\}@mail\.sustech\.edu\.cn, tangm3@sustech\.edu\.cn, huangch@montclair\.edu\.（通讯作者：Ming Tang）∗ 同等贡献

###### 摘要

边缘侧的大语言模型（LLM）微调在保护隐私的同时使模型适应特定场景的数据。尽管已有研究提出流水线并行来应对边缘设备有限的内存和计算资源，但它们通常依赖反向传播（BP）训练，而 BP 存在更新锁定的根本性局限，并可能遭遇严重的吞吐量和内存瓶颈。在本工作中，我们提出一种无 BP 算法，称为 ZeroLock，它通过局部目标构造将模型更新解耦为独立的块更新。它打破了 BP 的更新锁定，从而能在算法层面提升吞吐量，并通过减少激活存储来降低内存占用。据我们所知，我们首次为这种基于局部目标构造的方法在一般模型分块下建立了理论框架，将局部目标映射到全局目标。我们证明 ZeroLock 的收敛率为O~\(1/T\)\tilde\{\mathcal\{O\}\}\(1/\sqrt\{T\}\)，与 BP 相比仅相差多对数因子。我们为 ZeroLock 设计了一个系统，并构建了真实原型，结合了早期转发和故障恢复等技术以实现高效且鲁棒的实现。在原型上的实验表明，与基于 BP 的基线相比，ZeroLock 将内存降低了 26.5%，并将吞吐量提高了 4.9%。

## I引言

边缘智能在网络边缘部署大语言模型（LLM），实现隐私保护并提供实时推理服务。在实践中，在边缘侧微调 LLM 非常重要，原因有二。首先，数据分布通常具有场景和个体特异性，因此需要适配模型以保持推理精度。其次，数据隐私问题以及可能频繁的模型适配需求，使得将数据上传到云端进行微调变得困难。有许多这样的例子同时需要模型适配和隐私保护。例如，测试时训练[27 (https://arxiv.org/html/2608.07974#bib.bib41)]通过适配 LLM 以维护用户特定的记忆，从而提供个性化服务。基于脑电图信号的任务[2 (https://arxiv.org/html/2608.07974#bib.bib45)]（例如情绪识别、睡眠分期）需要对个体进行模型适配，因为这些信号表现出强烈的个体依赖模式。

然而，网络边缘的设备通常具有有限的计算和内存资源，因此单个设备或 GPU 很难微调整个 LLM。为了解决这个问题，已有研究提出了*流水线并行*方法[13 (https://arxiv.org/html/2608.07974#bib.bib10),20 (https://arxiv.org/html/2608.07974#bib.bib11),6 (https://arxiv.org/html/2608.07974#bib.bib1)]。其主要思想是将模型垂直划分为连续的层块。每个模型块的更新对应一个*阶段*，并被分配到不同的设备，这些设备以流水线方式协作更新各块。GPipe[13 (https://arxiv.org/html/2608.07974#bib.bib10)]是一种典型方法，按顺序处理微批次。1F1B[20 (https://arxiv.org/html/2608.07974#bib.bib11)]将前向传播和反向更新交错进行，以提高流水线利用率。在 1F1B 的基础上，PipeDream[20 (https://arxiv.org/html/2608.07974#bib.bib11),21 (https://arxiv.org/html/2608.07974#bib.bib12)]通过权重暂存实现了异步流水线训练。近期研究[7 (https://arxiv.org/html/2608.07974#bib.bib13),22 (https://arxiv.org/html/2608.07974#bib.bib14),16 (https://arxiv.org/html/2608.07974#bib.bib15),24 (https://arxiv.org/html/2608.07974#bib.bib16)]通过放置搜索、虚拟或双向阶段以及更细粒度的反向分解来减少流水线气泡。其他研究[17 (https://arxiv.org/html/2608.07974#bib.bib53),8 (https://arxiv.org/html/2608.07974#bib.bib54),29 (https://arxiv.org/html/2608.07974#bib.bib55),5 (https://arxiv.org/html/2608.07974#bib.bib52),25 (https://arxiv.org/html/2608.07974#bib.bib20)]则关注考虑抢占、落后者、设备异构性或通信效率的流水线规划。Confidant[6 (https://arxiv.org/html/2608.07974#bib.bib1)]在智能手机上实现了 1F1B。

参见说明

\(a\)

参见说明

\(b\) 参见说明参见说明

\(c\) \(d\)

图 1：\(a\) BP 基线和 \(b\) ZeroLock 的吞吐量；\(c\) BP 基线和 \(d\) ZeroLock 的内存。在图中，我们使用典型的 1F1B 框架作为 BP 基线，因为许多现有研究（例如[6 (https://arxiv.org/html/2608.07974#bib.bib1)]）都建立在它之上。S0、S1 和 S2 是阶段，每个阶段对应一个模型块的更新，并分配给一个设备执行。Fnn、Bnn 和 Annn分别表示第nn个微批次的前向传播、反向传播和激活。\(c\) 和 \(d\) 中的条纹表示相应激活在内存中保留的持续时间。例如，在 S0 的时间槽 3，\(c\) 中 A1-A3 被保留在内存中，而 \(d\) 中仅 A2 被保留在内存中。这些流水线并行工作（例如[13 (https://arxiv.org/html/2608.07974#bib.bib10),20 (https://arxiv.org/html/2608.07974#bib.bib11),21 (https://arxiv.org/html/2608.07974#bib.bib12),7 (https://arxiv.org/html/2608.07974#bib.bib13),22 (https://arxiv.org/html/2608.07974#bib.bib14),16 (https://arxiv.org/html/2608.07974#bib.bib15),24 (https://arxiv.org/html/2608.07974#bib.bib16),17 (https://arxiv.org/html/2608.07974#bib.bib53),8 (https://arxiv.org/html/2608.07974#bib.bib54),29 (https://arxiv.org/html/2608.07974#bib.bib55),6 (https://arxiv.org/html/2608.07974#bib.bib1),5 (https://arxiv.org/html/2608.07974#bib.bib52),4 (https://arxiv.org/html/2608.07974#bib.bib31),25 (https://arxiv.org/html/2608.07974#bib.bib20)]）大多依赖反向传播（BP）训练，并专注于系统级调度优化。

BP 训练具有更新锁定[14 (https://arxiv.org/html/2608.07974#bib.bib47),30 (https://arxiv.org/html/2608.07974#bib.bib8),31 (https://arxiv.org/html/2608.07974#bib.bib46)]的根本局限。也就是说，模型更新包含前向传播和随后的反向传播，因此上游层的更新需要等待下游层的前向和反向计算。因此，基于 BP 的方法可能存在关键局限：

- • 吞吐量瓶颈：如图 1 (https://arxiv.org/html/2608.07974#S1.F1)\(a\) 所示，由于更新锁定，上游阶段（例如 S0）需要等待所有下游阶段（例如 S1、S2）的反向传播完成后，才能完成自身的反向传播，从而产生气泡。尽管现有工作（例如[20 (https://arxiv.org/html/2608.07974#bib.bib11),21 (https://arxiv.org/html/2608.07974#bib.bib12)]）提出了消除气泡的方法，但由于 BP 训练，各阶段的更新仍然耦合，因此根本问题尚未解决。
- • 内存瓶颈：如图 1 (https://arxiv.org/html/2608.07974#S1.F1)\(c\) 所示，由于更新锁定，上游阶段（例如 S0）需要保留其激活，直到其所有下游阶段及其自身（例如 S0、S1、S2）完成反向传播。这导致上游阶段产生显著的内存浪费。

BP 训练的普遍使用以及由此产生的吞吐量和内存瓶颈引出了主要问题：

###### 主要问题。

我们如何设计一个系统，通过打破 BP 中的更新锁定，在算法层面从根本上克服吞吐量和内存瓶颈？

无 BP 训练被提出用于打破 BP 中的更新锁定。主要有两类。（I）反向梯度估计。例如，直接反馈对齐[23 (https://arxiv.org/html/2608.07974#bib.bib2),1 (https://arxiv.org/html/2608.07974#bib.bib3)]通过将目标误差直接传播到所有层来绕过链式法则。零阶优化[18 (https://arxiv.org/html/2608.07974#bib.bib4),28 (https://arxiv.org/html/2608.07974#bib.bib5)]在前向传播期间计算基于扰动的损失差异来估计梯度。然而，这类方法通常会导致复杂任务的模型精度显著下降或计算开销显著增加。（II）目标重构。例如，NoProp[15 (https://arxiv.org/html/2608.07974#bib.bib6)]将神经模块重新构建为独立的去噪单元，将带噪声的目标嵌入映射回干净目标。预测编码[19 (https://arxiv.org/html/2608.07974#bib.bib7)]在最小化逐块预测误差和使用局部误差更新权重之间交替进行。局部目标构造（例如深度渐进单调学习[30 (https://arxiv.org/html/2608.07974#bib.bib8)]）为每个层配备独立的局部目标，允许纯局部的梯度计算。在类别（II）中，局部目标构造将跨阶段的模块更新解耦，并且具有与 BP 方法相当的模型精度，因此它是解决主要问题的有前途的候选方案。具体来说，由于每个块使用其局部损失独立更新，其更新不需要等待其下游块的前向和反向传播，因此可以消除等待导致的气泡；同时，每个块只需要将激活保留到其自身更新完成，从而降低内存使用。

在本工作中，基于局部目标构造的无 BP 方法，我们旨在提出一种流水线并行 LLM 训练系统，以从根本上克服吞吐量和内存瓶颈。然而，这一设计并非直接可行，需要解决以下问题：

- Q1如何利用局部目标构造来微调 LLM？
- Q2局部目标构造的模块化解耦是否会在理论上损害模型收敛？
- Q3如何设计和构建真实原型系统（同时面向多 GPU 服务器和 Android 手机场景），以实现高效且鲁棒的实现？

回答 Q1 需要具体设计，将 LLM 特有特征（例如 token 序列、用于高效微调的低秩适配）融入局部目标构造框架。回答 Q2 非常具有挑战性。这是因为模块化更新解耦将全局目标分解为模型块的局部目标，而近期研究缺乏量化优化局部目标如何改进全局目标的框架。回答 Q3 也非易事，因为需要满足轻量、高吞吐量和鲁棒实现的真实工程要求。

现有工作尝试解决 Q1–Q3 中的一个或两个。PPLL[9 (https://arxiv.org/html/2608.07974#bib.bib39)]将梯度隔离的视觉块放置在不同的 GPU 上，并传输其特征以进行逐块更新。FluidPipe[3 (https://arxiv.org/html/2608.07974#bib.bib40)]专注于 LLM，并在两阶段模型的第一部分添加辅助任务头，避免第一部分的局部更新等待第二部分的梯度。SCPL[11 (https://arxiv.org/html/2608.07974#bib.bib42)]通过逐段监督对比损失解耦 BP，用于同步的多 GPU 模型并行。然而，这些工作[9 (https://arxiv.org/html/2608.07974#bib.bib39),3 (https://arxiv.org/html/2608.07974#bib.bib40),11 (https://arxiv.org/html/2608.07974#bib.bib42)]在回答 Q1 时考虑了隐藏状态对齐，未能在局部目标构造中刻画任务特定信息，并且没有解决 Q2。对于 Q3，尽管他们提出了高层流水线逻辑，但其设计仍停留在概念层面，缺乏关键实现细节，例如连续跨批次流水线、阶段 I/O 队列和显式 RPC 原语。同时，他们没有提供可部署的原型用于真实执行。尽管近期工作 LoPT[26 (https://arxiv.org/html/2608.07974#bib.bib48)]解决了 Q2，但其分析仅限于两个块，未能提供连接局部目标与全局目标的通用分析框架。

我们应对这些挑战并回答 Q1–Q3。我们的主要贡献总结如下：

- • ZeroLock 算法：为了回答 Q1，我们提出一种基于局部目标构造的无 BP 算法，称为 ZeroLock，用于 LLM 微调，该算法融入了低秩适配（LoRA）[12 (https://arxiv.org/html/2608.07974#bib.bib44)]以及 LLM 兼容的读出头和损失函数。它打破了模型块之间的更新耦合，从而在算法层面减少了流水线气泡和激活存储。
- • 理论分析：为了回答 Q2，我们建立了理论等价性，将模型块的局部最优表示为全局目标的形式。然后，模型块的局部更新可以等价地表示为全局更新，并由此推导全局收敛。据我们所知，这是第一个针对一般模型分块下局部目标构造算法的分析框架，为在解耦局部更新下分析全局收敛提供了系统方法。我们证明 ZeroLock 的收敛率为O~\(1/T\)\tilde\{\mathcal\{O\}\}\(1/\sqrt\{T\}\)，与 BP 相比仅相差多对数因子。
- • 真实系统设计：为了回答 Q3，基于 ZeroLock，我们为多 GPU 服务器和 Android 手机场景设计了系统并构建了原型。该系统可以实现并发吞吐量（见图 1 (https://arxiv.org/html/2608.07974#S1.F1)\(b\)），并通过减少激活存储来降低内存使用（见图 1 (https://arxiv.org/html/2608.07974#S1.F1)\(d\)）。我们提出了早期转发和故障恢复等技术，以确保轻量、高吞吐量和鲁棒的实现。据我们所知，我们在 Android 系统上构建了第一个基于局部目标构造的无 BP 算法原型。
- • 真实原型评估：在多 GPU 服务器上的实验表明，与基于 BP 的基线相比，ZeroLock 系统将内存使用降低了 26.5%，并将吞吐量提高了 4.9%。在 Android 手机上微调 TinyLlama 的峰值 PSS 低于 4000 MiB，电池温度约为 37∘C，墙钟时间为 1644.1 秒，表明该实现在实际中是可行的。我们的代码可在https://anonymous.4open.science/r/unlock_trainer-105B获取。

本文其余部分组织如下。第II (https://arxiv.org/html/2608.07974#S2)节和第III (https://arxiv.org/html/2608.07974#S3)节分别介绍 ZeroLock 算法及其系统设计。实验在第IV (https://arxiv.org/html/2608.07974#S4)节。第V (https://arxiv.org/html/2608.07974#S5)节总结本工作。

## IIZeroLock 算法与分析

### II-A ZeroLock 算法

考虑一个具有嵌入操作E\(x\)E\(x\)和一组变换的 LLM。