本文翻譯整理自 A Survey on Efficient Training of Transformers

來自：無數(shù)據(jù)不只能

在本文中，將探討了transformer高效訓練方法，從存儲效率、硬件算法協(xié)同設計和計算效率三個角度進行了闡述。

在計算效率方面，介紹加速優(yōu)化器、初始化、稀疏訓練、過參數(shù)化、大批次訓練和增量訓練等優(yōu)化策略，以及token掩碼和重要性抽樣等數(shù)據(jù)選擇方法。

在內(nèi)存效率方面，介紹并行化、量化訓練、再物化、卸載和參數(shù)效率微調(diào)等策略。

在硬件算法協(xié)同設計方面，介紹稀疏矩陣乘法、低精度硬件感知和高效注意力等方法。

計算效率

優(yōu)化

加速優(yōu)化器

AdamW是Adam的一個變體，其解耦了正則化2和權重衰減，是transformer中最常用的優(yōu)化器之一。

Lion只使用一階梯度跟蹤動量，更新只考慮符號方向，對每個參數(shù)具有相同的量級，與AdamW等自適應優(yōu)化器有很大不同。在實踐中，Lion一般收斂更快，在各種基準上訓練transformer時，Lion的內(nèi)存效率和準確性都比AdamW更高。

初始化

好的初始化對于穩(wěn)定訓練、加速收斂、提高學習率和泛化能力很重要。

Fixup重新縮放標準初始化以確保梯度范數(shù)適當，避免梯度爆炸或消失，從而可在不添加歸一化層的情況下使用非常深的網(wǎng)絡進行訓練。

ReZero和SkipInit只是將每個層初始化為執(zhí)行恒等操作，其基礎是歸一化殘差塊計算出的函數(shù)接近于恒等函數(shù)。它們在每個殘差塊的輸出上添加一個可學習的縮放乘法器。T-Fixup分析表明，Adam優(yōu)化器中二階動量的方差無界，導致部分優(yōu)化難度。因此，T-Fixup采用一種類似于Fixup對殘差塊的初始化的縮放方案。以上所有方法都從所有塊中刪除批/層歸一化，并在沒有學習率warm up的情況下進行訓練。

在訓練深度視覺transformer(ViT)時，提出了通道級可學習的縮放因子，并根據(jù)經(jīng)驗觀察到重新引入預熱和層歸一化技術可以使訓練更加穩(wěn)定。

稀疏訓練

稀疏訓練的關鍵思想是在不犧牲精度的情況下，直接訓練稀疏子網(wǎng)絡而不是全網(wǎng)絡?？煽啃允紫扔刹势奔僬f(LTH)證明，一個密集的、隨機初始化的網(wǎng)絡包含子網(wǎng)絡，這些子網(wǎng)絡可以單獨訓練，以匹配原始網(wǎng)絡的精度。然而，LTH需要以交替訓練-剪枝-再訓練的方式識別中獎彩票，這使得大型模型和數(shù)據(jù)集的訓練成本極其昂貴，限制了實際效益。鑒于此，具有更高訓練效率的后續(xù)工作可以大致分為三類：

(i) 通過測量連接對損失的重要性，在初始化時一次性找到稀疏網(wǎng)絡，消除了對復雜迭代優(yōu)化調(diào)節(jié)的需要；(ii) 通過低成本方案在非常早期的訓練階段識別transformer中的中獎彩票，然后僅僅訓練這些早期彩票，直到收斂；(iii) 在整個訓練過程中使用交替修剪和增長計劃來動態(tài)更新模型稀疏模式，適用于通用架構。

過參數(shù)化

實用的DNN是過于參數(shù)化的，其可學習的參數(shù)數(shù)量遠遠大于訓練樣本的數(shù)量。過參數(shù)化可以在一定程度上提高收斂性和泛化性，雖然不充分，但有理論支持。早期的研究證明在線性神經(jīng)網(wǎng)絡中通過過參數(shù)化增加深度可以加速SGD的收斂。進一步探索了兩層非線性神經(jīng)網(wǎng)絡，證明了SGD可以在多項式時間內(nèi)收斂到DNNs的訓練目標上的全局最小值，假設訓練樣本不重復，并且參數(shù)數(shù)量是多項式級別的。在泛化方面，理論證明了充分過參數(shù)化的神經(jīng)網(wǎng)絡可以泛化到種群風險，一個有趣的特性是，在SGD訓練軌跡上任意點的近鄰域內(nèi)存在一個精確的網(wǎng)絡，具有高概率的隨機初始化。需要注意的是，與LTH有很深的關聯(lián)，因為它部分解釋了為什么LTH在稀疏訓練中表現(xiàn)良好，因為過度參數(shù)化會產(chǎn)生許多低風險的好子網(wǎng)絡。在transformer中應用這一理論，可以獲得更快的收斂速度和更好的性能。

