大規(guī)模神經(jīng)網(wǎng)絡(luò)最新文獻(xiàn)綜述：訓(xùn)練高效DNN、節(jié)省內(nèi)存使用、優(yōu)化器設(shè)計(jì)

選自arXiv

作者：Julia Gusak等

機(jī)器之心編譯

編輯：杜偉、澤南

在本綜述論文中，研究者解釋了不同技術(shù)的工作原理、評(píng)估和比較，還分析了一些實(shí)現(xiàn)這些技術(shù)的框架。

現(xiàn)代深度學(xué)習(xí)和人工智能技術(shù)的發(fā)展涉及使用深度神經(jīng)網(wǎng)絡(luò)（DNN）來解決圖像、視頻、音頻、自然語言處理、圖像形式的內(nèi)容生成等各種問題，或生成給定格式主題的文本等任務(wù)。

俄羅斯斯科爾科沃科學(xué)技術(shù)研究所、法國里爾大學(xué)、波爾多大學(xué)、Inria 等科研機(jī)構(gòu)聯(lián)合發(fā)表了一篇論文《Survey on Large Scale Neural Network Training》，它試圖解決的問題是：若給定模型和計(jì)算平臺(tái)的情形下，如何訓(xùn)練才是最有效率的。為了使訓(xùn)練高效，其必須可行，最大程度地利用資源的計(jì)算能力，在并行情況下，它不能讓信息傳輸成為瓶頸。訓(xùn)練的效率從根本上取決于計(jì)算內(nèi)核在計(jì)算資源（CPU、TPU、GPU）上的有效實(shí)現(xiàn)以及 GPU 之間和不同內(nèi)存之間通信的有效實(shí)現(xiàn)。

論文鏈接：https://arxiv.org/abs/2202.10435

在這兩種情況下，人們?yōu)閮?yōu)化計(jì)算內(nèi)核的算術(shù)強(qiáng)度，及有效實(shí)現(xiàn)硬件網(wǎng)絡(luò)上的通信做了很多工作。對(duì)于使用者來說，已存在強(qiáng)大的分析工具來識(shí)別硬件瓶頸，并可用于判定本調(diào)查中描述哪些策略可用于解決算術(shù)強(qiáng)度、內(nèi)存和控制交換數(shù)據(jù)量的問題。

該綜述研究涵蓋了應(yīng)對(duì)這些限制的通用技術(shù)。如果由于模型、優(yōu)化器狀態(tài)和激活不適合內(nèi)存而無法先驗(yàn)執(zhí)行計(jì)算，則可以使用內(nèi)存交換計(jì)算（重新實(shí)現(xiàn)）或數(shù)據(jù)轉(zhuǎn)移（激活和權(quán)重卸載）。我們還可以通過近似優(yōu)化器狀態(tài)和梯度（壓縮、修剪、量化）來壓縮內(nèi)存使用。

并行方法（數(shù)據(jù)并行、模型并行、流水線模型并行）也可以將內(nèi)存需求分布到多個(gè)算力資源上。如果計(jì)算的算力強(qiáng)度不足以充分利用 GPU 和 TPU，一般是因?yàn)?mini-batch 太小，那么上述技術(shù)也可以增加 mini-batch 的大小。最后，如果使用數(shù)據(jù)并行引起的通信開銷昂貴到拖累計(jì)算速度，則可以使用其他形式的并行（模型并行、流水線模型并行），梯度壓縮也可以限制數(shù)據(jù)交換的數(shù)量。

在本次調(diào)查中，研究者解釋了這些不同技術(shù)是如何工作的，其中描述了評(píng)估和比較所提出方法的文獻(xiàn)，還分析了一些實(shí)施這些技術(shù)的框架。

下表 1為文章討論的不同技術(shù)及其對(duì)通信、內(nèi)存和計(jì)算效率的影響。

研究者根據(jù)目的區(qū)分了以下方法：首先討論減少 GPU 內(nèi)存使用，隨后考慮對(duì)不適合 GPU 的模型使用并行訓(xùn)練，最后討論為訓(xùn)練存儲(chǔ)在多個(gè)設(shè)備上的模型而開發(fā)的優(yōu)化器的設(shè)計(jì)。

單 GPU 情況下減少內(nèi)存使用

在前向傳播期間，神經(jīng)網(wǎng)絡(luò)存儲(chǔ)執(zhí)行反向傳播所需的激活。在某些情況下，這些激活會(huì)消耗大量內(nèi)存，讓模型無法訓(xùn)練。減少內(nèi)存使用的主要方法有兩種：重新實(shí)現(xiàn)（也稱為 checkpointing）和卸載。

激活的重新實(shí)現(xiàn)

