天天看點

大規模神經網絡最新文獻綜述:訓練高效DNN、節省記憶體使用、優化器設計

選自arXiv

作者:Julia Gusak等

機器之心編譯

編輯:杜偉、澤南

在本綜述論文中,研究者解釋了不同技術的工作原理、評估和比較,還分析了一些實作這些技術的架構。

現代深度學習和人工智能技術的發展涉及使用深度神經網絡(DNN)來解決圖像、視訊、音頻、自然語言處理、圖像形式的内容生成等各種問題,或生成給定格式主題的文本等任務。

俄羅斯斯科爾科沃科學技術研究所、法國裡爾大學、波爾多大學、Inria 等科研機構聯合發表了一篇論文《Survey on Large Scale Neural Network Training》,它試圖解決的問題是:若給定模型和計算平台的情形下,如何訓練才是最有效率的。為了使訓練高效,其必須可行,最大程度地利用資源的計算能力,在并行情況下,它不能讓資訊傳輸成為瓶頸。訓練的效率從根本上取決于計算核心在計算資源(CPU、TPU、GPU)上的有效實作以及 GPU 之間和不同記憶體之間通信的有效實作。

大規模神經網絡最新文獻綜述:訓練高效DNN、節省記憶體使用、優化器設計

論文連結:https://arxiv.org/abs/2202.10435

在這兩種情況下,人們為優化計算核心的算術強度,及有效實作硬體網絡上的通信做了很多工作。對于使用者來說,已存在強大的分析工具來識别硬體瓶頸,并可用于判定本調查中描述哪些政策可用于解決算術強度、記憶體和控制交換資料量的問題。

該綜述研究涵蓋了應對這些限制的通用技術。如果由于模型、優化器狀态和激活不适合記憶體而無法先驗執行計算,則可以使用記憶體交換計算(重新實作)或資料轉移(激活和權重解除安裝)。我們還可以通過近似優化器狀态和梯度(壓縮、修剪、量化)來壓縮記憶體使用。

并行方法(資料并行、模型并行、流水線模型并行)也可以将記憶體需求分布到多個算力資源上。如果計算的算力強度不足以充分利用 GPU 和 TPU,一般是因為 mini-batch 太小,那麼上述技術也可以增加 mini-batch 的大小。最後,如果使用資料并行引起的通信開銷昂貴到拖累計算速度,則可以使用其他形式的并行(模型并行、流水線模型并行),梯度壓縮也可以限制資料交換的數量。

在本次調查中,研究者解釋了這些不同技術是如何工作的,其中描述了評估和比較所提出方法的文獻,還分析了一些實施這些技術的架構。

下表 1為文章讨論的不同技術及其對通信、記憶體和計算效率的影響。

大規模神經網絡最新文獻綜述:訓練高效DNN、節省記憶體使用、優化器設計

研究者根據目的區分了以下方法:首先讨論減少 GPU 記憶體使用,随後考慮對不适合 GPU 的模型使用并行訓練,最後讨論為訓練存儲在多個裝置上的模型而開發的優化器的設計。

單 GPU 情況下減少記憶體使用

在前向傳播期間,神經網絡存儲執行反向傳播所需的激活。在某些情況下,這些激活會消耗大量記憶體,讓模型無法訓練。減少記憶體使用的主要方法有兩種:重新實作(也稱為 checkpointing)和解除安裝。

激活的重新實作

重新實作的政策僅在前向傳播期間存儲一小部分激活,并在反向傳播期間重新計算其餘部分。重新實作方法可以通過它們處理的計算圖來區分。第一組來自自動微分(AD),它們為同構順序網絡(多層按順序執行并具有相同計算和記憶體成本的 DNN)找到最佳排程。第二組專注于過渡模型,例如異構序列網絡(可以是由任意複雜子產品組成的任何序列神經網絡,如 CNN、ResNet、一些 transformer),它将解決方案從 AD 調整為異構設定。

一些方法可以對一般計算圖執行重新實作,盡管确切的計算成本可能指數級上升,如下表 2 所示。

大規模神經網絡最新文獻綜述:訓練高效DNN、節省記憶體使用、優化器設計

激活解除安裝

解除安裝(又被稱為記憶體交換)是一種通過在前向傳遞期間将激活轉移到 CPU 記憶體并将它們預取回 GPU 記憶體,以進行相應的向後計算來節省 GPU 記憶體的技術。

由于 CPU 和 GPU 之間 PCI 總線的帶寬有限,必須優化選擇傳輸激活,以及何時傳輸的選擇。

在 vDNN [Rhu et al., 2016] 研究中,作者通過僅解除安裝卷積層的輸入來遵循對 CNN 有效的啟發式方法,然而它不能很好地推廣到一般 DNN 上。另有研究 [Le et al., 2018] 考慮了激活生命周期來選擇解除安裝的内容,并使用圖搜尋方法來識别插入解除安裝 / 預取操作的時刻。AutoSwap [Zhang et al., 2019] 通過為每個變量配置設定優先級分數來決定解除安裝哪些激活。

