2021年 2 月份發布了 DeepSpeed。這是一個開源深度學習訓練優化庫，包含的一個新的顯存優化技術—— ZeRO（零冗余優化器），通過擴大規模，提升速度，控制成本，提升可用性，極大地推進了大模型訓練能力。DeepSpeed 已經幫助研究人員，開發了圖靈自然語言生成模型( Turing-NLG)，在發表時，為世界上最大的語言模型（擁有 170 億參數），有着最佳的精度。在2021年 5 月份發布了 ZeRO-2——支持有着 2000 億參數的模型訓練，與最新技術相比，訓練速度可達 10 倍——系列計算、IO 和收斂優化功能，從而助力最快速的 BERT 訓練。自那時起，持續高速進行創新，不斷突破深度學習模型訓練的速度和規模的邊界。
這些進展不僅會推動深度學習訓練走向極致，同時也讓這份技術的使用范圍更加廣泛——上至數據科學家們在超算上訓練，下至在低端集群甚至僅僅一張 GPU 上訓練。具體來說，DeepSpeed 加入了 4 項系統性新技術來進一步拓展 AI at Scale 倡議。推動了微軟的AI產品與平台的創新。這些技術提供了極為高效的計算、顯存和通信的利用效率，並助力訓練有着十億至萬億量級參數的模型。支持超長輸入序列，無論在單卡GPU、千卡GPU的高端集群上，還是在慢速以太網的低端集群上，均可以使用。
- 用 3D 並行化實現萬億參數模型訓練： DeepSpeed 實現了三種並行方法的靈活組合：ZeRO 支持的數據並行，流水線並行和張量切片模型並行。3D 並行性適應了不同工作負載的需求，支持具有萬億參數的超大型模型，實現了近乎完美的顯存擴展性和吞吐量擴展效率。提高的通信效率，使用戶可以在網絡帶寬有限的常規群集上，以 2-7 倍的速度訓練有數十億參數的模型。
- ZeRO-Offload 使 GPU 單卡，能夠訓練 10 倍大的模型： 為了同時利用 CPU 和 GPU 內存來訓練大型模型，擴展了 ZeRO-2。用戶在使用帶有單張英偉達 V100 GPU 的機器時，可以在不耗盡顯存的情況下，運行多達 130 億個參數的模型，模型規模擴展至現有方法的10倍，保持有競爭力的吞吐量。此功能使數十億參數的模型訓練，更加大眾化，，並為許多深度學習從業人員，打開了一扇探索更大更好的模型的窗戶。
- 通過 DeepSpeed Sparse Attention，用6倍速度執行10倍長的序列： DeepSpeed提供了稀疏 attention kernel ——一種工具性技術，可支持長序列的模型輸入，包括文本輸入，圖像輸入和語音輸入。與經典的稠密 Transformer 相比，支持的輸入序列長一個數量級，在保持相當的精度下，獲得最高 6 倍的執行速度提升。還比最新的稀疏實現快 1.5–3 倍。稀疏 kernel 靈活支持稀疏格式，使用戶能夠通過自定義稀疏結構進行創新。
- 1 比特 Adam 減少 5 倍通信量： Adam 是一個在大規模深度學習模型訓練場景下的，有效的（也許是最廣為應用的）優化器。與通信效率優化算法往往不兼容。在跨設備進行分布式擴展時，通信開銷可能成為瓶頸。推出了一種 1 比特 Adam 新算法，高效實現。該算法最多可減少 5 倍通信量，實現了與Adam相似的收斂率。在通信受限的場景下，觀察到分布式訓練速度，提升了 3.5 倍，這使得該算法，可以擴展到不同類型的 GPU 群集和網絡環境。
a screenshot of a cell phone
將深入探究這 4 項技術。已經將這些激動人心的優化技術公布在了開源項目 DeepSpeed中。
3D 並行：擴展至萬億參數模型
隨着現代 GPU 群集上計算量的快速增長，訓練具有驚人的功能的、強大的萬億參數模型，不再是遙不可及的，可能在不久的將來就能實現。DeepSpeed 結合了三項強大的技術，可以訓練數萬億規模的模型，擴展到數千個 GPU：數據並行訓練，模型並行訓練和流水線並行訓練。這三者的共生讓深度學習訓練的規模，遠遠超出了單獨使用每種策略。3D 並行解決了訓練萬億參數模型的兩個基本挑戰：顯存效率和計算效率。DeepSpeed 可以擴展至在顯存中，放下最巨大的模型，不會犧牲速度。
了解訓練巨大模型的顯存和計算效率的挑戰
顯存效率：訓練萬億參數模型所需的顯存，遠遠超出了單張 GPU 的顯存大小。在使用 Adam 優化器，進行混合精度訓練時，存儲模型狀態量（參數、梯度和優化器狀態量）需要約 16TB 的顯存。最先進的英偉達 A100 GPU 只有 40 GB 的顯存。僅僅為了存儲模型狀態，需要 400 張這樣的 GPU。
激活函數額外消耗的顯存，隨 batch 大小而增加。batch 設置為1的情況下，訓練萬億參數模型，產生超過 1 TB 的激活函數用的顯存（后文稱為激活顯存）。用 checkpoint 處理激活顯存，用計算換顯存，可以將顯存減少到大約20 GB，對於訓練而言，仍然過高了。
必須在多個 GPU 設備之間，有效地划分模型狀態量和激活顯存，才能讓這種大模型，在不耗盡顯存的情況下，開始訓練。
計算效率：經估算端到端訓練一個萬億參數的模型，大約需要 5000 Zflops（即 5 后面帶有 24 個零；這個估算結果基於 OpenAI 的研究 law of scaling）。訓練一個模型需要 4000 張 A100，以 50％ 的計算效率運行大約 100 天。
盡管大型超級計算 GPU 集群，可以擁有超過 4000 個 GPU，但是由於 batch 大小的限制，要在這種規模上實現高計算效率，仍然是一項挑戰。計算效率隨着計算時間，對通信時間的比例的增加而增加。該比例與 batch 大小成正比。訓練模型的 batch 大小，有一個上限——超過這個上限，收斂情況會明顯變差。
實際上最大的模型之一，GPT-3 的訓練 batch 大小約 1500。如果使用大約 4000 張 GPU, 即使可以自由設置 batch 大小為 4000，每張卡上的 batch 大小，也只有 1，這將影響擴展性。
理解數據並行、模型並行和流水線並行之間的權衡
數據並行是深度學習中的一種普遍使用的技術。每批輸入的訓練數據，都在數據並行的 worker 之間平分。反向傳播后，需要通信規約梯度，保證優化器在各個 worker 上進行相同的更新。數據並行性，具有幾個明顯的優勢，包括計算效率高和實現起來工作量小。數據並行的 batch 大小，隨 worker 數量提高，往往無法在不影響收斂性的情況下，一直增加 batch 大小。
- 顯存效率：數據並行會在所有 worker 之間，進行模型和優化器的復制，顯存效率不高。DeepSpeed 開發了 ZeRO ，一系列用於提高數據並行的顯存效率的優化器。依賴於 ZeRO 的 1 階段，在 worker 之間划分優化器狀態量，減少冗余。
- 計算效率：隨着提高並行度，每個 worker 執行的計算量是恆定的。數據並行，在小規模上實現近乎線性擴展。在 worker 之間規約梯度的通信開銷，跟模型大小成正相關，當模型很大或通信帶寬很低時，計算效率會受限。。梯度累積是一種均攤通信成本的一種常用策略。進一步增加batch大小，在本地使用 micro-batch，多次進行正向和反向傳播積累梯度后，進行梯度規約和優化器更新。
模型並行是包含范圍很廣的一類技術。在多個 worker 之間，划分模型的各個層。模型並行性的計算和通信因模型結構而異，在實現上有很大的工作量。DeepSpeed 借用了英偉達的 Megatron-LM，基於 Transformer 的語言模型，提供大規模模型並行功能。模型並行根據 worker 數量，成比例地減少顯存使用量，也是這三種並行度中顯存效率最高的。代價是計算效率最低。
- 顯存效率：模型並行會根據 worker 數量，成比例地減少顯存使用量。減少單個網絡層的激活顯存的唯一方法。DeepSpeed 通過在模型並行 worker 之間，划分激活顯存，進一步提高顯存效率。
- 計算效率：每次前向和反向傳播需要額外通信激活值，模型並行的計算效率很低。模型並行需要高通信帶寬，不能很好地擴展到通信帶寬受限的節點。每個模型並行worker，都會減少每個通信階段之間執行的計算量，影響計算效率。模型並行性通常與數據並行性結合使用，在內存和計算效率之間進行權衡。
流水線並行訓練引擎，包含在了這次發布的DeepSpeed中！流水線並行將模型的各層划分，可以並行處理。當完成一個 micro-batch 的正向傳遞時，激活內存將被通信至流水線的下一個階段。當下一階段完成反向傳播時，將通過管道反向通信梯度。必須同時計算多個 micro-batch，確保流水線的各個階段能並行計算。已經開發出了幾種用於權衡內存和計算效率收斂行為的方法，例如 PipeDream。DeepSpeed通過梯度累積來實現並行，保持與傳統數據並行和模型並行訓練，在相同的總 batch 大小下收斂。
- 顯存效率：流水線並行減少的顯存與流水線的階段數成正比，使模型的大小，可以隨 worker 的數量線性擴展。流水線並行不會減少每一層的激活函數的顯存占用量。每個 worker 必須存儲，同時運行的各個 micro-batch激活值。導致流水線第一階段的激活內存，與單個 mirco batch 的總激活內存，大致相同。一個萬億參數模型，需要為一個 micro batch，提供大約 19 GB 的顯存的激活內存，幾乎占到新推出的英偉達 A100 GPU 總顯存的一半。
- 計算效率：流水線並行具有最低的通信量，通信量只和在各階段邊界的各層的激活值大小成正比。但是不能無限擴展。像模型並行一樣，增加流水線大小，會減少每個流水線階段的計算量，降低計算與通信的比率。如果要實現好的計算效率，流水線並行，要求其每個階段的計算負載完美均衡。
此外，流水線並行性，在每個 batch 的開始和結束時，因為需要重新填充，或排空流水線，產生 bubble overhead。使用流水線階段數的 4 倍或 8 倍的梯度，累積步驟（以及 batch 大小）進行訓練，相較於只有一個流水線階段，分別達到了 81％ 和 90％ 的擴展性。
通過3D並行，同時實現高內存效率和高計算效率
數據，模型和流水線並行，在提高內存和計算效率方面，均起到特定的作用。圖 1 說明了 3D 策略。
顯存效率：先將模型的各層，划分到不同的流水線階段，把每個階段的層通過模型並行，進行划分。這種 2D 組合同時減少了模型、優化器和激活函數所消耗的內存。不能在不引入通信開銷的情況下，無限划分模型，通信開銷會限制計算效率。
計算效率：為了在不犧牲計算效率，將 worker 數量擴展至超出模型和流水線並行規模，使用了 ZeRO 支持的數據並行功能（ZeRO-DP）。ZeRO-DP 不僅可以通過划分優化器狀態量，進一步提高顯存利用效率，通過利用基於通信拓撲的映射關系，以最小的通信開銷，擴展到任意數量的 GPU。
基於通信拓撲的 3D 映射（圖2）：通過利用兩個關鍵的架構屬性，將 3D 並行中的每個維度，仔細映射到 worker 上，實現最大的計算效率。
- 優化節點內和節點間的通信帶寬：模型並行是這三種策略中通信開銷最大的，優先考慮將模型並行 worker 組放置在節點內，利用更大的節點內帶寬。基於英偉達 Megatron-LM，進行了張量切分式的模型並行。當模型並行組不占滿節點內的所有 worker 時，選擇將數據並行組放置在節點內。不然，就跨節點進行數據並行。流水線並行的通信量最低，可以跨節點調度流水線的各個階段，不受通信帶寬的限制。
- 通過並行通信增大帶寬：每個數據並行組，需要通信的梯度量，隨着流水線和模型並行的規模線性減小，總通信量少於單純使用數據並行。此外，每個數據並行組，在局部的一小部分 worker 內部，獨立進行通信，組間通信可以相互並行。通過減少通信量和增加局部性與並行性，數據並行通信的有效帶寬被增大了。
該圖顯示了一個有 32 個 worker 進行 3D 並行的例子。神經網絡的各層分為四個流水線階段。每個流水線階段中的層在四個模型並行 worker 之間進一步划分。最后，每個流水線階段有兩個數據並行實例，且 ZeRO 在這 2 個副本之間划分優化器狀態量。
圖 1：一個有 32 個 worker 進行 3D 並行的例子。神經網絡的各層分為四個流水線階段。每個流水線階段中的層在四個模型並行 worker 之間，進一步划分。每個流水線階段，有兩個數據並行實例，ZeRO 在這 2 個副本之間，划分優化器狀態量。
彩色塊顯示圖 1 中的 worker 到八個節點（每個節點有四個 GPU）的系統上的 GPU 的映射。同一顏色的 GPU 在同一節點上。
圖 2：圖 1 中的 worker 到八個節點（每個節點有四個 GPU）的系統上的 GPU 的映射。同一顏色的 GPU 在同一節點上。
了解關於 3D 並行訓練萬億參數模型的更多信息
使用 8 路模型並行，64 路流水線並行和 8 路數據並行，在 4096 個英偉達 A100 GPU 上，擴展訓練一個萬億參數模型。
通過結合模型並行和流水線並行，3D 並行可實現出色的內存效率和跨多個節點的高效計算效率。模型並行性，提高了節點內的激活內存和模型狀態量的存儲效率，流水線並行，相較於僅使用模型並行，可以在不犧牲計算效率的情況下，跨節點高效存儲模型狀態。在 micro-batch 大小為 1 的萬億參數例子中，在使用激活值 checkpoint及上述 3D 並行后，模型狀態量會消耗 30 GB 的顯存，划分后的激活值消耗 2.5 GB 的內存。總顯存占用為 32.5 GB，就能夠使用具有 40 GB 內存的英偉達 A100 GPU，容納和訓練這樣的模型了。
結合模型並行與流水線並行，可以使流水線並行在非常小的 batch 下，以最小的 bubble overhead，實現高計算效率。在 8 路模型並行下，每個模型使用 micro-batch 為 1 個微批處理，導致每個 GPU 的有效 micro-batch 大小為 1/8。使用 8 倍於管道並行度的梯度累加步驟，讓每張 GPU 上的總累計 batch 大小為 1，流水並行處理，可以實現 90％ 的計算效率。與數據並行性結合使用時，讓 4096 張 GPU 上的總有效 batch 大小為 4096，仍然可以達到 90％ 的流水線效率。
數據並行怎樣影響計算效率呢？難道數據並行不是需要每張 GPU，都有大 batch 才能保持高效嗎？
模型並行可以將每張GPU上的有效 batch 大小，減小到小於 1。使流水線並行，即使在小 batch 下，仍可以隱藏流水線 bubble overhead。通過跨節點使用流水線並行性，可以讓流水線每個階段的數據並行節點之間的獨立進行通信，與其它流水線並行進行。實際上，在高端 GPU 集群中常見的，完全連接的網絡拓撲中，對可用於數據並行訓練的有效通信帶寬，具有重要意義。流水線階段中的每個節點，都可以與其對應的數據並行節點，並行通信，有效的通信帶寬與流水線階段數成正比。通過設置64個並行流水線，有效帶寬將變為往返單個節點的帶寬的 64 倍。流水線並行帶給數據並行，如此大的有效帶寬，使數據並行在計算與通信比率非常低的小 batch 情況下，實現高效擴展。
在線性擴展性下，訓練萬億參數模型
DeepSpeed 可以只用 800 張英偉達 V100 GPU，訓練具有一個萬億參數的語言模型（圖 3）。展示了模型大小和訓練吞吐量，可以觀察到顯存和計算效率，隨模型的大小的擴展線性增長。在各種配置中，在每個 GPU 上訓練，大約 14 億個參數，這是單個 GPU 在不耗盡內存的情況下，可以支持的最大模型大小，表明了完美的顯存擴展性。獲得了接近完美的線性計算效率擴展，每張 V100 GPU 的吞吐量為 47 Tflops。對於上述的硬件，這是令人印象深刻的擴展性和吞吐量。
圖3：模型大小（以十億個參數為單位）和訓練吞吐量（以 Pflops 為單位）隨 GPU 數量變化趨勢的圖表。DeepSpeed 可以使用 800 張，具有 32 GB 內存的英偉達 V100 Tensor Core GPU，訓練有 1 萬億個參數的模型。每種配置都使用 NVIDIA Megatron-LM 提供的，16路模型並行性，剩余的GPU，負責進行流水線並行。萬億參數模型，具有 298 層 Transformer，隱藏層大小為 17408，訓練的序列長度為 2048，batch 大小 2048。對於較小的模型，根據 GPU 數量，按比例減少了 Transformer 層的數量和 batch 大小。
深入研究 3D 並行，如何加速訓練 GPT-3 規模的模型
圖 4：使用 2D 和 3D 並行，使用 800 個 GPU，訓練具有 1800 億參數的 GPT-3 規模模型的系統性能。模型具有 100 個 Transformer 層，隱藏層尺寸為 12288，有 96 個 attention head。訓練使用的 batch 大小為 2048，序列長度為 2048。ZeRO-1可以跟數據並行結合使用。P、M 和 D 分別表示流水線，模型和數據並行維度。
在圖 4 中，使用具有超過 1,750 億個參數的最新 GPT-3 模型架構，作為 3D 並行性的基准：
- 首先評估了 2D 配置（C1-C3）。配置 C1 和 C2 僅使用流水線和模型並行——可以訓練模型，由於過度分解模型，導致吞吐量較低，GPU 利用率較低。C3 嘗試僅使用流水線和數據並行，不通過 Megatron 的模型並行，減少激活量，無法解決顯存不足的問題。
- 3D 配置（C4-C10）依次增加了流水線並行度；中間的平衡了並行性的配置，實現最佳性能，實現了顯存，計算和通信效率三高。
- 最佳的 3D 方法，每個GPU可實現 49 Tflops，超過硬件的理論峰值的 40％。
混合並行如何在低帶寬集群上，7 倍加速訓練 GPT-2
訓練了一個 15 億參數的 GPT-2 模型，在圖 5，展示了混合並行的通信優勢。為了突出展示訓練的通信階段，訓練在節點間帶寬較低的四節點的群集上進行：
- 模型並行在這種情況下沒有優勢，因為模型較小，節點內帶寬較低。
- 流水線並行的通信量，比配置數據和模型並行的情況，小一個數量級。在 batch 較小時，訓練速度快 7 倍。
- 數據並行使用通過梯度累積增加 batch 大小，均攤通信開銷，在更大的 batch 大小下，配置了流水線並行的情況的性能，仍是數據並行的兩倍。
- 混合流水線和數據並行配置，通過將數據並行組限制在節點內的 GPU 上，避免了梯度通信瓶頸，因此梯度通信受益於更快的節點內帶寬。
圖 5：在訓練序列長度為 1024 的 GPT-2（1.5B 參數）時，吞吐量與 batch 大小的關系。使用四個節點，每個節點配備四個具有 16 GB 內存的 V100 GPU 訓練。GPU 之間用每秒 50 Gbps 的節點內帶寬和 4 Gbps 的節點間帶寬連接。DP 表示啟用 ZeRO-1 的數據並行性。所有方法都通過增加梯度累積的步數，擴展批量大小。
ZeRO-Offload：單 GPU 訓練 10 倍大的模型
ZeRO-Offload 通過同時利用GPU和宿主機 CPU 的計算和存儲資源，提升了較少的 GPU 資源下，可以高效訓練的最大模型規模。在單張 V100 上，進行最高至 1300 億參數的模型訓練，10 倍於當前最高水平，保持每 GPU 30Tflop 的高訓練吞吐量。
通過使單 GPU，具備訓練數十億參數的模型的能力，ZeRO-Offload 讓大模型訓練變得親民，讓硬件資源有限的深度學習，從業者也能參與其中。
在單 GPU 上使用默認的 PyTorch 和 ZeRO-Offload，能訓練的最大模型規模的柱狀圖。
圖 6：可以在單 GPU 上使用默認的 PyTorch 和 ZeRO-Offload 訓練的最大的模型規模。
ZeRO-Offload 背后的核心技術，在 ZeRO-2 的基礎上，將優化器狀態和梯度，卸至 CPU 內存。讓 ZeRO-Offload 能最大程度，降低拷貝至 CPU 導致的計算效率損失，達到和 ZeRO-2 相同，甚至有時超過的效率。下圖展示了 Zero-OffLoad 的架構：
圖7: ZeRO-Offload 概述。
了解 ZeRO-Offload 是如何在單GPU上，訓練數十億參數模型的
訓練 GPT 和 T5，這樣有數十億參數的模型，需要多個 GPU 來存儲模型和狀態量。大模型訓練大多通過跨 GPU 的模型並行，解決顯存限制問題。最近，發布了 ZeRO，一個高效利用顯存的優化器，將模型狀態量（優化器狀態量、梯度和模型參數），分布在多個並行 GPU 上，讓數十億參數模型，在不使用模型並行的情況下，進行訓練。ZeRO 還是需要大量數據並行的 GPU，保存划分后的模型狀態量，只有少數人，有條件進行這種模型訓練。
ZeRO-Offload 讓單 GPU 可以進行大模型訓練，使這種訓練變得平民化。為了在不使用多個 GPU 的情況下，訓練數十億個參數的模型，ZeRO-Offload 繼承了 ZeRO-2 的划分優化器狀態量和梯度的方法。與ZeRO-2 不同，ZeRO-Offload沒有在每個 GPU 上保存一部分優化器狀態量和梯度，把兩者都移到了本機內存上。Optimizer 狀態，在整體訓練過程中，都保存在內存中。梯度則是在反向計算過程中，在 GPU 上進行計算，通過 reduce-scatter 進行平均，每個數據並行進程，把平均后的梯度，卸到 CPU 上（圖7中的 g offload），棄掉不屬於自己負責的部分。
一旦梯度到了 CPU 上，划分后的優化狀態量，就會並行地在 CPU 上，進行更新(圖7中的 p update)。在更新進行完后，划分后的參數就被移回GPU並用 all gather 操作進行更新 (圖7中的 g swap）。Zero-Offload 也通過使用不同 CUDA stream，重疊通信（如 g offload 和 g swap）和計算（如反向傳播和 p update)，提高訓練效率。
從模型規模，訓練速度和擴展性看 ZeRO-Offload 的優勢
10 倍模型擴展：在單張 32GB V100 GPU 上，圖 6 顯示 PyTorch 能最多訓練有 13 億個參數的模型，ZeRO-Offload 能訓練 130 億個參數的模型， PyTorch 的 10 倍。ZeRO-Offload 在整個訓練過程中，將消耗了大部分 GPU 顯存的優化器狀態，保留在本機內存中，在反向傳播過程中，將計算出來的梯度移至 CPU。節省的 GPU 顯存，可用於訓練更大的模型。
高效的訓練吞吐量：如圖 8 所示，在訓練 100 億參數模型時，即使僅使用單個 GPU 進行訓練，使用 ZeRO-Offload，可讓每個 GPU 有超過 30 Tflops 的吞吐量，吞吐量隨 GPU 數量增長，呈近完美的線性增長。
ZeRO-Offload 是 ZeRO-2 的完美補充，支持在少量 GPU 上，高效訓練大型模型。通過利用 CPU 內存，減少了模型所需的 GPU 顯存，ZeRO-Offload 讓在 1 到 16 個 GPU，訓練大模型變得可行。在 32 個 GPU 上，ZeRO-Offload 的性能，略高於 ZeRO-2; 性能提升來源於 ZeRO-Offload 節省的 GPU 顯存，可以在更大 batch 下，訓練了模型，盡管存在拷貝至 CPU 的開銷，GPU 計算效率仍然可以提高。在有更多的 GPU（例如 64 和 128）的情況下，ZeRO-2 的性能，優於 ZeRO-Offload，兩者都可以運行類似大小的batch，ZeRO-2 沒有將數據移至 CPU 的開銷，GPU 上進行優化器更新，比 CPU 上快得多。總之，ZeRO-Offload 是 ZeRO-2 的補充，擴展了 ZeRO 家族的優化范圍，從單台設備到數千台設備，有大型模型訓練的優化方案。
使用 ZeRO-Offload 和 ZeRO-2 在 128 張 GPU 上，訓練有 100 億參數的 GPT-2 模型的的吞吐量的柱狀圖。
圖 8：使用 128 張 GPU，訓練 100 億參數 GPT-2 模型的 ZeRO-Offload 和 ZeRO-2 的訓練吞吐量比較。
DeepSpeed 稀疏注意力機制：以 6 倍快的速度，執行 10 倍長的序列
基於注意力機制的深度學習模型（例如，Transformers），在捕獲輸入序列中的 token 之間的關系（即使是兩者之間距離很長）方面，非常有效。與文本，圖像和語音相關的輸入配合使用。這些輸入的序列長度可至數千 token。注意力模塊有效地捕獲了長序列內的依賴關系，對長序列輸入的支持，計算量和顯存的限制。計算量和顯存需求，關於序列長度\(n\)呈二次方級增長。
為了解決此限制，DeepSpeed 提供了一套稀疏注意力 kernel——一種工具性技術，通過塊狀稀疏計算，將注意力計算的計算和顯存需求，降低幾個數量級。這套工具不僅緩解了注意力計算的內存瓶頸，稀疏計算非常高效。API 可以方便地集成進，任何基於 Transformer 的模型。除了提供各種稀疏結構外，可以靈活處理任何用戶，自定義的塊狀稀疏結構。
稀疏注意力（SA）可以設計計算，靠近的 token 之間的局部注意力，或通過使用局部注意力計算，得到 summary token，進而得到全局注意力。SA 既支持隨機注意力，也支持局部、全局和隨機注意力的任意組合，如圖 10 中的藍色，橙色和綠色塊。使SA將內存占用減小到\(O(wn)\)，其中1\(<w≤n \)是一個參數，取決於注意力結構。
彩色小方塊顯示可變的稀疏度結構
圖 10：可變稀疏結構
在 GPU 上的高效實現：盡管稀疏注意力的基本實現，節省顯存，但在計算上，可能會比稠密計算要差。稀疏數據導致了內存訪問的分散性。開發高效的稀疏內核，通常是頗具挑戰性的，尤其是在 GPU 上。DeepSpeed 提供了在 Triton 中開發的高效的稀疏注意力 kernel。這些 kernel 呈塊狀稀疏范式結構，實現對齊的內存訪問，減少GPU線程分支平衡處理器上的工作負載。
系統性能：如圖11所示，SA 支持 10 倍長的序列和最高 6.3 倍的計算提速。左圖顯示了可在 BERT-Base 和 BERT-Large 中，運行的最長序列長度。實驗有以下三種設置：稠密模式，具有激活 checkpoint 的稠密模式和具有激活 checkpoint 的稀疏（SA）模式。與 BERT-Base 和 BERT-Large 的稠密模式相比，SA 的序列分別長 10 倍和 16 倍。與稠密模式相比，SA 減少了總計算量，提高了訓練速度：提高的效率隨着序列長度的增加而提高，對於 BERT-Base 而言，提升速度高達 6.3 倍，對於 BERT-Large，高達 5.3 倍。
圖11：BERT 模型的可支持的最大序列長度（左）；在單英偉達 V100 GPU 上，訓練具有不同序列長度的 BERT-Base（中）和 BERT-Large（右）的時間。
了解 SA 如何使其准確率與全稠密注意力相當，甚至比它更高
涉及稀疏注意力的相關工作（Sparse Transformer，Longformer，BigBird），均顯示出比全注意力更高的准確性，與經驗一致。除了降低內存開銷和加快計算速度外，在生產模型中觀察到 SA，有更高准確性，更快收斂的情況。下圖說明了訓練基於 BERT 的長文本，理解（序列長度 2048）生產模型的准確性。實驗在以下三種設置中進行：從頭開始進行稠密訓練，從頭開進行 SA 訓練，從使用序列長度為 512 的密集型 checkpoint，繼續進行 SA 訓練。對於從頭開始進行預訓練，SA較於稠密設置收斂的速度更高，精度更好。時間和准確性，從用 SA 繼續訓練，預先訓練好的 checkpoint 的效果，甚至更好。
圖12：長文本理解應用的准確性
了解 SA 與最新的 LongFormer 的比較情況
將 SA 與 Longformer（一種最新的稀疏結構及其實現）進行了比較。在實驗中，SA 使用“Fixed”稀疏性。兩種實現的准確性相當。在系統性能方面，SA在訓練和推斷方面均優於Longformer：
- 運行 Wikitext103 上的預訓練MLM的速度，提高了 1.5 倍
- BERT-Base 的推理速度，提高3倍（batch 大小 1，序列長度 2,048）
處理任何塊狀稀疏結構的靈活性： DeepSpeed 稀疏注意力套件，不針對任何特定的稀疏結構，能有效支持模型研究人員探索任何塊狀稀疏結構。添加了流行的稀疏結構，例如 Fixed（來自OpenAI稀疏Transformer），[BigBird](https://arxiv.org/pdf/2007.14062 .pdf)（來自Google）和BSLongformer（AI2 Longformer的塊稀疏實現）。定義了一個具有“可變”結構的模板，如圖 10 所示，模板可用於簡單地自定義任何隨機，局部或全局注意力模式的塊狀稀疏結構。
1．比特 Adam：減少5倍的通信量，提升 3.4 倍的訓練速度
大型模型（如 BERT 和 GPT-3）的擴展訓練，需要基於模型設計，體系結構和系統功能的細致優化。從系統的角度來看，通信效率已成為主要的瓶頸，在使用標准 TCP 且網絡帶寬有限的商用系統上。
通信壓縮是減少在此類系統上的訓練時間的重要技術。壓縮通信的最有效方法之一，誤差補償壓縮，即使在1比特壓縮下，可以提供穩定的收斂速度。但是，最新的誤差補償技術，僅適用於一些和梯度線性相關的簡單優化器，隨機梯度下降（SGD）和 Momentum SGD。這些技術無法和 Adam 之類的非線性優化器整合，在許多任務（包括訓練類似 BERT 的模型）中，帶來了最好的收斂率和精度。
對於像 Adam 之類的強大優化器，依賴於梯度的非線性特征（在方差項上），開發基於誤差補償的壓縮技術，頗具挑戰性的工作，限制了先進的通信壓縮技術的實用價值。
理解經典壓縮技術的背景
通信壓縮的一種方法是1比特壓縮，可以被表示為：
在這種壓縮中，用 1 比特表示每個數字，將內存需求減少 32 倍。這種直接的方法會大大降低收斂速度，沒什么實用價值。研究表明，通過使用誤差補償壓縮，有望在通信壓縮下，保證幾乎相同的收斂率。
誤差補償的思想，可以概括為：1）進行壓縮，2）記憶壓縮誤差，3）在下一次迭代中，把壓縮誤差加回來。對於 SGD，誤差壓縮相當於：
其中\(C(⋅)\)是1比特壓縮算子。這種誤差壓縮的優點在於壓縮誤差的歷史值\(e_t\)和\(e_t-1\)最終會相互抵消, 這使得：
該策略已經被證明，適用於所有線性依賴於梯度的優化算法，例如 SGD 和 Momentum SGD。
了解將誤差補償應用於 Adam 的挑戰
在下面提供了 Adam 算法的概述。更新規則如下：
如上圖的公式所示，方差項 \(v_t\) 和梯度 \(g_t\) 呈非線程關系。如果對 Adam 進行普通的誤差補償，發現（見圖 13）Adam 將無法收斂。
圖13：由於對梯度的非線性依賴，誤差補償壓縮不適用於 Adam
用 1 比特 Adam 壓縮通信
為了在使用 Adam 優化器時壓縮通信，開發了 1 比特 Adam，通過預處理解決了梯度中的非線性依賴問題。觀察到非線性項方差（\(v_t\)）的變化幅度，在幾個訓練周期后顯著降低，將 \(v_t\) 設置為常數不會改變收斂速度。1 位 Adam 優化器，由兩部分組成（如圖 14 所示）：預熱階段，本質上就是原始的 Adam 算法。壓縮階段，使方差項保持恆定，將剩余的線性項（即動量）壓縮為 1 位表示形式。
算法的壓縮階段，由閾值參數控制（如圖 14 所示）。當檢測到“方差”的變化，降至某個閾值以下時，切換到壓縮階段。研究表明，熱身階段，需要全部訓練步驟的 15-20％。
進一步了解 1 比特 Adam 的底層機制
1 比特 Adam 的權重按以下公式進行更新。對於第 i 個 worker，在壓縮階段：
a screenshot of text
a screenshot of a cell phone
圖 14：使用經典 Adam 算法和使用 1 比特壓縮 Adam 算法，進行分布式訓練的流程對比
應對 1 比特 Adam 的系統挑戰
除了算法上的挑戰外，在訓練系統中應用 1 比特 Adam，有兩個系統挑戰。首先，需要具備將動量轉換為 1 比特表示形式的功能的高效 kernel。其次，需要高效的通信方案，在不同的 GPU 之間，傳輸壓縮后的動量。壓縮的目的是減少總體訓練時間，使帶寬受限的商品系統，用來訓練大型模型。在 DeepSpeed 中，解決了這些具有挑戰性的問題，針對在通信效率受限的系統上，進行訓練的場景，對 1 比特 Adam 實現，進行了全面的優化。
1 比特 Adam 在通信受限系統上的優勢
1 比特 Adam 提供了和 Adam 相同的收斂能力，最多可以減少 5 倍的通信量，進行 BERT-Large 預訓練任務時，可達最高 3.5 倍的吞吐量，用於 SQuAD fine-tuning 任務時，可達 2.7 倍的高吞吐量。端到端吞吐量的提高，來源於在壓縮階段觀察到的 6.6 倍（圖 15 左）和 6.2 倍（圖 15 右）速度提升。1 位 Adam 優化器，在 40 Gb 以太網系統上的擴展性非常好，性能可與 Adam 在 40 Gb InfiniBand QDR 系統上的擴展性相媲美。基於 iPerf 基准，40 Gb 以太網上的有效帶寬為 4.1 Gbps，基於 InfiniBand perftest 微基准，InfiniBand 提供了 32 Gbps 的近峰帶寬。
圖 15：NVIDIA V100 GPU 上的 BERT-Large 預訓練（左）和 SQuAD fine-tuning（右）的 1 比特 Adam 擴展性。BERT 預訓練的 batch 大小為 16/GPU，SQuAD fine-tuning 為 3/GPU。
深入研究 1 比特 Adam 的評測結果
與 Adam 相同的收斂性：使用 1 比特 Adam 的一個主要問題是收斂速度。發現在使用相同數量的訓練樣本時，1 比特 Adam，可以達到相同的收斂速度和相當的性能，見圖 16。
圖 16：使用相同數量的訓練樣本，1 比特 Adam，可以像 Adam 一樣收斂。
表 1 顯示了 BERT-Base 和 BERT-Large 的詳細結果。看到，對於未壓縮和壓縮情況，1 比特 Adam 的性能，均與原始模型相當，有些則優於原始模型。
表 1：在各種測試任務上驗證 1 比特 Adam 的正確性
最多可減少 5 倍的通信量： 1 比特 Adam，提供了與 Adam 相同的收斂能力，在壓縮階段（對於 16 位（FP16）訓練），將通信量減少了 16 倍。 對於 BERT 預訓練模型，觀察到預熱階段，僅為端到端訓練時間的 15％，總體通信減少了 5 倍。
原始 Adam 和 1 比特 Adam 的通信量之比的公式如下：
1 / (warmup + (1 – warmup)/16)
1 比特 Adam 使訓練 BERT-Large 的速度快 3.5 倍： 提供了在兩個具有有限帶寬限制的系統上，訓練 BERT-Large 的結果：1）40 Gbps 以太網（圖 17 左）和 2）40 Gbps InfiniBand QDR（圖 17 右）。在壓縮階段，發現使用以太網的系統吞吐量，提高了 6.6 倍，使用 InfiniBand的系統吞吐量，提高了 2 倍，端到端的速度（包括預熱和壓縮階段），分別提高了 3.5 倍和 2.7 倍。1 比特 Adam 主要得益於通信量的減少（因為對動量通信的壓縮實現），及自定義的 allreduce 操作，操作通過高效的 1 比特，無阻塞 gather 和一個 allgather 操作實現。
使用 LAMB 而不是 Adam 優化器，進行 BERT 預訓練，通過增加總 batch 大小，減少通信量。1 比特的 Adam，避免了這種要求嚴格的超參數調參。根據經驗，大 batch 下進行調參，通常會更加困難。此外，1 比特 Adam，對於臨界批處理量較小（無法在大 batch 下良好收斂，例如許多 fine-tuning 任務）的工作也非常適用。
圖 17：在壓縮階段，使用 1 比特 Adam 在 40 Gbps 以太網（左）和 InfiniBand（右）上，進行 BERT-Large 訓練時的性能
1 比特 Adam 使 SQuAD fine-tuning 任務加速 2.7 倍： 1 比特 Adam，不僅在大規模訓練任務上提供擴展性，而且在 SQuAD 微調之類的任務上，也有效果。如圖 18 所示，1 比特 Adam，在基於以太網和基於 InfiniBand 的系統上很好地擴展，在基於以太網的系統上，提供高達 6.2 倍的高吞吐量（在壓縮階段），帶來端到端的 2.7 倍提速（預熱階段占 25％，壓縮階段占 75％）。對於 SQuAD fine-tuning，觀察到總 batch 大小為 96 時，F1 得分最高。 batch 大小大於此值降低收斂率，需要額外的超參數調整。為了擴展到 32 個 GPU，在每個 GPU 上運行值為 3-4 的小 batch。使得 fine-tuning 任務的通信強度大，難以擴展。1 比特 Adam，很好地解決了擴展性的難題，在不增大 batch 的情況下，減少了 3.4 倍的通信量，實現了 2.7 倍的端到端加速。
圖 18：在 40 Gbps 以太網（左）和 InfiniBand（右）上的 SQuAD fine-tuning 任務中，使用 1 比特 Adam 時，壓縮階段的性能。
DeepSpeed團隊：
一群熱衷於大規模系統性能優化的研究員和工程師——Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Reza Yazdani Aminabadi, Elton Zheng, Arash Ashari, Jing Zhao, Minjia Zhang, Niranjan Uma Naresh, Shaden Smith, Ammar Ahmad Awan, Conglong Li, Yuxiong He (team lead)。
參考鏈接：
https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/
https://zhuanlan.zhihu.com/p/343570325