蘋果提出新型反向傳播：一臺(tái)iPhone 15 Pro Max就能微調(diào)LLM

機(jī)器之心報(bào)道

編輯：Panda

用 iPhone 本地跑大模型已經(jīng)不是新鮮事了，但能不能在 iPhone 上微調(diào)模型呢？

最近，蘋果親自上場(chǎng)，用一篇論文展示了其可行性。在這篇論文中，蘋果提出了一種內(nèi)存高效型反向傳播（MeBP）。該方法可在內(nèi)存使用量和計(jì)算時(shí)間之間提供比零階優(yōu)化（ZO/zeroth-order optimization）更好的權(quán)衡，同時(shí)還比 ZO 基線收斂更快、性能更優(yōu)。他們還在 iPhone 15 Pro Max 上驗(yàn)證了 MeBP 的有效性。

這個(gè)蘋果團(tuán)隊(duì)（宋叢崢與 Xinyu Tang）也在論文中表示會(huì)發(fā)布一個(gè) MeBP 實(shí)現(xiàn)，但其公開的鏈接目前還空無(wú)一碼。

• 論文標(biāo)題：Memory-Efficient Backpropagation for Fine-Tuning LLMs on Resource-Constrained Mobile Devices
• 論文地址：https://arxiv.org/abs/2510.03425
• 倉(cāng)庫(kù)地址：https://github.com/apple/ml-mebp

內(nèi)存高效型反向傳播（MeBP）

在這篇論文中，蘋果團(tuán)隊(duì)的研究重點(diǎn)是使用 LoRA 微調(diào) LLM。因此，主要的內(nèi)存瓶頸在于模型參數(shù)和中間激活值。該團(tuán)隊(duì)的目標(biāo)是將微調(diào)的內(nèi)存使用量保持在現(xiàn)代移動(dòng)設(shè)備可接受的范圍內(nèi)，例如 PocketLLM 所建議的「低于 1GB」。

使用 MeBP 在設(shè)備上微調(diào) LLM 包含三個(gè)步驟：

• 壓縮模型基礎(chǔ)權(quán)重（凍結(jié)的參數(shù)）以減少磁盤空間占用
• 編譯包含反向傳播和梯度檢查點(diǎn)的訓(xùn)練圖（training graph）以優(yōu)化內(nèi)存
• 實(shí)現(xiàn)一個(gè)內(nèi)存高效的運(yùn)行時(shí)（runtime）來(lái)執(zhí)行編譯后的訓(xùn)練圖。

下面將詳細(xì)描述每個(gè)步驟。

基礎(chǔ)模型權(quán)重壓縮

在設(shè)備上部署 LLM 時(shí)，壓縮基礎(chǔ)模型權(quán)重以減少磁盤空間使用是一種常見做法。

在該團(tuán)隊(duì)的實(shí)現(xiàn)中，他們對(duì)包括嵌入在內(nèi)的非 LoRA 參數(shù)使用了 4-bit 對(duì)稱模式 INT4 量化。

梯度檢查點(diǎn)編譯

也就是說(shuō)，反向圖的輸入是：已被檢查點(diǎn)的激活值、來(lái)自前一個(gè)檢查點(diǎn)的梯度、以及相應(yīng)的可訓(xùn)練權(quán)重；其輸出則是這些輸入的梯度。

隨后，所有塊的前向圖和反向圖被序列化為設(shè)備運(yùn)行時(shí)兼容的格式，例如模型中間語(yǔ)言（MIL）表示或 MLX 導(dǎo)出的函數(shù)。

在運(yùn)行時(shí)，這些序列化后的圖將被反序列化并編譯以進(jìn)行計(jì)算。

運(yùn)行時(shí)實(shí)現(xiàn)

算法 1 概述了 MeBP 的運(yùn)行時(shí)實(shí)現(xiàn)。

模型首先使用 InitializeModel 函數(shù)進(jìn)行初始化，之后訓(xùn)練循環(huán)中的每個(gè)數(shù)據(jù)點(diǎn)都會(huì)調(diào)用 Backpropagation 函數(shù)。在 InitializeModel 期間，壓縮后的基礎(chǔ)模型權(quán)重被內(nèi)存映射（memory-mapped）。為最小化內(nèi)存占用，基礎(chǔ)模型權(quán)重在訓(xùn)練循環(huán)開始前不會(huì)被解壓。相反，它們會(huì)在計(jì)算需要時(shí)才被按需（on demand）延遲解壓和加載。注意，對(duì)于支持使用量化權(quán)重進(jìn)行計(jì)算的設(shè)備運(yùn)行時(shí)框架，解壓步驟可以被跳過(guò)，屆時(shí)只需按需加載壓縮后的權(quán)重。

在 Backpropagation 函數(shù)中，系統(tǒng)首先執(zhí)行已編譯的前向子圖（subgraphs）以存儲(chǔ)所有必要的檢查點(diǎn)；隨后，按相反順序執(zhí)行已編譯的反向子圖，使用存儲(chǔ)的檢查點(diǎn)來(lái)計(jì)算梯度。在前向傳播過(guò)程中，這些檢查點(diǎn)被內(nèi)存映射，而不是保留在內(nèi)存中。

在每次前向和反向傳播之前，只有必需的基礎(chǔ)模型權(quán)重會(huì)被解壓和加載。如此一來(lái)，總內(nèi)存使用量被限制為：所需基礎(chǔ)模型權(quán)重的大小，加上每個(gè)子圖中操作的峰值內(nèi)存使用量。這個(gè)總和遠(yuǎn)小于基礎(chǔ)模型權(quán)重的完整大小。該函數(shù)描述的是單個(gè)數(shù)據(jù)點(diǎn)的梯度計(jì)算。對(duì)于批量輸入，可以使用梯度累積來(lái)計(jì)算梯度，而不會(huì)增加內(nèi)存占用。

