研究者表示，他們將邊緣訓(xùn)練看作一個優(yōu)化問題，從而發(fā)現(xiàn)了在給定內(nèi)存預(yù)算下實現(xiàn)最小能耗的最優(yōu)調(diào)度。??

目前，智能手機和嵌入式平臺等邊緣設(shè)備上已經(jīng)廣泛部署深度學(xué)習(xí)模型來進行推理。其中，訓(xùn)練仍然主要是在具有 GPU 等高通量加速器的大型云服務(wù)器上完成。集中式云訓(xùn)練模型需要將照片和按鍵等敏感數(shù)據(jù)從邊緣設(shè)備傳輸?shù)皆贫耍瑥亩鵂奚擞脩綦[私并導(dǎo)致了額外的數(shù)據(jù)移動成本。

圖注：推特 @Shishir Patil?

因此，為了使用戶在不犧牲隱私的情況下個性化他們的模型，聯(lián)邦學(xué)習(xí)等基于設(shè)備的訓(xùn)練方法不需要將數(shù)據(jù)整合到云端，也能執(zhí)行本地訓(xùn)練更新。這些方法已被部署在谷歌 Gboard 鍵盤上以個性化鍵盤建議，也被 iPhones 手機用來提升自動語音識別。同時，當(dāng)前基于設(shè)備的訓(xùn)練方法不支持訓(xùn)練現(xiàn)代架構(gòu)和大模型。在邊緣設(shè)備上訓(xùn)練更大的模型不可行，主要是有限的設(shè)備內(nèi)存無法存儲反向傳播激活。ResNet-50 的單次訓(xùn)練迭代所需的內(nèi)存是推理的 200 多倍。

以往工作提出的策略包括分頁到輔助內(nèi)存和重新實現(xiàn)，以減少云端訓(xùn)練的內(nèi)存占用。但是，這些方法會顯著增加整體能耗。與分頁方法相關(guān)的數(shù)據(jù)傳輸通常需要比重計算數(shù)據(jù)更多的能量。隨著內(nèi)存預(yù)算的縮減，重新實現(xiàn)會以 O(n^2 ) 的速度增加能耗。

在 UC 伯克利最近的一篇論文中，幾位研究者表明分頁和重新實現(xiàn)是高度互補的。通過對簡單操作重新實現(xiàn)，同時將復(fù)雜操作的結(jié)果分頁到閃存或 SD 卡等輔助存儲器上，他們能夠以最小的能耗擴展有效的內(nèi)存容量。并且，通過這兩種方法的結(jié)合，研究者還證明了在移動級邊緣設(shè)備上訓(xùn)練 BERT 等模型是可能的。通過將邊緣訓(xùn)練看作一個優(yōu)化問題，他們發(fā)現(xiàn)了在給定內(nèi)存預(yù)算下實現(xiàn)最小能耗的最優(yōu)調(diào)度。

• 論文地址：https://arxiv.org/pdf/2207.07697.pdf
• 項目主頁：https://poet.cs.berkeley.edu/
• GitHub 地址：https://github.com/shishirpatil/poet

研究者提出了 POET（Private Optimal Energy Training），這是一種在內(nèi)存受限邊緣設(shè)備上對現(xiàn)代神經(jīng)網(wǎng)絡(luò)進行能量最優(yōu)訓(xùn)練的算法，其架構(gòu)如下圖 1 所示。鑒于為反向傳播緩存所有激活張量的成本極高，POET 對激活進行優(yōu)化分頁和重新實現(xiàn)，因而可以將內(nèi)存消耗最高減少兩倍。他們將邊緣訓(xùn)練問題重新表述為整數(shù)線性程規(guī)劃（ILP），發(fā)現(xiàn)可以通過求解器在 10 分鐘內(nèi)將其求解到最優(yōu)。

圖注：POET 在邊緣設(shè)備上對 SOTA 機器學(xué)習(xí)模型的訓(xùn)練進行優(yōu)化。

對于部署在真實世界邊緣設(shè)備上的模型，當(dāng)邊緣設(shè)備出現(xiàn)空閑并可以計算周期時就會進行訓(xùn)練，例如谷歌 Gboard 會在手機充電時安排模型更新。因此，POET 也包含了嚴(yán)格的訓(xùn)練限制。給定內(nèi)存限制和訓(xùn)練 epoch 的數(shù)量，POET 生成的解決方案也能滿足給定的訓(xùn)練截止期限。此外，研究者還利用 POET 開發(fā)了一個全面的成本模型，并證明它在數(shù)學(xué)上是保值的（即不做近似），適用于現(xiàn)有的開箱即用架構(gòu)。

論文一作 Shishir Patil 在演示視頻中表示，POET 算法可以在智能手機等商用邊緣設(shè)備上訓(xùn)練任何需要極大內(nèi)存的 SOTA 模型。他們也成為了首個展示在智能手機和 ARM Cortex-M 設(shè)備上訓(xùn)練 BERT 和 ResNet 等 SOTA 機器學(xué)習(xí)模型的研究團隊。

集成分頁和重新實現(xiàn)

重新實現(xiàn)和分頁是降低大型 SOTA ML 模型內(nèi)存消耗的兩種技術(shù)。在重新實現(xiàn)中，一旦不再需要激活張量就會被刪除，最常見的是在前向傳播期間。從而釋放了寶貴的內(nèi)存，可用于存儲后續(xù)層的激活。當(dāng)再次需要刪除的張量時，該方法會根據(jù)譜系的規(guī)定從其他相關(guān)的激活中重新計算。而分頁，也稱為 offloading，是一種減少內(nèi)存的補充技術(shù)。在分頁中，不是立即需要的激活張量從主存儲器調(diào)出到二級存儲器，例如閃存或 SD 卡。當(dāng)再次需要張量時，將其分頁。

圖 2 顯示了一個八層神經(jīng)網(wǎng)絡(luò)的執(zhí)行時間表。沿著 X 軸，每個單元對應(yīng)神經(jīng)網(wǎng)絡(luò)的每一層（共 8 層 L8）。Y 軸表示一個 epoch 內(nèi)的邏輯時間步長。圖中占用的單元（用顏色填充）表示在相應(yīng)的時間步執(zhí)行的操作（前向 / 后向傳播計算、重新實現(xiàn)或分頁）。

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

假設(shè) L5 和 L6 層是計算密集型運算，例如卷積、密集矩陣乘法等。對于此類運算，重新實現(xiàn)將導(dǎo)致運行時間和能量的增加，并且這種方式是次優(yōu)的。對于這些層，最好將激活張量分頁到輔助存儲（{T6，L5}，{T7，L6}），并在需要時分頁到（{T10，L6}，{T11，L5 }）。

分頁的一個主要優(yōu)點是，根據(jù)內(nèi)存總線的占用情況，可以進行 pipelin 處理，以隱藏延遲。這是因為現(xiàn)代系統(tǒng)具有 DMA（直接內(nèi)存訪問）特性，它可以在計算引擎并行運行時將激活張量從輔助存儲移動到主內(nèi)存。例如，在時間步 T7，可以同時將 L6 調(diào)出并計算 L7。但是，重新實現(xiàn)是計算密集型的，不能并行化，這導(dǎo)致運行時間增加。例如，我們必須將時間步 T14 用于重新計算 L3，從而延遲其余反向傳播執(zhí)行。

POET

該研究提出了 POET，這是一種用于深度神經(jīng)網(wǎng)絡(luò)的圖形級編譯器，它重寫了大型模型的訓(xùn)練 DAG，以適應(yīng)邊緣設(shè)備的內(nèi)存限制，同時保持高能效。

POET 是硬件感知的，它首先跟蹤前向和后向傳播的執(zhí)行以及相關(guān)的內(nèi)存分配請求、運行時間以及每次操作的內(nèi)存和能源消耗。對于給定的硬件，每個工作負(fù)載的這種細(xì)粒度分析只發(fā)生一次，具有自動化、便宜等特性，并且為 POET 提供了最準(zhǔn)確的成本模型。

POET 然后生成可以有效求解的混合整數(shù)線性規(guī)劃 (MILP)。POET 優(yōu)化器搜索有效的重新實現(xiàn)和分頁調(diào)度，以最大限度地減少受內(nèi)存限制的端到端能源消耗。然后使用得到的調(diào)度生成一個新的 DAG，在邊緣設(shè)備上執(zhí)行。

雖然 MILP 是在商用硬件上解決的，但發(fā)送到邊緣設(shè)備的調(diào)度表只有幾百字節(jié)，因此內(nèi)存效率很高。

對于計算成本低但內(nèi)存密集型的操作，重新實現(xiàn)是最有效的。然而，分頁最適合于計算密集型操作，在這種操作中，重新實現(xiàn)將導(dǎo)致大量的能量開銷。POET 在一個集成搜索空間中共同考慮重新實現(xiàn)和分頁。

本文方法可擴展到復(fù)雜、現(xiàn)實的架構(gòu)中，POET 優(yōu)化器算法如下。

該研究在優(yōu)化問題中引入了一個新的目標(biāo)函數(shù)，以最小化計算、page-in 和 page-out 的綜合能耗，分頁和重新實現(xiàn)能耗結(jié)合的新目標(biāo)函數(shù)為：

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

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

 實驗結(jié)果?

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

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

下圖 3 顯示了單次訓(xùn)練 epoch 的能耗，每列分別對應(yīng)不同的硬件平臺。研究者發(fā)現(xiàn)，POET 在所有平臺上生成節(jié)能耗最優(yōu)的調(diào)度（Y 軸），同時減少峰值內(nèi)存消耗（X 軸）并符合時間預(yù)算。

在下圖 5 中，研究者在 A72 上訓(xùn)練 ResNet-18 時對 POET 和 Capuchin 進行了基準(zhǔn)測試。隨著 RAM 預(yù)算的減少，Capuchin 比具有完整內(nèi)存的基線多了 73% 到 141% 的能耗。相比之下，POET 產(chǎn)生的能耗不到 1%。這種趨勢適用于測試的所有架構(gòu)和平臺。

表 3 中，該研究在 Nvidia 的 Jetson TX2 上訓(xùn)練 ResNet-18 時對 POET 和 POFO 進行了基準(zhǔn)測試。研究發(fā)現(xiàn) POET 找到了一個集成的重新實現(xiàn)和分頁調(diào)度，可將峰值內(nèi)存消耗降低 8.3%，并將吞吐量提高 13%。這展示了 POET 的 MILP 求解器的優(yōu)勢，它能夠在更大的搜索空間上進行優(yōu)化。雖然 POFO 僅支持線性模型，但 POET 可以推廣到非線性模型，如圖 3 所示。

圖 4 強調(diào)了 POET 在不同時間約束下采用集成策略的好處。對于每個運行時，下圖繪制了總能耗圖。? ?