首次在智慧型手機上訓練BERT和ResNet，耗能降35%

研究者表示，他們將邊緣訓練視為一個最佳化問題，從而發現了在給定記憶體預算下實現最小能耗的最優調度。

目前，智慧型手機和嵌入式平台等邊緣裝置上已經廣泛部署深度學習模型來進行推理。其中，訓練仍然主要是在具有 GPU 等高通量加速器的大型雲端伺服器上完成。集中式雲端訓練模型需要將照片和按鍵等敏感資料從邊緣設備傳輸到雲端，從而犧牲了用戶隱私並導致了額外的資料移動成本。

圖註：推特@Shishir Patil

因此，為了使用戶在不犧牲隱私的情況下個性化他們的模型，聯邦學習等基於設備的訓練方法不需要將資料整合到雲端，也能執行本地訓練更新。這些方法已部署在Google Gboard 鍵盤上以個人化鍵盤建議，也被 iPhones 手機用來提升自動語音辨識。同時，目前基於設備的訓練方法不支援訓練現代架構和大模型。在邊緣設備上訓練更大的模型不可行，主要是有限的設備記憶體無法儲存反向傳播啟動。 ResNet-50 的單次訓練迭代所需的記憶體是推理的 200 倍以上。

過去工作提出的策略包括分頁到輔助記憶體和重新實現，以減少雲端訓練的記憶體佔用。但是，這些方法會顯著增加整體能耗。與分頁方法相關的資料傳輸通常需要比重計算資料更多的能量。隨著記憶體預算的縮減，重新實現會以 O(n^2 ) 的速度增加能耗。

在 UC 伯克利最近的一篇論文中，幾位研究者表明分頁和重新實現是高度互補的。透過對簡單操作重新實現，同時將複雜操作的結果分頁到快閃記憶體或 SD 卡等輔助記憶體上，他們能夠以最小的能耗擴展有效的記憶體容量。並且，透過這兩種方法的結合，研究者也證明了在移動級邊緣設備上訓練 BERT 等模型是可能的。透過將邊緣訓練看作一個最佳化問題，他們發現了在給定記憶體預算下實現最小能耗的最優調度。

• 論文網址：https://arxiv.org/pdf/2207.07697.pdf
• 專案首頁：https://poet.cs .berkeley.edu/
• GitHub 網址：https://github.com/shishirpatil/poet

##研究者提出了POET（Private Optimal Energy Training ），這是一種在記憶體受限邊緣設備上對現代神經網路進行能量最優訓練的演算法，其架構如下圖1 所示。鑑於為反向傳播快取所有啟動張量的成本極高，POET 對啟動進行最佳化分頁和重新實現，因而可以將記憶體消耗最高減少兩倍。他們將邊緣訓練問題重新表述為整數線性程規劃（ILP），發現可以透過解算器在 10 分鐘內將其求解到最優。

圖註：POET 在邊緣設備上對 SOTA 機器學習模型的訓練進行最佳化。

對於部署在真實世界邊緣設備上的模型，當邊緣設備出現空閒並可以計算週期時就會進行訓練，例如Google Gboard 會在手機充電時安排模型更新。因此，POET 也包含了嚴格的訓練限制。給定記憶體限制和訓練 epoch 的數量，POET 產生的解決方案也能滿足給定的訓練截止期限。此外，研究者也利用 POET 開發了一個全面的成本模型，並證明它在數學上是保值的（即不做近似），適用於現有的開箱即用架構。

論文一作 Shishir Patil 在示範影片中表示，POET 演算法可以在智慧型手機等商用邊緣設備上訓練任何需要極大記憶體的 SOTA 模型。他們也成為了首個展示在智慧型手機和 ARM Cortex-M 設備上訓練 BERT 和 ResNet 等 SOTA 機器學習模型的研究團隊。

集成分頁和重新實作

重新實作和分頁是降低大型 SOTA ML 模型記憶體消耗的兩種技術。在重新實作中，一旦不再需要啟動張量就會被刪除，最常見的是在前向傳播期間。從而釋放了寶貴的內存，可用於儲存後續層的激活。當再次需要刪除的張量時，該方法會根據譜系的規定從其他相關的活化中重新計算。而分頁，也稱為 offloading，是一種減少記憶體的補充技術。在分頁中，不是立即需要的啟動張量從主記憶體調出到二級記憶體，例如快閃記憶體或 SD 卡。當再次需要張量時，將其分頁。

圖 2 顯示了一個八層神經網路的執行時間表。沿著 X 軸，每個單元對應神經網路的每一層（共 8 層 L8）。 Y 軸表示一個 epoch 內的邏輯時間步長。圖中佔用的單元（以顏色填滿）表示在對應的時間步執行的操作（前向 / 後向傳播計算、重新實現或分頁）。