在 MeBP 中，內(nèi)存中僅為優(yōu)化器保留一份 LoRA 權(quán)重及其梯度的副本。

對(duì)于參數(shù)量從 0.5B 到 4B 的 LLM，LoRA 權(quán)重的大小通常在幾十 MB 的范圍內(nèi)，這在內(nèi)存中存儲(chǔ)是合理的。優(yōu)化器狀態(tài)（例如動(dòng)量）可以像基礎(chǔ)模型權(quán)重一樣，被內(nèi)存映射并延遲加載。

實(shí)驗(yàn)表現(xiàn)如何？

MeBP 表現(xiàn)如何，還得看實(shí)踐，而作為對(duì)比的基線，他們選擇了 MeZO，因?yàn)樗悄壳耙阎奈ㄒ粦?yīng)用于移動(dòng)設(shè)備 LLM 微調(diào)的優(yōu)化方法。該團(tuán)隊(duì)通過(guò)服務(wù)器端的模擬來(lái)評(píng)估 MeZO 和 MeBP 的效用（utility），并在移動(dòng)設(shè)備上比較它們的性能。

效用（Utility）比較

配置上，這個(gè)蘋果團(tuán)隊(duì)使用了 Gemma-3 和 Qwen-2.5，在 WikiText-2 數(shù)據(jù)集上進(jìn)行語(yǔ)言建模任務(wù)實(shí)驗(yàn)，以此比較一階（FO）優(yōu)化（即通過(guò)反向傳播獲得梯度）和零階（ZO）優(yōu)化的效用。該團(tuán)隊(duì)專注于參數(shù)量不超過(guò) 4B 的模型，因?yàn)橐苿?dòng)設(shè)備的計(jì)算資源有限。該團(tuán)隊(duì)的評(píng)估指標(biāo)是評(píng)估集上的損失（loss）和下一 token 準(zhǔn)確度。其它配置見原論文，下面重點(diǎn)關(guān)注結(jié)果。

如圖 1 所示，盡管 ZO 的損失和下一 token 準(zhǔn)確度呈現(xiàn)收斂趨勢(shì)，但 ZO 的收斂速度明顯慢于 FO。FO 方法在最初的 100 步內(nèi)就顯著改善了這兩項(xiàng)指標(biāo)，而 ZO 在 1,000 步后僅顯示出輕微的改善。即使在 100,000 步之后（即比 FO 多 100 倍的優(yōu)化步數(shù)），對(duì)于同一模型，ZO 的測(cè)試損失仍然高于 FO，測(cè)試準(zhǔn)確度也低于 FO。

目前 AI 社區(qū)已經(jīng)提出了幾種方法，可以改善 ZO 方法的收斂速度。該團(tuán)隊(duì)也在 Qwen2.5-0.5B 上使用了這些改進(jìn)版 ZO 方法進(jìn)行實(shí)驗(yàn)，結(jié)果見下圖。

盡管這些方法比「純」 ZO 收斂得更快，但其損失和下一 token 準(zhǔn)確度仍然劣于使用 FO 微調(diào)的模型。此外，這些方法通常每次迭代需要更多的計(jì)算時(shí)間，因?yàn)樗鼈冃枰~外的前向傳播來(lái)更準(zhǔn)確地估計(jì)梯度。

效用結(jié)果表明，在語(yǔ)言建模任務(wù)的 LLM 微調(diào)上，按「每一步」（per-step）來(lái)看，反向傳播的收斂速度明顯快于 ZO 方法。這使得它在計(jì)算時(shí)間方面更適合移動(dòng)部署 —— 前提是每個(gè) FO 優(yōu)化步驟都能被高效地實(shí)現(xiàn)。

性能比較

蘋果使用 Swift 在 iOS 中實(shí)現(xiàn)了 MeBP，并在配備 8GB RAM 的 iPhone 15 Pro Max 上評(píng)估了其性能。對(duì)于 MeZO 基線實(shí)現(xiàn)，其前向圖被拆分為多個(gè)子圖，并應(yīng)用了延遲解壓來(lái)減少基礎(chǔ)模型權(quán)重的總內(nèi)存使用。每個(gè) MeZO 優(yōu)化步驟涉及兩次前向傳播。其它設(shè)置見原論文。

結(jié)果見下表。

總體而言，與 MeZO 相比，MeBP 每個(gè)梯度步驟的計(jì)算時(shí)間要多 43% 到 94%。但是，正如前面的效用對(duì)比所示，MeZO 所需的步數(shù)是一階優(yōu)化的 10 倍到 100 倍以上，因此在時(shí)間方面，MeBP 的收斂速度要快得多。在最壞情況下，MeBP 的內(nèi)存使用量比 MeZO 多出 20%，但其總訓(xùn)練內(nèi)存使用量比以往的移動(dòng)設(shè)備實(shí)現(xiàn)大約小 10 倍。所有測(cè)試的 LLM 均可在 1GB 內(nèi)存內(nèi)高效微調(diào)，使其適合在手機(jī)上進(jìn)行后臺(tái)訓(xùn)練。

此外，該團(tuán)隊(duì)還測(cè)試了解壓開銷與序列長(zhǎng)度的影響，并還分析了每一層的性能；詳見原論文。