重新實(shí)現(xiàn)的策略僅在前向傳播期間存儲(chǔ)一小部分激活，并在反向傳播期間重新計(jì)算其余部分。重新實(shí)現(xiàn)方法可以通過它們處理的計(jì)算圖來區(qū)分。第一組來自自動(dòng)微分（AD），它們?yōu)橥瑯?gòu)順序網(wǎng)絡(luò)（多層按順序執(zhí)行并具有相同計(jì)算和內(nèi)存成本的 DNN）找到最佳調(diào)度。第二組專注于過渡模型，例如異構(gòu)序列網(wǎng)絡(luò)（可以是由任意復(fù)雜模塊組成的任何序列神經(jīng)網(wǎng)絡(luò)，如 CNN、ResNet、一些 transformer），它將解決方案從 AD 調(diào)整為異構(gòu)設(shè)置。

一些方法可以對(duì)一般計(jì)算圖執(zhí)行重新實(shí)現(xiàn)，盡管確切的計(jì)算成本可能指數(shù)級(jí)上升，如下表 2 所示。

激活卸載

卸載（又被稱為內(nèi)存交換）是一種通過在前向傳遞期間將激活轉(zhuǎn)移到 CPU 內(nèi)存并將它們預(yù)取回 GPU 內(nèi)存，以進(jìn)行相應(yīng)的向后計(jì)算來節(jié)省 GPU 內(nèi)存的技術(shù)。

由于 CPU 和 GPU 之間 PCI 總線的帶寬有限，必須優(yōu)化選擇傳輸激活，以及何時(shí)傳輸?shù)倪x擇。

在 vDNN [Rhu et al., 2016] 研究中，作者通過僅卸載卷積層的輸入來遵循對(duì) CNN 有效的啟發(fā)式方法，然而它不能很好地推廣到一般 DNN 上。另有研究 [Le et al., 2018] 考慮了激活生命周期來選擇卸載的內(nèi)容，并使用圖搜索方法來識(shí)別插入卸載 / 預(yù)取操作的時(shí)刻。AutoSwap [Zhang et al., 2019] 通過為每個(gè)變量分配優(yōu)先級(jí)分?jǐn)?shù)來決定卸載哪些激活。

權(quán)重卸載

前面提到的很多方法也適用于卸載權(quán)重，這是因?yàn)樾遁d權(quán)重依賴于適用于任何張量的通用技術(shù)，比如 TFLMS、AutoSwap 或者 SwapAdvisor。

不適合單個(gè) GPU 的模型的并行性

在模型并行化中，只需要傳達(dá)激活信息，并且傳輸只發(fā)生在分配給不同處理器的連續(xù)層之間。本章節(jié)提到的工作如下表 4 所示。

如果多個(gè)小批量被 pipeline 化 ，則可以加快模型并行化中的執(zhí)行速度，從而同時(shí)激活了多個(gè)訓(xùn)練迭代，具體可見 [Huang et al., 2019]。一旦在所有這些小批量上計(jì)算了前向和后向階段，權(quán)重就會(huì)更新。這種方法實(shí)現(xiàn)起來相當(dāng)簡單，但也導(dǎo)致計(jì)算資源大部分處于空置狀態(tài)。[Narayanan et al., 2019] 中提出的 PipeDream 方法僅強(qiáng)制前向和后向任務(wù)針對(duì)給定的小批量使用相同的模型權(quán)重，改進(jìn)了這一訓(xùn)練過程。

減少執(zhí)行更新的頻率也已被證明有助于限制權(quán)重過期（Narayanan et al., 2021a）。[Yang et al., 2021] 提出的 PipeMare 根據(jù) pipeline 階段向后調(diào)整學(xué)習(xí)率和模型權(quán)重。

對(duì) pipeline 方法中激活導(dǎo)致的存儲(chǔ)成本進(jìn)行建模是一項(xiàng)艱巨的任務(wù)（Beaumont et al., 2021b）。例如，[Fan et al., 2021] 中的 DAPPLE 、 [Li and Hoefler, 2021] 中的 Chimera 使用 1F1B（One-Forward-One-Backward）調(diào)度來減少與激活相關(guān)的內(nèi)存消耗。1F1B 是一種同步權(quán)重更新技術(shù)，盡可能早地安排每個(gè)微批次的反向傳遞，以釋放激活占用的內(nèi)存。

有些論文專門處理具有挑戰(zhàn)性的拓?fù)?。比如，為了解決高通信成本和異構(gòu)網(wǎng)絡(luò)能力的問題，[Zhan and Zhang, 2019] 中的 Pipe-torch 提出了一種更新的動(dòng)態(tài)規(guī)劃策略，該策略假設(shè)計(jì)算和通信之間沒有重疊。[Park et al., 2020] 中的 Pipe 解決了異構(gòu) GPU 的其他問題，采用的方法是將這些異構(gòu) GPU 分成虛擬 worker，并在每個(gè)虛擬 worker 中運(yùn)行 pipeline 并行化，同時(shí)依賴 worker 之間的數(shù)據(jù)并行化。