大批次訓練

另一種加速訓練的流行方法是使用大批量，減少每個epoch的迭代次數(shù)，提高計算資源利用率。從統(tǒng)計學的角度來看，大批量訓練降低了隨機梯度估計的方差，因此需要調(diào)整可靠的步長以獲得更好的收斂。在卷積神經(jīng)網(wǎng)絡時代，利用學習率的線性縮放，在ImageNet上以8192的批次大小在1小時內(nèi)訓練ResNet-50。隨后提出了更先進的步長估計方法。廣泛使用的方法是SGD的LARS和Adam的LAMB，它們建議分別對ResNet和Transformers使用分層自適應學習率。

它配備了一個歸一化項，對梯度爆炸和平臺提供魯棒性，縮放項確保更新的參數(shù)與參數(shù)的范數(shù)相同階次，促進更快的收斂。最近，根據(jù)經(jīng)驗，針對大批量訓練定制的更強大的優(yōu)化方法表現(xiàn)良好。例如，表明，將一定數(shù)量的最新檢查點的權重進行平均，可以促進更快的訓練。DeepMind在中訓練了400多個具有不同規(guī)模的模型大小和訓練token的Transformer語言模型，達到了一個實用的假設，即模型大小和訓練token的數(shù)量應該被均勻縮放以進行計算最優(yōu)的LLM訓練。

增量訓練

增量學習的高級概念是將原來具有挑戰(zhàn)性的優(yōu)化問題放松為一系列易于優(yōu)化的子問題，其中一個子問題的解可以作為后續(xù)子問題的良好初始化，以規(guī)避訓練難度，類似于退火。一些工作提出通過逐步堆疊層來加速BERT預訓練，從一個較小的模型適當?shù)爻跏蓟粋€較大的模型。以相反的方向通過層丟棄來訓練具有隨機深度的transformer，其中它沿著時間維度和深度維度逐步增加丟棄率。為ViT定制，AutoProg建議在使用神經(jīng)架構搜索的漸進式學習過程中自動決定模型是否增長、在哪里增長以及應該增長多少。一個關鍵的觀察結果是，逐步增加輸入圖像的分辨率（減少patch大小）可以顯著加快ViT訓練，這與眾所周知的訓練動態(tài)一致，即在早期集中于低頻結構，在后期集中于高頻語義。

數(shù)據(jù)選擇

除了模型效率，數(shù)據(jù)效率也是高效訓練的關鍵因素。

token掩碼

Toke n m asking 是自監(jiān)督預訓練任務中的一種主要方法，例如掩碼語言建模和掩碼圖像建模。標記掩碼的精神是隨機掩碼一些輸入標記，并訓練模型用可見標記中的上下文信息來預測缺失的內(nèi)容，例如詞匯表ID或像素。由于壓縮序列長度以二次方式降低了計算和存儲復雜度，跳過處理掩碼token為MLM和MIM帶來了可觀的訓練效率增益。對于MLM，提出聯(lián)合預訓練語言生成任務的編碼器和解碼器，同時刪除解碼器中的掩碼標記，以節(jié)省內(nèi)存和計算成本。對于MIM，代表性工作表明，在視覺中，在編碼器之前移除掩碼圖像塊顯示出更強的性能，并且比保留掩碼token的整體預訓練時間和內(nèi)存消耗低3倍或更多。在中也發(fā)現(xiàn)了類似的現(xiàn)象，對于語言-圖像預訓練，隨機掩碼并去除掩碼圖像塊顯示了比原始剪輯快3.7倍的整體預訓練時間。

重要性抽樣

在數(shù)據(jù)上采樣的重要性，也稱為數(shù)據(jù)剪枝，理論上可以通過優(yōu)先考慮信息量大的訓練樣本來加速監(jiān)督學習的隨機梯度算法，主要受益于方差縮減。對于DNNs來說，估計每個樣本重要性的一種主要方法是使用梯度范數(shù)，并且使用不同的近似使計算這些規(guī)范變得容易進一步加快類似于early-bird LTH的采樣過程，但在數(shù)據(jù)域，在幾個權重初始化上的簡單平均梯度范數(shù)或誤差2范數(shù)可以用于在訓練的非常早期階段識別重要的示例。最近，展示了一個令人興奮的分析理論，即測試誤差與數(shù)據(jù)集大小的縮放可以突破冪次縮放定律，如果配備了優(yōu)越的數(shù)據(jù)修剪度量，則可以減少到至少指數(shù)縮放，并且它采用了使用k-means聚類的自監(jiān)督度量。它展示了一個有希望的方向，即基于數(shù)據(jù)重要性采樣的更有效的神經(jīng)縮放定律。