權重解除安裝

前面提到的很多方法也适用于解除安裝權重,這是因為解除安裝權重依賴于适用于任何張量的通用技術,比如 TFLMS、AutoSwap 或者 SwapAdvisor。

不适合單個 GPU 的模型的并行性

在模型并行化中,隻需要傳達激活資訊,并且傳輸隻發生在配置設定給不同處理器的連續層之間。本章節提到的工作如下表 4 所示。

大規模神經網絡最新文獻綜述:訓練高效DNN、節省記憶體使用、優化器設計

如果多個小批量被 pipeline 化 ,則可以加快模型并行化中的執行速度,進而同時激活了多個訓練疊代,具體可見 [Huang et al., 2019]。一旦在所有這些小批量上計算了前向和後向階段,權重就會更新。這種方法實作起來相當簡單,但也導緻計算資源大部分處于空置狀态。[Narayanan et al., 2019] 中提出的 PipeDream 方法僅強制前向和後向任務針對給定的小批量使用相同的模型權重,改進了這一訓練過程。

減少執行更新的頻率也已被證明有助于限制權重過期(Narayanan et al., 2021a)。[Yang et al., 2021] 提出的 PipeMare 根據 pipeline 階段向後調整學習率和模型權重。

對 pipeline 方法中激活導緻的存儲成本進行模組化是一項艱巨的任務(Beaumont et al., 2021b)。例如,[Fan et al., 2021] 中的 DAPPLE 、 [Li and Hoefler, 2021] 中的 Chimera 使用 1F1B(One-Forward-One-Backward)排程來減少與激活相關的記憶體消耗。1F1B 是一種同步權重更新技術,盡可能早地安排每個微批次的反向傳遞,以釋放激活占用的記憶體。

有些論文專門處理具有挑戰性的拓撲。比如,為了解決高通信成本和異構網絡能力的問題,[Zhan and Zhang, 2019] 中的 Pipe-torch 提出了一種更新的動态規劃政策,該政策假設計算和通信之間沒有重疊。[Park et al., 2020] 中的 Pipe 解決了異構 GPU 的其他問題,采用的方法是将這些異構 GPU 分成虛拟 worker,并在每個虛拟 worker 中運作 pipeline 并行化,同時依賴 worker 之間的資料并行化。

用于跨裝置模型訓練的優化器

零備援優化器

2020 年, Rajbhandari, S. 等人在論文《 ZeRO: Memory Optimizations toward Training Trillion Parameter Models》中提出了零備援優化器(Zero Redundancy Optimizer, ZeRO),将它作為一種減少記憶體使用的資料并行化實作。根據在裝置上劃分的張量,該算法具有三個階段,即階段 1 - 優化器狀态、階段 2 - 優化器狀态和梯度和階段 3 - 優化器狀态、梯度和模型超參數。

2021 年, Ren, J. 等人在論文《 ZeRO-Offload: Democratizing Billion-Scale Model Training》中将 ZeRO 與 Zero-Offload 内部參數更新的 CPU 端計算統一起來,其中梯度被遷移至存儲參數副本的 CPU,更新的權重遷移回 GPU。

低精度優化器

為了進一步減少記憶體使用,低精度優化器(low-precision optimizer)有了用武之地。這些方法使用低精度格式拉力表示優化器狀态以及狀态的輔助向量。并且,誤差補償技術可以被用來維持跟蹤統計的近似準确率。

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

收斂加速

另一種加速大規模深度學習模型的方法是減少節點之間的通信時間以及在适當局部最小值收斂所需的 epoch 數量。

關于通信成本的降低。在将梯度在計算節點之間遷移之前對它們進行壓縮已經出現了不同的方法,具體有三類,分别是分裂(sparsification)、量化(quantization)和低秩(low-rank)方法。

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

另一種方法是基于遷移梯度的量化,該方法隻遷移一定數量的 bit、從這些 bit 中重建整個梯度向量并更新參數向量的所有元素。這種方法對于一些神經網絡架構和實驗設定得到了不錯的結果,代表工作有 Alistarh, D. 等人 2017 年的論文《QSGD: Communication-Efficient SGD via Gradient Quantization and Encoding》。

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

Cho et al., 2019。

大批量訓練。另一種加速優化器收斂的方法是針對每個批使用大量的樣本。這種訓練設定可以減少每個 epoch 中的疊代次數,并提升 GPU 的使用率。在 Goyal, P 等人 2017 年的論文《Accurate, Large Minibatch SGD》中,研究者提出使用線性縮放規則來更新學習率和批大小。這一設定可以穩定優化過程,并将模型的最終性能收斂至相同。

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

繼續閱讀