例如，我們可以看到 L1 的活化是在第一個時間步 (T1) 計算的。在 T2 和 T3 時刻，分別計算 L2 和 L3 的活化量。假設層 L2 和 L3 恰好是記憶體密集但計算成本較低的運算，例如非線性 (tanH、ReLU 等)，那麼重新實作就成為了最佳選擇。我們可以刪除激活({T3, L2}， {T4, L3}) 來釋放內存，當後向傳播過程中需要這些激活時，可以再重新實現它們({T14, L3}， {T16, L2}) 。

假設 L5 和 L6 層是計算密集型運算，例如卷積、密集矩陣乘法等。對於此類運算，重新實現將導致運行時間和能量的增加，並且這種方式是次優的。對於這些層，最好將啟動張量分頁到輔助儲存（{T6，L5}，{T7，L6}），並在需要時分頁到（{T10，L6}，{T11，L5 }）。

分頁的一個主要優點是，根據記憶體匯流排的佔用情況，可以進行 pipelin 處理，以隱藏延遲。這是因為現代系統具有 DMA（直接記憶體存取）特性，它可以在計算引擎並行運行時將啟動張量從輔助儲存移動到主記憶體。例如，在時間步 T7，可以同時將 L6 調出並計算 L7。但是，重新實作是計算密集型的，不能並行化，這導致運行時間增加。例如，我們必須將時間步 T14 用於重新計算 L3，從而延遲其餘反向傳播執行。

POET

該研究提出了POET，這是一種用於深度神經網路的圖形級編譯器，它重寫了大型模型的訓練DAG，以適應邊緣設備的內存限制，同時保持高能效。

POET 是硬體感知的，它首先追蹤前向和後向傳播的執行以及相關的記憶體分配請求、運行時間以及每次操作的記憶體和能源消耗。對於給定的硬件，每個工作負載的這種細粒度分析只發生一次，具有自動化、便宜等特性，並且為 POET 提供了最準確的成本模型。

POET 然後產生可以有效求解的混合整數線性規劃 (MILP)。 POET 優化器搜尋有效的重新實現和分頁調度，以最大限度地減少受記憶體限制的端到端能源消耗。然後使用所得的調度產生一個新的 DAG，在邊緣設備上執行。

雖然 MILP 是在商用硬體上解決的，但發送到邊緣設備的調度表只有幾百字節，因此記憶體效率很高。

對於計算成本低但記憶體密集型的操作，重新實作是最有效的。然而，分頁最適合計算密集型操作，在這種操作中，重新實作將導致大量的能量開銷。 POET 在一個整合搜尋空間中共同考慮重新實現和分頁。

本文方法可擴展到複雜、現實的架構中，POET 優化器演算法如下。

該研究在最佳化問題中引入了一個新的目標函數，以最小化計算、page-in 和page-out 的綜合能耗，分頁和重新實現能耗結合的新目標函數為：

#其中Φ_compute、Φ_pagein 和Φ_pageout 分別表示每個節點在計算、page-in 和page-out 時所消耗的能量。

POET 根據圖的哪些節點 (k) 進行了重新實現，以及在每個時間步長 (t) 將哪些節點 page-in 或 page-out 來輸出 DAG 調度。

 實驗結果

在對 POET 的評估中，研究者試圖回答三個關鍵問題。首先，POET 在不同的模型和平台上能夠減少多少能耗？其次，POET 如何從混合分頁和重新實現策略中獲益？最後，POET 如何適應不同的運行時預算？

研究者在下表 2 中列出四種不同的硬體設備，分別為 ARM Cortex M0 MKR1000、ARM Cortex M4F nrf52840、A72 Raspberry Pi 4B 和 Nvidia Jetson TX2。 POET 是完全硬體感知的，依賴於細粒度的分析。

下圖 3 顯示了單次訓練 epoch 的能耗，每列分別對應不同的硬體平台。研究者發現，POET 在所有平台上產生節能最優的調度（Y 軸），同時減少峰值記憶體消耗（X 軸）並符合時間預算。

在下圖 5 中，研究者在 A72 上訓練 ResNet-18 時對 POET 和 Capuchin 進行了基準測試。隨著 RAM 預算的減少，Capuchin 比具有完整記憶體的基準多了 73% 到 141% 的能耗。相比之下，POET 產生的能耗不到 1%。這種趨勢適用於測試的所有架構和平台。

表 3 中，研究在 Nvidia 的 Jetson TX2 上訓練 ResNet-18 時對 POET 和 POFO 進行了基準測試。研究發現 POET 找到了一個整合的重新實現和分頁調度，可將峰值記憶體消耗降低 8.3%，並將吞吐量提高 13%。這展示了 POET 的 MILP 求解器的優勢，它能夠在更大的搜尋空間上進行最佳化。雖然 POFO 僅支援線性模型，但 POET 可以推廣到非線性模型，如圖 3 所示。

圖 4 強調了 POET 在不同時間限制下採用整合策略的好處。對於每個運行時，下圖繪製了總能耗圖。

#

以上是首次在智慧型手機上訓練BERT和ResNet，耗能降35%的詳細內容。更多資訊請關注PHP中文網其他相關文章！