內(nèi)存效率

除了計算負擔之外，大型Transformer模型不斷增長的模型大小，例如從BERT 345M參數(shù)模型到1.75萬億參數(shù)的GPT-3，是訓練的一個關鍵瓶頸，因為它們不適合單個設備的內(nèi)存。我們首先分析了現(xiàn)有模型訓練框架的內(nèi)存消耗，這些內(nèi)存消耗由模型狀態(tài)，包括優(yōu)化器狀態(tài)、梯度和參數(shù)；以及激活。我們在表1中總結了內(nèi)存有效的訓練方法。在下文中，我們將討論優(yōu)化內(nèi)存使用的主要解決方案。

并行化

跨設備并行訓練大型DNN是滿足內(nèi)存需求的常見做法?；旧嫌袃煞N范式：數(shù)據(jù)并行（DP），它在不同的設備上分配一個小批量的數(shù)據(jù)，模型并行(MP)，它在多個worker上分配一個模型的子圖。對于DP來說，隨著可用worker的增加，批處理大小接近線性縮放。第2節(jié)中討論的大批量培訓就是針對這種情況開發(fā)的。然而，很明顯，DP具有較高的通信/計算效率，但內(nèi)存效率較差。當模型變得很大時，單個設備無法存儲模型副本，并且針對梯度的同步通信會阻礙DP的可擴展性。因此，DP本身只適合訓練小到中等規(guī)模的模型。為了提高DP的可擴展性，Transformer的一種解決方案是參數(shù)共享 ，即Albert，但它限制了表征能力。最近，ZeRO將均勻劃分策略與DP結合在一起，其中每個數(shù)據(jù)并行過程僅處理模型狀態(tài)的一個劃分，在混合精度制度下工作。為了處理非常大的DNN，人們總是需要利用模型并行性，以“垂直”的方式在多個加速器上分配不同的層。雖然MP具有良好的存儲效率，但由于跨設備的大量數(shù)據(jù)傳輸和較差的PE利用率，其通信和計算效率較低。幸運的是，在an中有兩種策略可以進一步提高MP的效率

正交的“水平”維度，包括張量并行(TP)和流水線并行(PP)。TP跨worker將張量操作劃分在一個層中，以實現(xiàn)更快的計算和更多的內(nèi)存節(jié)省。根據(jù)基于transformer的模型定制，Megatron-LM跨GPU對MSA和FFN進行切片，并且在前向和后向傳遞中只需要一些額外的All-Reduce操作，使他們能夠使用512個GPU訓練模型多達83億個參數(shù)。在PP方面，它最初是在GPipe中提出的，它將輸入的mini-batch分割為多個較小的micro-batch，使不同的加速器（在加速器上劃分順序?qū)樱┩瑫r在不同的micro-batch上工作，然后對整個mini-batch應用單一的同步梯度更新。然而，它仍然受到管道氣泡（加速器空閑時間）的影響，這會降低效率。特別是，PyTorch實現(xiàn)了torchgpipe，它使用檢查點執(zhí)行微批次PP，允許擴展到大量的微批次，以最小化氣泡開銷。請注意，DP和MP是正交的，因此可以同時使用兩者來訓練具有更高計算和內(nèi)存容量的更大模型。例如，Megatron-LM和DeepSpeed組成了張量、管道和數(shù)據(jù)并行，以將訓練擴展到數(shù)千個GPU。

量化訓練

訓練神經(jīng)網(wǎng)絡的標準程序采用全精度（即FP32）。相比之下，量化訓練通過將激活值/權重/梯度壓縮為低比特值（例如FP16或INT8），以降低的精度從頭開始訓練神經(jīng)網(wǎng)絡。在之前的工作中已經(jīng)表明，降低精度的訓練可以加速神經(jīng)網(wǎng)絡的訓練，并具有良好的性能。對于Transformer來說，采用最廣泛的方法是自動混合精度（AMP）訓練。具體來說，AMP以全精度存儲權重的主副本用于更新，而激活值、梯度和權重則存儲在FP16中用于算術。與全精度訓練相比，AMP能夠?qū)崿F(xiàn)更快的訓練/推理速度，并減少網(wǎng)絡訓練期間的內(nèi)存消耗。例如，基于64的批量大小和224×224的圖像分辨率，在AMP下的RTX3090上訓練DeiT-B比全精度訓練快2倍（305張圖像vs124張圖像/s），同時消耗22%的峰值GPU內(nèi)存（7.9GBvs10.2GB）。雖然人們普遍認為，至少需要16位才能訓練網(wǎng)絡而不影響模型精度，但最近在NVIDIA H100上對FP8訓練的支持在Transformer訓練上顯示出了可喜的結果，在FP8下訓練DeiT-S和GPT可以與16位訓練相媲美。除了同時量化激活/權重/梯度的精度降低訓練外，激活壓縮訓練（ACT）在精確計算前向傳遞的同時存儲低精度的激活近似副本，這有助于減少訓練期間的整體內(nèi)存消耗。然后將保存的激活解量化為向后傳遞中的原始精度，以計算梯度。最近的工作進一步建議自定義ACT以支持內(nèi)存高效的Transformer訓練。