用于跨設(shè)備模型訓(xùn)練的優(yōu)化器

零冗余優(yōu)化器

2020 年， Rajbhandari, S. 等人在論文《 ZeRO: Memory Optimizations toward Training Trillion Parameter Models》中提出了零冗余優(yōu)化器（Zero Redundancy Optimizer, ZeRO），將它作為一種減少內(nèi)存使用的數(shù)據(jù)并行化實(shí)現(xiàn)。根據(jù)在設(shè)備上劃分的張量，該算法具有三個(gè)階段，即階段 1 - 優(yōu)化器狀態(tài)、階段 2 - 優(yōu)化器狀態(tài)和梯度和階段 3 - 優(yōu)化器狀態(tài)、梯度和模型超參數(shù)。

2021 年， Ren, J. 等人在論文《 ZeRO-Offload: Democratizing Billion-Scale Model Training》中將 ZeRO 與 Zero-Offload 內(nèi)部參數(shù)更新的 CPU 端計(jì)算統(tǒng)一起來，其中梯度被遷移至存儲(chǔ)參數(shù)副本的 CPU，更新的權(quán)重遷移回 GPU。

低精度優(yōu)化器

為了進(jìn)一步減少內(nèi)存使用，低精度優(yōu)化器（low-precision optimizer）有了用武之地。這些方法使用低精度格式拉力表示優(yōu)化器狀態(tài)以及狀態(tài)的輔助向量。并且，誤差補(bǔ)償技術(shù)可以被用來維持跟蹤統(tǒng)計(jì)的近似準(zhǔn)確率。

2021 年， Dean, J. 等人在論文《Large Scale Distributed Deep Networks》中提出了一種將 Adam 優(yōu)化器存儲(chǔ)在 8-bit 的方法，同時(shí)在使用 32-bit 格式時(shí)保持整體性能不變。2020 年， Sun, X. 等人在論文《Ultra-Low Precision 4-bit Training of Deep Neural Networks》中提出了更激進(jìn)的精度降低，其中開發(fā)了處理 4-bit 表示的特定路徑。

收斂加速

另一種加速大規(guī)模深度學(xué)習(xí)模型的方法是減少節(jié)點(diǎn)之間的通信時(shí)間以及在適當(dāng)局部最小值收斂所需的 epoch 數(shù)量。

關(guān)于通信成本的降低。在將梯度在計(jì)算節(jié)點(diǎn)之間遷移之前對(duì)它們進(jìn)行壓縮已經(jīng)出現(xiàn)了不同的方法，具體有三類，分別是分裂（sparsification）、量化（quantization）和低秩（low-rank）方法。

分裂方法只遷移完整梯度元素的一些子集，并在參數(shù)向量中更新相應(yīng)的元素。這種近似方法能夠顯著降低通信成本，同時(shí)保持訓(xùn)練模型的性能，代表工作有 2017 年 Aji, A. F. 和 Heafield, K 的論文《 Sparse Communication for Distributed Gradient Descent 》和 2019 年 Alistarh, D. 等的論文《The Convergence of Sparsified Gradient Methods》。

另一種方法是基于遷移梯度的量化，該方法只遷移一定數(shù)量的 bit、從這些 bit 中重建整個(gè)梯度向量并更新參數(shù)向量的所有元素。這種方法對(duì)于一些神經(jīng)網(wǎng)絡(luò)架構(gòu)和實(shí)驗(yàn)設(shè)置得到了不錯(cuò)的結(jié)果，代表工作有 Alistarh, D. 等人 2017 年的論文《QSGD: Communication-Efficient SGD via Gradient Quantization and Encoding》。

最后一種降低通信成本的方法是低秩方法，其中在更新參數(shù)向量之前構(gòu)建、遷移和使用梯度的低秩近似來恢復(fù)完整格式的梯度。低秩近似可以通過塊能量（block power）方法或者最小化策略來構(gòu)建，各自的代表工作分別是 Vogels et al., 2019 和

Cho et al., 2019。

大批量訓(xùn)練。另一種加速優(yōu)化器收斂的方法是針對(duì)每個(gè)批使用大量的樣本。這種訓(xùn)練設(shè)置可以減少每個(gè) epoch 中的迭代次數(shù)，并提升 GPU 的利用率。在 Goyal, P 等人 2017 年的論文《Accurate, Large Minibatch SGD》中，研究者提出使用線性縮放規(guī)則來更新學(xué)習(xí)率和批大小。這一設(shè)置可以穩(wěn)定優(yōu)化過程，并將模型的最終性能收斂至相同。

封面來源：https://www.youtube.com/watch?v=RSRkp8VAavQ

關(guān)鍵詞： 大規(guī)模神經(jīng)網(wǎng)絡(luò)最新文獻(xiàn)綜述訓(xùn)練高效DNN節(jié)省內(nèi)存使用