当前位置:首页 > 科技资讯 > 正文

苹果MeBP技术:iPhone上高效微调大型语言模型

在iPhone上本地运行大型模型已不是新鲜事,但能否在iPhone上进行模型微调呢?

最近,苹果亲自上阵,通过一篇论文展示了其可行性。在这篇论文中,苹果提出了一种内存高效型反向传播(MeBP)。该方法在内存使用量和计算时间之间提供了比零阶优化(ZO)更好的权衡,同时比ZO基线收敛更快、性能更优。他们还在iPhone 15 Pro Max上验证了MeBP的有效性。

这个由宋丛峥与Xinyu Tang组成的苹果团队也在论文中表示会发布一个MeBP实现,但其公开的链接目前还空无一码。

苹果MeBP技术:iPhone上高效微调大型语言模型 MeBP iPhone 模型微调 高效反向传播 第1张

论文标题: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上执行自动微分:

苹果MeBP技术:iPhone上高效微调大型语言模型 MeBP iPhone 模型微调 高效反向传播 第2张

其中E表示最终要优化的损失。接着,可以生成一个反向图

苹果MeBP技术:iPhone上高效微调大型语言模型 MeBP iPhone 模型微调 高效反向传播 第3张

,其中⊙表示哈达玛积(Hardmard product),而

苹果MeBP技术:iPhone上高效微调大型语言模型 MeBP iPhone 模型微调 高效反向传播 第4张

由反向图B_{i+1}输出。

也就是说,反向图的输入是:已被检查点的激活值、来自前一个检查点的梯度、以及相应的可训练权重;其输出则是这些输入的梯度。

随后,所有块的前向图和反向图被序列化为设备运行时兼容的格式,例如模型中间语言(MIL)表示或MLX导出的函数。

在运行时,这些序列化后的图将被反序列化并编译以进行计算。

运行时实现

算法 1 概述了MeBP的运行时实现。

苹果MeBP技术:iPhone上高效微调大型语言模型 MeBP iPhone 模型微调 高效反向传播 第5张

模型首先使用InitializeModel函数进行初始化,之后训练循环中的每个数据点都会调用Backpropagation函数。在InitializeModel期间,压缩后的基础模型权重被内存映射(memory-mapped)。为最小化内存占用,基础模型权重在训练循环开始前不会被解压。相反,它们会在计算需要时才被按需(on demand)延迟解压和加载。注意,对于支持使用量化权重进行计算的设备运行时框架,解压步骤可以被跳过,届时只需按需加载压缩后的权重。

在Backpropagation函数中,系统首先执行已编译的前向子图(subgraphs)以存储所有必要的检查点;随后,按相反顺序执行已编译的反向子图,使用存储的检查点来计算梯度。在前向传播过程中,这些检查点被内存映射,而不是保留在内存中。

在每次前向和反向传播之前,只有必需的基础模型权重会被解压和加载。如此一来,总内存使用量被限制为:所需基础模型权重的大小,加上每个子图中操作的峰值内存使用量。这个总和远小于基础模型权重的完整大小。该函数描述的是单个数据点的梯度计算。对于批量输入,可以使用梯度累积来计算梯度,而不会增加内存占用。

在MeBP中,内存中仅为优化器保留一份LoRA权重及其梯度的副本。

实验表现如何?

MeBP表现如何,还得看实践,而作为对比的基线,他们选择了MeZO,因为它是目前已知的唯一应用于移动设备LLM微调的优化方法。该团队通过服务器端的模拟来评估MeZO和MeBP的效用(utility),并在移动设备上比较它们的性能。

效用(Utility)比较