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

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

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

圖注:Twitter @Shishir Patil

圖注:Twitter @Shishir Patil

因此,為了使使用者在不犧牲隱私的情況下個性化他們的模型,聯邦學習等基於設備的訓練方法不需要將資料整合到雲端,也能執行本地訓練更新。這些方法已被部署在谷歌 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 機器學習模型的訓練進行最佳化。

對於部署在真實世界邊緣設備上的模型,當邊緣設備出現空閒並可以計算週期時就會進行訓練,例如谷歌 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 在不同時間約束下采用集成策略的好處。對於每個運行時,下圖繪製了總能耗圖。

相關文章

使用 TVMC 編譯和最佳化模型

使用 TVMC 編譯和最佳化模型

內容一覽:本節講解使用 TVMC 編譯和最佳化模型。TVMC 是 TVM 的命令驅動程序,通過命令列執行 TVM 功能。本節是了解 TVM ...

ChatGPT 的「神功」,是如何煉成的?

ChatGPT 的「神功」,是如何煉成的?

最強對話式 AI,站在巨人的肩膀上。 AIGC 在這幾個月成了一個大熱的話題。 頗有些風水輪流轉的感覺,如同年初大火特火的 web3 一樣,...