繁體
  • 简体中文
  • 繁體中文

熱門資訊> 正文

蘋果提出新型反向傳播:一臺iPhone 15 Pro Max就能微調LLM

2025-10-30 11:01

用 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氪經授權發佈。

風險及免責提示:以上內容僅代表作者的個人立場和觀點,不代表華盛的任何立場,華盛亦無法證實上述內容的真實性、準確性和原創性。投資者在做出任何投資決定前,應結合自身情況,考慮投資產品的風險。必要時,請諮詢專業投資顧問的意見。華盛不提供任何投資建議,對此亦不做任何承諾和保證。