再物化和卸載

再物化 (R e materialization)，也被稱為檢查點，是一種廣泛使用的時空權衡技術。在正向傳遞時只存儲一部分激活/權重，在反向傳遞時重新計算其余部分。提供了一個在PyTorch3中實現(xiàn)的簡單的周期性調(diào)度，但它僅適用于同構的順序網(wǎng)絡。更先進的方法，如實現(xiàn)了異構網(wǎng)絡4的最優(yōu)檢查點。在卸載方面，它是一種利用CPU內(nèi)存等外部內(nèi)存作為GPU內(nèi)存的擴展，通過GPU和CPU之間的通信，在訓練過程中增加內(nèi)存容量的技術。模型狀態(tài)和激活，都可以卸載到CPU，但最優(yōu)選擇需要最小化通信成本（即數(shù)據(jù)移動）到GPU，減少CPU計算和最大化GPU內(nèi)存節(jié)省。一個代表性的工作是ZERO-Offload，它提供了使用Adam優(yōu)化器定制的混合精度訓練的最優(yōu)卸載策略。它在CPU內(nèi)存上卸載所有fp32模型狀態(tài)和fp16梯度，并在CPU上計算fp32參數(shù)更新。fp16參數(shù)保存在GPU上，前向和后向計算在GPU上。為了兩全兼顧，建議聯(lián)合優(yōu)化激活卸載和再物化。

參數(shù)效率微調(diào)

以 HuggingFace 為代表的公共模型動物園包含豐富的預訓練模型，可以隨時下載和執(zhí)行，正在為降低訓練成本做出顯著貢獻。對這些現(xiàn)成的模型進行有效調(diào)優(yōu)，正成為一種大幅削減訓練成本的流行方式。作為香草全微調(diào)的強大替代方案，

參數(shù)高效調(diào)優(yōu)（PET）在凍結預訓練模型的同時只更新少量額外參數(shù)，以顯著減少存儲負擔。它可以隨動態(tài)部署場景擴展，而不需要為每種情況存儲單獨的模型實例。一般的PET方法可以分為基于添加的方法和基于重新參數(shù)化的方法。前者將額外的可訓練參數(shù)附加到預訓練模型上，并且只調(diào)整這些參數(shù)。例如，在輸入空間中添加可訓練參數(shù)，在MSA和FFN之后在每個Transformer塊中添加兩次適配器模塊。然而，額外的參數(shù)會在推理過程中引入額外的計算和內(nèi)存開銷。為了應對這一挑戰(zhàn)，后者建議調(diào)整模型中固有的參數(shù)或可以重新參數(shù)化到模中的新參數(shù)，從而不會對推理效率產(chǎn)生犧牲。受觀察到大型語言預訓練模型具有較低的內(nèi)在維度的啟發(fā)，代表性工作LoRA將自注意力權重的更新近似為兩個低秩矩陣，可以在推理過程中將其合并到預訓練權重中。值得注意的是，在實現(xiàn)LLM民主化方面最受認可的工作之一是Stanford Alpaca，它使用從Chat GPT生成的52K指令遵循數(shù)據(jù)從開源的LLaMA模型中進行微調(diào)。為了廉價有效地對其進行微調(diào)，其變體Alpaca-LoRA5進一步采用低秩LoRA來實現(xiàn)在客戶硬件上對羊駝進行指令微調(diào)，表明可以在單個RTX 4090上在數(shù)小時內(nèi)完成訓練。

開源框架

