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

機(jī)器之心編輯部

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

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

圖注：推特 @Shishir Patil

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

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

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

論文地址：https://arxiv.org/pdf/2207.07697.pdf

項(xiàng)目主頁(yè)：https://poet.cs.berkeley.edu/

GitHub 地址：https://github.com/shishirpatil/poet

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

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

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

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

集成分頁(yè)和重新實(shí)現(xiàn)

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

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

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

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

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

POET

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

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

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

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

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

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

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

其中 Φ_compute、Φ_pagein 和 Φ_pageout 分別表示每個(gè)節(jié)點(diǎn)在計(jì)算、page-in 和 page-out 時(shí)所消耗的能量。

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

實(shí)驗(yàn)結(jié)果

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

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

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

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

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

圖 4 強(qiáng)調(diào)了 POET 在不同時(shí)間約束下采用集成策略的好處。對(duì)于每個(gè)運(yùn)行時(shí)，下圖繪制了總能耗圖。