苹果提出新型反向传播:一台iPhone 15 Pro Max就能微调LLM
用 iPhone 本地跑大模型已经不是新鲜事了,但能不能在 iPhone 上微调模型呢?
最近,苹果亲自上场,用一篇论文展示了其可行性。在这篇论文中,苹果提出了一种内存高效型反向传播(MeBP)。该方法可在内存使用量和计算时间之间提供比零阶优化(ZO/zeroth-order optimization)更好的权衡,同时还比 ZO 基线收敛更快、性能更优。他们还在 iPhone 15 Pro Max 上验证了 MeBP 的有效性。
这个苹果团队(宋丛峥与 Xinyu Tang)也在论文中表示会发布一个 MeBP 实现,但其公开的链接目前还空无一码。
论文标题:Memory-Efficient Backpropagation for Fine-Tuning LLMs on Resource-Constrained Mobile Devices
论文地址:https://arxiv.org/abs/2510.03425
仓库地址:https://github.com/apple/ml-mebp
内存高效型反向传播(MeBP)
在这篇论文中,苹果团队的研究重点是使用 LoRA 微调 LLM。因此,主要的内存瓶颈在于模型参数和中间激活值。该团队的目标是将微调的内存使用量保持在现代移动设备可接受的范围内,例如 PocketLLM 所建议的「低于 1GB」。
使用 MeBP 在设备上微调 LLM 包含三个步骤:
压缩模型基础权重(冻结的参数)以减少磁盘空间占用
编译包含反向传播和梯度检查点的训练图(training graph)以优化内存
实现一个内存高效的运行时(runtime)来执行编译后的训练图。
下面将详细描述每个步骤。
基础模型权重压缩
在设备上部署 LLM 时,压缩基础模型权重以减少磁盘空间使用是一种常见做法。
在该团队的实现中,他们对包括嵌入在内的非 LoRA 参数使用了 4-bit 对称模式 INT4 量化。
梯度检查点编译
为在 MeBP 中实现梯度检查点,该团队首先将 LLM 拆分为多个块,确保对单个块(例如一个 transformer 层)执行反向传播的内存消耗在设备内存限制之内。对于每个产生待检查点的激活值的块 F,通过对 F 输出应用自动微分来生成反向图(backward graph)。例如,假设 y = F_i (x, w)是块 F_i 的前向图,则在标量 s 上执行自动微分:
其中 E 表示最终要优化的损失。接着,可以生成一个反向图
,其中 ⊙ 表示哈达玛积(Hardmard product),而
由反向图 B_{i+1} 输出。
也就是说,反向图的输入是:已被检查点的激活值、来自前一个检查点的梯度、以及相应的可训练权重;其输出则是这些输入的梯度。
随后,所有块的前向图和反向图被序列化为设备运行时兼容的格式,例如模型中间语言(MIL)表示或 MLX 导出的函数。
在运行时,这些序列化后的图将被反序列化并编译以进行计算。
运行时实现
算法 1 概述了 MeBP 的运行时实现。
模型首先使用 InitializeModel 函数进行初始化,之后训练循环中的每个数据点都会调用 Backpropagation 函数。在 InitializeModel 期间,压缩后的基础模型权重被内存映射(memory-mapped)。为最小化内存占用,基础模型权重在训练循环开始前不会被解压。相反,它们会在计算需要时才被按需(on demand)延迟解压和加载。注意,对于支持使用量化权重进行计算的设备运行时框架,解压步骤可以被跳过,届时只需按需加载压缩后的权重。
在 Backpropagation 函数中,系统首先执行已编译的前向子图(subgraphs)以存储所有必要的检查点;随后,按相反顺序执行已编译的反向子图,使用存储的检查点来计算梯度。在前向传播过程中,这些检查点被内存映射,而不是保留在内存中。
在每次前向和反向传播之前,只有必需的基础模型权重会被解压和加载。如此一来,总内存使用量被限制为:所需基础模型权重的大小,加上每个子图中操作的峰值内存使用量。这个总和远小于基础模型权重的完整大小。该函数描述的是单个数据点的梯度计算。对于批量输入,可以使用梯度累积来计算梯度,而不会增加内存占用。
在 MeBP 中,内存中仅为优化器保留一份 LoRA 权重及其梯度的副本。
对于参数量从 0.5B 到 4B 的 LLM,LoRA 权重的大小通常在几十 MB 的范围内,这在内存中存储是合理的。优化器状态(例如动量)可以像基础模型权重一样,被内存映射并延迟加载。
实验表现如何?
MeBP 表现如何,还得看实践,而作为对比的基线,他们选择了 MeZO,因为它是目前已知的唯一应用于移动设备 LLM 微调的优化方法。该团队通过服务器端的模拟来评估 MeZO 和 MeBP 的效用(utility),并在移动设备上比较它们的性能。
效用(Utility)比较
配置上,这个苹果团队使用了 Gemma-3 和 Qwen-2.5,在 WikiText-2 数据集上进行语言建模任务实验,以此比较一阶(FO)优化(即通过反向传播获得梯度)和零阶(ZO)优化的效用。该团队专注于参数量不超过 4B 的模型,因为移动设备的计算资源有限。该团队的评估指标是评估集上的损失(loss)和下一 token 准确度。其它配置见原论文,下面重点关注结果。
如图 1 所示,尽管 ZO 的损失和下一 token 准确度呈现收敛趋势,但 ZO 的收敛速度明显慢于 FO。FO 方法在最初的 100 步内就显著改善了这两项指标,而 ZO 在 1,000 步后仅显示出轻微的改善。即使在 100,000 步之后(即比 FO 多 100 倍的优化步数),对于同一模型,ZO 的测试损失仍然高于 FO,测试准确度也低于 FO。
目前 AI 社区已经提出了几种方法,可以改善 ZO 方法的收敛速度。该团队也在 Qwen2.5-0.5B 上使用了这些改进版 ZO 方法进行实验,结果见下图。
尽管这些方法比「纯」 ZO 收敛得更快,但其损失和下一 token 准确度仍然劣于使用 FO 微调的模型。此外,这些方法通常每次迭代需要更多的计算时间,因为它们需要额外的前向传播来更准确地估计梯度。
效用结果表明,在语言建模任务的 LLM 微调上,按「每一步」(per-step)来看,反向传播的收敛速度明显快于 ZO 方法。这使得它在计算时间方面更适合移动部署 —— 前提是每个 FO 优化步骤都能被高效地实现。
性能比较
苹果使用 Swift 在 iOS 中实现了 MeBP,并在配备 8GB RAM 的 iPhone 15 Pro Max 上评估了其性能。对于 MeZO 基线实现,其前向图被拆分为多个子图,并应用了延迟解压来减少基础模型权重的总内存使用。每个 MeZO 优化步骤涉及两次前向传播。其它设置见原论文。
结果见下表。
总体而言,与 MeZO 相比,MeBP 每个梯度步骤的计算时间要多 43% 到 94%。但是,正如前面的效用对比所示,MeZO 所需的步数是一阶优化的 10 倍到 100 倍以上,因此在时间方面,MeBP 的收敛速度要快得多。在最坏情况下,MeBP 的内存使用量比 MeZO 多出 20%,但其总训练内存使用量比以往的移动设备实现大约小 10 倍。所有测试的 LLM 均可在 1GB 内存内高效微调,使其适合在手机上进行后台训练。
此外,该团队还测试了解压开销与序列长度的影响,并还分析了每一层的性能;详见原论文。
本文来自微信公众号“机器之心”(ID:almosthuman2014),作者:机器之心,编辑:Panda,36氪经授权发布。