有幾個被廣泛采用的原型用于大規(guī)模訓練大型Transformer模型，其中微軟DeepSpeed、HPC-AI Tech ColossalAI和Nvidia Megatron是先驅(qū)。具體來說，DeepSpeed主要基于和ZERO系列作品實現(xiàn)，Colossal-AI基于構建，Megatron-LM實現(xiàn)。所有這些都支持混合精度的數(shù)據(jù)和模型并行，以及其他通用實踐，如卸載和再物化。更多用于高效分布式訓練的庫包括但不限于HuggingFace Transformers，MosaicML Composer，百度PaddlePaddle，Bytedance Lightspeed，EleutherAI GPT-NeoX等。

硬件算法協(xié)同設計

除了計算和內(nèi)存負擔，設計高效的硬件加速器可以使DNNs更快地訓練和推理。具體來說，與中央處理單元（CPU）相比，圖形處理單元（GPU）由于高度并行性，在執(zhí)行矩陣乘法時更加強大。對于專注于特定計算任務的應用，專用集成電路（ASIC）具有低功耗、高訓練/推理速度的優(yōu)勢。例如，谷歌設計的張量處理單元（TPU）比當代CPU和GPU提供了30~80倍的性能/瓦特。然而，ASIC不容易重新編程或適應新任務。相比之下，現(xiàn)場可編程門陣列（FPGA）被設計成可以根據(jù)需要重新編程以執(zhí)行不同的功能，在最終確定設計之前也可以作為ASIC的原型。為了進一步優(yōu)化DNNs特別是Transformer的訓練效率，硬件-算法協(xié)同設計在設計算法時考慮了硬件的約束和能力，這將在以下小節(jié)中介紹。

稀疏矩陣乘法

為了降低 Transformer 的計算開銷，涉及將稀疏矩陣與稠密矩陣相乘的稀疏通用矩陣乘法（SpGEMM），利用注意力矩陣的稀疏性來減少計算次數(shù)。目前有幾個流行的稀疏矩陣計算庫，如Intel Math Kernel library CPU和cuSPARSE15、CUSP16和24 GPU上的結構化稀疏性。然而，由于不規(guī)則的稀疏性，SpGEMM往往對CPU、GPU等通用處理器的硬件不友好。為了解決這個問題，需要專門的硬件加速器，如FPGA和ASIC，來處理糟糕的數(shù)據(jù)局部性問題。例如，OuterSPACE 將矩陣乘法轉(zhuǎn)換為外積過程，并通過將乘法與累加解耦來消除冗余內(nèi)存訪問。為了在不引入顯著開銷的情況下充分利用這一減少，OuterSPACE構建了一個具有可重構存儲層次的自定義加速器，相對于運行英特爾數(shù)學內(nèi)核庫的CPU實現(xiàn)了7.9倍的平均加速，相對于運行CUSP的GPU實現(xiàn)了14.0倍的平均加速。此外，為了緩解高稀疏性造成的數(shù)據(jù)移動瓶頸，ViTCoD使用可學習的自編碼器將稀疏注意力壓縮為更緊湊的表示，并設計編碼器和解碼器引擎以提高硬件利用率。

低精度硬件感知

降低計算精度會減少內(nèi)存和計算量，這可以用硬件友好的定點或整數(shù)表示來實現(xiàn)，而不是浮點表示。因此，我們可以使用精度較低的乘法器、加法器和內(nèi)存塊，從而在功耗和加速比方面得到顯著改善。此外，低精度的算法可以與其他技術相結合，如剪枝和低秩近似，以實現(xiàn)進一步的加速。例如，Sanger使用4位查詢和鍵來計算稀疏注意力矩陣的量化預測。然后，稀疏注意力掩碼被重新排列成結構化的塊，并由可重構硬件處理。以下工作DOTA使用低秩變換和低精度計算識別注意力中的不重要連接。通過結合token級并行和亂序執(zhí)行，DOTA實現(xiàn)了GPU上152.6倍的加速。

高效注意力

除了稀疏的矩陣乘法和低精度的計算，幾項開創(chuàng)性的工作都專注于在硬件上實現(xiàn)高效和輕量級的注意力。具體來說，一個3只選擇那些可能與給定查詢具有高相似性的鍵，以減少注意力的計算量。ELSA根據(jù)哈希相似度過濾掉與特定查詢無關的鍵，以節(jié)省計算量。借助高效的硬件加速器，ELSA實現(xiàn)了58.1倍的加速，與配備16GB內(nèi)存的Nvidia V100 GPU相比，在能效方面有三個數(shù)量級的提升。值得注意的是，F(xiàn)lash Attention提出利用tiling來減少GPU高帶寬內(nèi)存(HBM)和片上SRAM之間的I/O通信，這正在成為一個默認的快速和內(nèi)存高效的注意力模塊來加速。

審核編輯：湯梓紅