一、Deepspeed

核心是zero優(yōu)化器

1荚虚、總覽

微軟開(kāi)發(fā)ZeRO是為了克服數(shù)據(jù)并行性和模型并行性的限制洼滚，同時(shí)實(shí)現(xiàn)兩者的優(yōu)點(diǎn)歪脏。ZeRO通過(guò)在數(shù)據(jù)并行進(jìn)程中劃分模型狀態(tài)（參數(shù)良蒸，梯度和優(yōu)化器狀態(tài)）技扼，而不是復(fù)制它們，從而消除了數(shù)據(jù)并行進(jìn)程中的內(nèi)存冗余嫩痰。它在訓(xùn)練期間使用動(dòng)態(tài)通信計(jì)劃剿吻，以在分布式設(shè)備之間共享必要的狀態(tài)，以保持計(jì)算粒度和數(shù)據(jù)并行性的通信量串纺。

ZeRO驅(qū)動(dòng)的數(shù)據(jù)并行性丽旅，它允許每個(gè)設(shè)備的內(nèi)存使用量隨數(shù)據(jù)并行性的程度線性擴(kuò)展，并產(chǎn)生與數(shù)據(jù)并行性相似的通信量纺棺。 ZeRO支持的數(shù)據(jù)并行性可以適合任意大小的模型榄笙，只要聚合的設(shè)備內(nèi)存足夠大以共享模型狀態(tài)即可

2、ZeRO的三個(gè)階段及其優(yōu)勢(shì)

ZeRO具有三個(gè)主要的優(yōu)化階段（如圖1所示）祷蝌，分別對(duì)應(yīng)于優(yōu)化器狀態(tài)茅撞，梯度和參數(shù)的劃分。 累計(jì)啟用時(shí)：

• 1. 優(yōu)化器狀態(tài)分區(qū)(Pos) -減少了4倍的內(nèi)存杆逗，通信容量與數(shù)據(jù)并行性相同
• 2. 增加梯度分區(qū)(Pos+g) - 8x內(nèi)存減少乡翅，通信容量與數(shù)據(jù)并行性相同
• 3.增加參數(shù)分區(qū)(Pos+g+p) -內(nèi)存減少與數(shù)據(jù)并行度和復(fù)雜度成線性關(guān)系鳞疲。

例如罪郊，跨64個(gè)gpu (Nd = 64)進(jìn)行拆分將減少64倍的內(nèi)存。通信容量適度增加了50%尚洽。

ZeRO消除了內(nèi)存冗余悔橄，并使群集的全部聚合內(nèi)存容量可用。 啟用所有三個(gè)階段后腺毫，ZeRO可以僅在1024個(gè)NVIDIA GPU上訓(xùn)練一個(gè)萬(wàn)億參數(shù)模型癣疟。 一個(gè)具有16位精度的類似Adam之類的優(yōu)化器的參數(shù)化萬(wàn)億模型需要大約16 TB的內(nèi)存來(lái)保存優(yōu)化器的狀態(tài)，梯度和參數(shù)潮酒。 16TB除以1024即為16GB睛挚，這在GPU的合理范圍內(nèi)。

ZeRO-Offload 通過(guò)同時(shí)利用GPU和宿主機(jī) CPU 的計(jì)算和存儲(chǔ)資源急黎，提升了較少的 GPU 資源下可以高效訓(xùn)練的最大模型規(guī)模扎狱。它讓我們可以在單張 V100 上進(jìn)行最高至 1300 億參數(shù)的模型訓(xùn)練，10 倍于當(dāng)前最高水平勃教，同時(shí)保持每 GPU 30Tflop 的高訓(xùn)練吞吐量淤击。

通過(guò)使單 GPU 具備訓(xùn)練數(shù)十億參數(shù)的模型的能力，ZeRO-Offload 讓大模型訓(xùn)練變得親民故源，讓硬件資源有限的深度學(xué)習(xí)從業(yè)者也能參與其中污抬。

在單 GPU 上使用默認(rèn)的 PyTorch 和 ZeRO-Offload 能訓(xùn)練的最大模型規(guī)模的柱狀圖

圖 6：可以在單 GPU 上使用默認(rèn)的 PyTorch 和 ZeRO-Offload 訓(xùn)練的最大的模型規(guī)模。

ZeRO-Offload 背后的核心技術(shù)是在 ZeRO-2 的基礎(chǔ)上將優(yōu)化器狀態(tài)和梯度卸至 CPU 內(nèi)存绳军。這個(gè)方法讓 ZeRO-Offload 能最大程度降低拷貝至 CPU 導(dǎo)致的計(jì)算效率損失印机，同時(shí)達(dá)到和 ZeRO-2 相同矢腻，甚至有時(shí)超過(guò)的效率。下圖展示了 Zero-OffLoad 的架構(gòu)：

圖7: ZeRO-Offload 概述耳贬。

了解 ZeRO-Offload 是如何在單GPU上訓(xùn)練數(shù)十億參數(shù)模型的

訓(xùn)練 GPT 和 T5 這樣有數(shù)十億參數(shù)的模型需要多個(gè) GPU 來(lái)存儲(chǔ)模型和狀態(tài)量踏堡。大模型訓(xùn)練大多通過(guò)跨 GPU 的模型并行來(lái)解決顯存限制問(wèn)題。最近咒劲，我們發(fā)布了 ZeRO顷蟆，一個(gè)高效利用顯存的優(yōu)化器，它會(huì)將模型狀態(tài)量（優(yōu)化器狀態(tài)量腐魂、梯度和模型參數(shù)）分布在多個(gè)并行 GPU 上帐偎，讓數(shù)十億參數(shù)模型可以在不使用模型并行的情況下進(jìn)行訓(xùn)練。然而蛔屹，ZeRO 還是需要大量數(shù)據(jù)并行的 GPU 來(lái)保存劃分后的模型狀態(tài)量削樊，因此只有少數(shù)人有條件進(jìn)行這種模型訓(xùn)練。

ZeRO-Offload 讓單 GPU 可以進(jìn)行大模型訓(xùn)練兔毒，從而使這種訓(xùn)練變得平民化漫贞。為了在不使用多個(gè) GPU 的情況下訓(xùn)練數(shù)十億個(gè)參數(shù)的模型，ZeRO-Offload 繼承了 ZeRO-2 的劃分優(yōu)化器狀態(tài)量和梯度的方法育叁。和 ZeRO-2 不同之處在于迅脐，ZeRO-Offload 并沒(méi)有在每個(gè) GPU 上保存一部分優(yōu)化器狀態(tài)量和梯度，而是把兩者都移到了本機(jī)內(nèi)存上豪嗽。Optimizer 狀態(tài)在整體訓(xùn)練過(guò)程中都保存在內(nèi)存中谴蔑。梯度則是在反向計(jì)算過(guò)程中在 GPU 上進(jìn)行計(jì)算并通過(guò) reduce-scatter 進(jìn)行平均，之后每個(gè)數(shù)據(jù)并行進(jìn)程把自己的那份平均后的梯度卸到 CPU 上（圖7中的 g offload）并棄掉不屬于自己負(fù)責(zé)的部分龟梦。

一旦梯度到了 CPU 上隐锭，劃分后的優(yōu)化狀態(tài)量就會(huì)并行地在 CPU 上進(jìn)行更新(圖7中的 p update)。在更新進(jìn)行完后计贰，劃分后的參數(shù)就被移回GPU并用 all gather 操作進(jìn)行更新 (圖7中的 g swap）钦睡。Zero-Offload 也通過(guò)使用不同 CUDA stream 來(lái)重疊通信（如 g offload 和 g swap）和計(jì)算（如反向傳播和 p update) 以提高訓(xùn)練效率。

從模型規(guī)模躁倒，訓(xùn)練速度和擴(kuò)展性看 ZeRO-Offload 的優(yōu)勢(shì)

10 倍模型擴(kuò)展：在單張 32GB V100 GPU 上荞怒，圖 6 顯示 PyTorch 能最多訓(xùn)練有 13 億個(gè)參數(shù)的模型，而 ZeRO-Offload 能訓(xùn)練 130 億個(gè)參數(shù)的模型樱溉，是 PyTorch 的 10 倍挣输。這是因?yàn)?ZeRO-Offload 在整個(gè)訓(xùn)練過(guò)程中將消耗了大部分 GPU 顯存的優(yōu)化器狀態(tài)保留在本機(jī)內(nèi)存中，同時(shí)還在反向傳播過(guò)程中將計(jì)算出來(lái)的梯度移至 CPU福贞。因此撩嚼，節(jié)省的 GPU 顯存可用于訓(xùn)練更大的模型。

高效的訓(xùn)練吞吐量：如圖 8 所示，在訓(xùn)練 100 億參數(shù)模型時(shí)完丽，即使僅使用單個(gè) GPU 進(jìn)行訓(xùn)練恋技，使用 ZeRO-Offload 仍可讓每個(gè) GPU 有超過(guò) 30 Tflops 的吞吐量，并且其吞吐量隨 GPU 數(shù)量增長(zhǎng)呈近完美的線性增長(zhǎng)逻族。

ZeRO-Offload 是 ZeRO-2 的完美補(bǔ)充蜻底，支持在少量 GPU 上高效訓(xùn)練大型模型。通過(guò)利用 CPU 內(nèi)存來(lái)減少了模型所需的 GPU 顯存聘鳞，ZeRO-Offload 讓在 1 到 16 個(gè) GPU 上訓(xùn)練大模型變得可行薄辅。在 32 個(gè) GPU 上，ZeRO-Offload 的性能略高于 ZeRO-2; 性能提升來(lái)源于 ZeRO-Offload 節(jié)省的 GPU 顯存抠璃，它們讓我們可以在更大 batch 下訓(xùn)練了模型站楚，因此盡管存在拷貝至 CPU 的開(kāi)銷忧侧，GPU 計(jì)算效率仍然可以提高东囚。在有更多的 GPU（例如 64 和 128）的情況下，ZeRO-2 的性能優(yōu)于 ZeRO-Offload政勃，因?yàn)閮烧攥F(xiàn)在都可以運(yùn)行類似大小的batch采盒，ZeRO-2 沒(méi)有將數(shù)據(jù)移至 CPU 的開(kāi)銷旧乞，并且 GPU 上進(jìn)行優(yōu)化器更新要比 CPU 上快得多“醢保總而言之尺栖，ZeRO-Offload 是 ZeRO-2 的補(bǔ)充，并擴(kuò)展了 ZeRO 家族的優(yōu)化范圍悍赢，從單臺(tái)設(shè)備到數(shù)千臺(tái)設(shè)備决瞳，都有大型模型訓(xùn)練的優(yōu)化方案货徙。

使用 ZeRO-Offload 和 ZeRO-2 在 128 張 GPU 上訓(xùn)練有 100 億參數(shù)的 GPT-2 模型的的吞吐量的柱狀圖左权。

3、DeepSpeed

實(shí)現(xiàn)了ZeRO的第一階段（優(yōu)化器狀態(tài)分區(qū)（簡(jiǎn)稱ZeRO-OS））痴颊，具有支持1000億參數(shù)模型的強(qiáng)大能力赏迟。 該代碼將與我們的培訓(xùn)優(yōu)化庫(kù)DeepSpeed一起發(fā)布。 DeepSpeed通過(guò)與PyTorch兼容的輕量級(jí)API帶來(lái)了最新的培訓(xùn)技術(shù)蠢棱，例如ZeRO锌杀，分布式培訓(xùn)，混合精度和檢查點(diǎn)泻仙。 只需對(duì)PyTorch模型進(jìn)行幾行代碼更改糕再，就可以利用DeepSpeed解決潛在的性能挑戰(zhàn)，并提高培訓(xùn)速度和規(guī)模玉转。 DeepSpeed在四個(gè)方面表現(xiàn)出色(如圖2所示):

（1）突想、3大基礎(chǔ)并行方式

• 數(shù)據(jù)并行是深度學(xué)習(xí)中的一種普遍使用的技術(shù)。在該技術(shù)中，每批輸入的訓(xùn)練數(shù)據(jù)都在數(shù)據(jù)并行的 worker 之間平分猾担。反向傳播后需要通信并規(guī)約梯度袭灯，以保證優(yōu)化器在各個(gè) worker 上進(jìn)行相同的更新。數(shù)據(jù)并行性具有幾個(gè)明顯的優(yōu)勢(shì)绑嘹，包括計(jì)算效率高和實(shí)現(xiàn)起來(lái)工作量小稽荧。但是，數(shù)據(jù)并行的 batch 大小隨 worker 數(shù)量提高工腋，而我們往往無(wú)法在不影響收斂性的情況下一直增加 batch 大小姨丈。

• 顯存效率：數(shù)據(jù)并行會(huì)在所有 worker 之間進(jìn)行模型和優(yōu)化器的復(fù)制，因此顯存效率不高擅腰。DeepSpeed 開(kāi)發(fā)了 ZeRO 构挤，它是一系列用于提高數(shù)據(jù)并行的顯存效率的優(yōu)化器。 這項(xiàng)工作依賴于 ZeRO 的 1 階段惕鼓，該階段在 worker 之間劃分優(yōu)化器狀態(tài)量以減少冗余筋现。
• 計(jì)算效率：隨著我們提高并行度，每個(gè) worker 執(zhí)行的計(jì)算量是恒定的箱歧。數(shù)據(jù)并行可以在小規(guī)模上實(shí)現(xiàn)近乎線性擴(kuò)展矾飞。但是，在 worker 之間規(guī)約梯度的通信開(kāi)銷跟模型大小成正相關(guān)呀邢，所以當(dāng)模型很大或通信帶寬很低時(shí)洒沦，計(jì)算效率會(huì)受限。价淌。梯度累積是一種用來(lái)均攤通信成本的一種常用策略申眼。它會(huì)進(jìn)一步增加batch大小，在本地使用 micro-batch 多次進(jìn)行正向和反向傳播積累梯度后蝉衣，再進(jìn)行梯度規(guī)約和優(yōu)化器更新括尸。

• 模型并行是包含范圍很廣的一類技術(shù)。它會(huì)在多個(gè) worker 之間劃分模型的各個(gè)層病毡。就其本質(zhì)而言濒翻，模型并行性的計(jì)算和通信因模型結(jié)構(gòu)而異，因此在實(shí)現(xiàn)上有很大的工作量啦膜。DeepSpeed 借用了英偉達(dá)的 Megatron-LM 來(lái)為基于 Transformer 的語(yǔ)言模型提供大規(guī)模模型并行功能有送。模型并行會(huì)根據(jù) worker 數(shù)量成比例地減少顯存使用量，也是這三種并行度中顯存效率最高的僧家。但是其代價(jià)是計(jì)算效率最低雀摘。

• 顯存效率：模型并行會(huì)根據(jù) worker 數(shù)量成比例地減少顯存使用量。至關(guān)重要的是八拱，這是減少單個(gè)網(wǎng)絡(luò)層的激活顯存的唯一方法阵赠。DeepSpeed 通過(guò)在模型并行 worker 之間劃分激活顯存來(lái)進(jìn)一步提高顯存效率烁落。
• 計(jì)算效率：由于每次前向和反向傳播中都需要額外通信激活值，模型并行的計(jì)算效率很低豌注。模型并行需要高通信帶寬伤塌，并且不能很好地?cái)U(kuò)展到通信帶寬受限的節(jié)點(diǎn)。此外轧铁，每個(gè)模型并行worker 都會(huì)減少每個(gè)通信階段之間執(zhí)行的計(jì)算量每聪，從而影響計(jì)算效率。模型并行性通常與數(shù)據(jù)并行性結(jié)合使用齿风，以在內(nèi)存和計(jì)算效率之間進(jìn)行權(quán)衡药薯。

• 流水線并行訓(xùn)練引擎也被包含在了這次發(fā)布的DeepSpeed中！流水線并行將模型的各層劃分為可以并行處理的階段救斑。當(dāng)一個(gè)階段完成一個(gè) micro-batch 的正向傳遞時(shí)童本，激活內(nèi)存將被通信至流水線的下一個(gè)階段。類似地脸候，當(dāng)下一階段完成反向傳播時(shí)穷娱，將通過(guò)管道反向通信梯度。必須同時(shí)計(jì)算多個(gè) micro-batch 以確保流水線的各個(gè)階段能并行計(jì)算运沦。目前已經(jīng)開(kāi)發(fā)出了幾種用于權(quán)衡內(nèi)存和計(jì)算效率以及收斂行為的方法泵额，例如 PipeDream。DeepSpeed 采用的方法是通過(guò)梯度累積來(lái)實(shí)現(xiàn)并行携添，并保持與傳統(tǒng)數(shù)據(jù)并行和模型并行訓(xùn)練在相同的總 batch 大小下收斂情況相同嫁盲。

• 顯存效率：模型并行會(huì)根據(jù) worker 數(shù)量成比例地減少顯存使用量。至關(guān)重要的是烈掠，這是減少單個(gè)網(wǎng)絡(luò)層的激活顯存的唯一方法羞秤。DeepSpeed 通過(guò)在模型并行 worker 之間劃分激活顯存來(lái)進(jìn)一步提高顯存效率。
• 計(jì)算效率：由于每次前向和反向傳播中都需要額外通信激活值左敌，模型并行的計(jì)算效率很低瘾蛋。模型并行需要高通信帶寬，并且不能很好地?cái)U(kuò)展到通信帶寬受限的節(jié)點(diǎn)母谎。此外瘦黑，每個(gè)模型并行worker 都會(huì)減少每個(gè)通信階段之間執(zhí)行的計(jì)算量京革，從而影響計(jì)算效率奇唤。模型并行性通常與數(shù)據(jù)并行性結(jié)合使用，以在內(nèi)存和計(jì)算效率之間進(jìn)行權(quán)衡匹摇。

此外咬扇，流水線并行性會(huì)在每個(gè) batch 的開(kāi)始和結(jié)束時(shí)因?yàn)樾枰匦绿畛浠蚺趴樟魉€而產(chǎn)生 bubble overhead。使用流水線階段數(shù)的 4 倍或 8 倍的梯度累積步驟（以及 batch 大欣炔）進(jìn)行訓(xùn)練懈贺，相較于只有一個(gè)流水線階段分別達(dá)到了 81％ 和 90％ 的擴(kuò)展性经窖。

（2）、3D混合并行

數(shù)據(jù)梭灿，模型和流水線并行在提高內(nèi)存和計(jì)算效率方面均起到特定的作用画侣。圖 1 說(shuō)明了我們的 3D 策略。

• 顯存效率：先將模型的各層劃分到不同的流水線階段堡妒，并進(jìn)一步把每個(gè)階段的層通過(guò)模型并行進(jìn)行劃分配乱。這種 2D 組合同時(shí)減少了模型、優(yōu)化器和激活函數(shù)所消耗的內(nèi)存皮迟。不過(guò)搬泥，我們不能在不引入通信開(kāi)銷的情況下無(wú)限地劃分模型，而通信開(kāi)銷會(huì)限制計(jì)算效率伏尼。
• 計(jì)算效率：為了在不犧牲計(jì)算效率的情況下將 worker 數(shù)量擴(kuò)展至超出模型和流水線并行能支持的規(guī)模忿檩，我們使用了 ZeRO 支持的數(shù)據(jù)并行功能（ZeRO-DP）。ZeRO-DP 不僅可以通過(guò)劃分優(yōu)化器狀態(tài)量進(jìn)一步提高顯存利用效率爆阶，而且還可以通過(guò)利用基于通信拓?fù)涞挠成潢P(guān)系燥透，以最小的通信開(kāi)銷擴(kuò)展到任意數(shù)量的 GPU。

1. 基于通信拓?fù)涞?3D 映射（圖2）：通過(guò)利用兩個(gè)關(guān)鍵的架構(gòu)屬性辨图，我們將 3D 并行中的每個(gè)維度仔細(xì)地映射到 worker 上兽掰，以實(shí)現(xiàn)最大的計(jì)算效率。
   優(yōu)化節(jié)點(diǎn)內(nèi)和節(jié)點(diǎn)間的通信帶寬：模型并行是這三種策略中通信開(kāi)銷最大的徒役，因此我們優(yōu)先考慮將模型并行 worker 組放置在節(jié)點(diǎn)內(nèi)以利用更大的節(jié)點(diǎn)內(nèi)帶寬孽尽。這里我們基于英偉達(dá) Megatron-LM 進(jìn)行了張量切分式的模型并行。當(dāng)模型并行組不占滿節(jié)點(diǎn)內(nèi)的所有 worker 時(shí)忧勿，我們選擇將數(shù)據(jù)并行組放置在節(jié)點(diǎn)內(nèi)杉女。不然就跨節(jié)點(diǎn)進(jìn)行數(shù)據(jù)并行。流水線并行的通信量最低鸳吸，因此我們可以跨節(jié)點(diǎn)調(diào)度流水線的各個(gè)階段熏挎，而不受通信帶寬的限制。

2. 通過(guò)并行通信增大帶寬：每個(gè)數(shù)據(jù)并行組需要通信的梯度量隨著流水線和模型并行的規(guī)模線性減小晌砾，因此總通信量少于單純使用數(shù)據(jù)并行坎拐。此外，每個(gè)數(shù)據(jù)并行組會(huì)在局部的一小部分 worker 內(nèi)部獨(dú)立進(jìn)行通信养匈，組間通信可以相互并行哼勇。這樣的結(jié)果是，通過(guò)減少通信量和增加局部性與并行性呕乎，數(shù)據(jù)并行通信的有效帶寬被增大了

   
   
   該圖顯示了一個(gè)有 32 個(gè) worker 進(jìn)行 3D 并行的例子积担。神經(jīng)網(wǎng)絡(luò)的各層分為四個(gè)流水線階段。每個(gè)流水線階段中的層在四個(gè)模型并行 worker 之間進(jìn)一步劃分猬仁。最后帝璧，每個(gè)流水線階段有兩個(gè)數(shù)據(jù)并行實(shí)例先誉，且 ZeRO 在這 2 個(gè)副本之間劃分優(yōu)化器狀態(tài)量。
   
   

   圖 1：一個(gè)有 32 個(gè) worker 進(jìn)行 3D 并行的例子的烁。神經(jīng)網(wǎng)絡(luò)的各層分為四個(gè)流水線階段褐耳。每個(gè)流水線階段中的層在四個(gè)模型并行 worker 之間進(jìn)一步劃分。最后渴庆，每個(gè)流水線階段有兩個(gè)數(shù)據(jù)并行實(shí)例漱病，且 ZeRO 在這 2 個(gè)副本之間劃分優(yōu)化器狀態(tài)量。

   
   
   彩色塊顯示圖 1 中的 worker 到八個(gè)節(jié)點(diǎn)（每個(gè)節(jié)點(diǎn)有四個(gè) GPU）的系統(tǒng)上的 GPU 的映射把曼。同一顏色的 GPU 在同一節(jié)點(diǎn)上杨帽。

具體理解方式：
首先stage0先算，stage0算完之后嗤军，stage1算注盈，同時(shí)stage0再算，完成流水線并行叙赚。
另外老客，stage0內(nèi)部算的時(shí)候，7層模型會(huì)通過(guò)megatron進(jìn)行模型并行的計(jì)算震叮。
圖 2：圖 1 中的 worker 到八個(gè)節(jié)點(diǎn)（每個(gè)節(jié)點(diǎn)有四個(gè) GPU）的系統(tǒng)上的 GPU 的映射胧砰。同一顏色的 GPU 在同一節(jié)點(diǎn)上。
了解關(guān)于 3D 并行訓(xùn)練萬(wàn)億參數(shù)模型的更多信息

使用 8 路模型并行苇瓣，64 路流水線并行和 8 路數(shù)據(jù)并行尉间，可以在 4096 個(gè)英偉達(dá) A100 GPU 上擴(kuò)展訓(xùn)練一個(gè)萬(wàn)億參數(shù)模型。

通過(guò)結(jié)合模型并行和流水線并行击罪，3D 并行可實(shí)現(xiàn)出色的內(nèi)存效率和跨多個(gè)節(jié)點(diǎn)的高效計(jì)算效率哲嘲。模型并行性提高了節(jié)點(diǎn)內(nèi)的激活內(nèi)存和模型狀態(tài)量的存儲(chǔ)效率，而流水線并行媳禁，相較于僅使用模型并行眠副，則可以在不犧牲計(jì)算效率的情況下，跨節(jié)點(diǎn)高效存儲(chǔ)模型狀態(tài)竣稽。在 micro-batch 大小為 1 的萬(wàn)億參數(shù)例子中囱怕，在使用激活值 checkpoint 以及上述 3D 并行后，模型狀態(tài)量會(huì)消耗 30 GB 的顯存毫别，劃分后的激活值消耗 2.5 GB 的內(nèi)存娃弓。這樣總顯存占用為 32.5 GB，就能夠使用具有 40 GB 內(nèi)存的英偉達(dá) A100 GPU 來(lái)容納和訓(xùn)練這樣的模型了拧烦。

結(jié)合模型并行與流水線并行忘闻，可以使流水線并行在非常小的 batch 下以最小的 bubble overhead 實(shí)現(xiàn)高計(jì)算效率。在 8 路模型并行下恋博，每個(gè)模型使用 micro-batch 為 1 個(gè)微批處理將導(dǎo)致每個(gè) GPU 的有效 micro-batch 大小為 1/8齐佳。因此，使用 8 倍于管道并行度的梯度累加步驟债沮，只會(huì)讓每張 GPU 上的總累計(jì) batch 大小為 1炼吴，并且流水并行處理可以實(shí)現(xiàn) 90％ 的計(jì)算效率。與數(shù)據(jù)并行性結(jié)合使用時(shí)疫衩，這讓 4096 張 GPU 上的總有效 batch 大小為 4096硅蹦，并仍然可以達(dá)到 90％ 的流水線效率。

思考：但是數(shù)據(jù)并行會(huì)怎樣影響計(jì)算效率呢闷煤？難道數(shù)據(jù)并行不是需要每張 GPU 都有大 batch 才能保持高效嗎童芹？

模型并行可以將每張GPU上的有效 batch 大小減小到小于 1。這使流水線并行即使在小 batch 下仍可以隱藏流水線 bubble overhead鲤拿。請(qǐng)注意假褪，通過(guò)跨節(jié)點(diǎn)使用流水線并行性，我們就可以讓流水線每個(gè)階段的數(shù)據(jù)并行節(jié)點(diǎn)之間的獨(dú)立進(jìn)行通信近顷，并且與其他流水線階段并行進(jìn)行生音。實(shí)際上，在高端 GPU 集群中常見(jiàn)的完全連接的網(wǎng)絡(luò)拓?fù)渲兄仙@對(duì)可用于數(shù)據(jù)并行訓(xùn)練的有效通信帶寬具有重要意義缀遍。由于流水線階段中的每個(gè)節(jié)點(diǎn)都可以與其對(duì)應(yīng)的數(shù)據(jù)并行節(jié)點(diǎn)并行通信，因此有效的通信帶寬與流水線階段數(shù)成正比饱须。通過(guò)設(shè)置64個(gè)并行流水線階段域醇，有效帶寬將變?yōu)橥祮蝹€(gè)節(jié)點(diǎn)的帶寬的 64 倍。流水線并行帶給數(shù)據(jù)并行如此大的有效帶寬蓉媳，這使數(shù)據(jù)并行在計(jì)算與通信比率非常低的小 batch 情況下歹苦，也能實(shí)現(xiàn)高效擴(kuò)展。

（3）督怜、DeepSpeed 稀疏注意力機(jī)制：以 6 倍快的速度執(zhí)行 10 倍長(zhǎng)的序列

基于注意力機(jī)制的深度學(xué)習(xí)模型（例如殴瘦，Transformers）在捕獲輸入序列中的 token 之間的關(guān)系（即使是兩者之間距離很長(zhǎng)）方面非常有效。因此号杠，它們常與文本蚪腋，圖像和語(yǔ)音相關(guān)的輸入配合使用。這些輸入的序列長(zhǎng)度可至數(shù)千 token姨蟋。然而屉凯，盡管注意力模塊有效地捕獲了長(zhǎng)序列內(nèi)的依賴關(guān)系，在實(shí)際應(yīng)用中眼溶，對(duì)長(zhǎng)序列輸入的支持受計(jì)算量和顯存的限制悠砚。計(jì)算量和顯存需求關(guān)于序列長(zhǎng)度(n)呈二次方級(jí)增長(zhǎng)。

為了解決此限制堂飞，DeepSpeed 提供了一套稀疏注意力 kernel——它是一種工具性技術(shù)灌旧，可以通過(guò)塊狀稀疏計(jì)算將注意力計(jì)算的計(jì)算和顯存需求降低幾個(gè)數(shù)量級(jí)绑咱。這套工具不僅緩解了注意力計(jì)算的內(nèi)存瓶頸，而且其稀疏計(jì)算非常高效枢泰。它的 API 可以方便地集成進(jìn)任何基于 Transformer 的模型描融。除了提供各種稀疏結(jié)構(gòu)外，它還可以靈活處理任何用戶自定義的塊狀稀疏結(jié)構(gòu)衡蚂。

更具體地說(shuō)窿克，稀疏注意力（SA）可以設(shè)計(jì)計(jì)算靠近的 token 之間的局部注意力，或通過(guò)使用局部注意力計(jì)算得到 summary token毛甲，進(jìn)而得到全局注意力年叮。此外，SA 既支持隨機(jī)注意力玻募，也支持局部只损、全局和隨機(jī)注意力的任意組合，如圖 10 中的藍(lán)色补箍，橙色和綠色塊改执。這使SA將內(nèi)存占用減小到(O(wn))，其中1(<w≤n )是一個(gè)參數(shù)坑雅，其值取決于注意力結(jié)構(gòu)辈挂。

彩色小方塊顯示可變的稀疏度結(jié)構(gòu)

系統(tǒng)性能：如圖11所示苞笨，SA 支持 10 倍長(zhǎng)的序列和最高 6.3 倍的計(jì)算提速债朵。左圖顯示了可在 BERT-Base 和 BERT-Large 中運(yùn)行的最長(zhǎng)序列長(zhǎng)度。我們的實(shí)驗(yàn)有以下三種設(shè)置：稠密模式瀑凝，具有激活 checkpoint 的稠密模式和具有激活 checkpoint 的稀疏（SA）模式序芦。與 BERT-Base 和 BERT-Large 的稠密模式相比，SA 的序列分別長(zhǎng) 10 倍和 16 倍粤咪。 此外谚中，與稠密模式相比，SA 減少了總計(jì)算量，并提高了訓(xùn)練速度：提高的效率隨著序列長(zhǎng)度的增加而提高宪塔，對(duì)于 BERT-Base 而言磁奖，提升速度高達(dá) 6.3 倍，而對(duì)于 BERT-Large蝌麸，則高達(dá) 5.3 倍点寥。

圖11：BERT 模型的可支持的最大序列長(zhǎng)度（左）艾疟；在單英偉達(dá) V100 GPU 上訓(xùn)練具有不同序列長(zhǎng)度的 BERT-Base（中）和 BERT-Large（右）的時(shí)間来吩。

涉及稀疏注意力的相關(guān)工作（Sparse Transformer，Longformer蔽莱，BigBird）均顯示出比全注意力更高的準(zhǔn)確性弟疆，與我們的經(jīng)驗(yàn)一致。除了降低內(nèi)存開(kāi)銷和加快計(jì)算速度外盗冷，我們還在生產(chǎn)模型中觀察到 SA 有更高準(zhǔn)確性并更快收斂的情況怠苔。下圖說(shuō)明了訓(xùn)練基于 BERT 的長(zhǎng)文本理解（序列長(zhǎng)度 2048）生產(chǎn)模型的準(zhǔn)確性。該實(shí)驗(yàn)在以下三種設(shè)置中進(jìn)行：從頭開(kāi)始進(jìn)行稠密訓(xùn)練仪糖，從頭開(kāi)始進(jìn)行 SA 訓(xùn)練柑司，以及從使用序列長(zhǎng)度為 512 的密集型 checkpoint 繼續(xù)進(jìn)行 SA 訓(xùn)練。我們已經(jīng)觀察到锅劝，對(duì)于從頭開(kāi)始進(jìn)行預(yù)訓(xùn)練攒驰，SA較于稠密設(shè)置收斂的速度更高，精度更好故爵。此外玻粪，就時(shí)間和準(zhǔn)確性而言，從用 SA 繼續(xù)訓(xùn)練預(yù)先訓(xùn)練好的 checkpoint 的效果甚至更好诬垂。

圖12：長(zhǎng)文本理解應(yīng)用的準(zhǔn)確性

（4）1 比特 Adam：減少5倍的通信量并提升 3.4 倍的訓(xùn)練速度

大型模型（如 BERT 和 GPT-3）的擴(kuò)展訓(xùn)練需要基于模型設(shè)計(jì)劲室，體系結(jié)構(gòu)和系統(tǒng)功能的細(xì)致優(yōu)化。從系統(tǒng)的角度來(lái)看结窘，通信效率已成為主要的瓶頸很洋，尤其是在使用標(biāo)準(zhǔn) TCP 且網(wǎng)絡(luò)帶寬有限的商用系統(tǒng)上。

通信壓縮是減少在此類系統(tǒng)上的訓(xùn)練時(shí)間的重要技術(shù)隧枫。壓縮通信的最有效方法之一是誤差補(bǔ)償壓縮喉磁，即使在1比特壓縮下，它也可以提供穩(wěn)定的收斂速度悠垛。但是线定，最新的誤差補(bǔ)償技術(shù)僅適用于一些和梯度線性相關(guān)的簡(jiǎn)單優(yōu)化器，例如隨機(jī)梯度下降（SGD）和 Momentum SGD确买。這些技術(shù)無(wú)法和 Adam 之類的非線性優(yōu)化器整合斤讥，后者在許多任務(wù)（包括訓(xùn)練類似 BERT 的模型）中帶來(lái)了最好的收斂率和精度。

對(duì)于像 Adam 之類的強(qiáng)大優(yōu)化器而言，由于它依賴于梯度的非線性特征（在方差項(xiàng)上）芭商，針對(duì)它來(lái)開(kāi)發(fā)基于誤差補(bǔ)償?shù)膲嚎s技術(shù)是一項(xiàng)頗具挑戰(zhàn)性的工作派草，因此限制了先進(jìn)的通信壓縮技術(shù)的實(shí)用價(jià)值。

理解經(jīng)典壓縮技術(shù)的背景

通信壓縮的一種方法是1比特壓縮铛楣，它可以被表示為：

image.png

在這種壓縮中近迁，我們用 1 比特表示每個(gè)數(shù)字，從而將內(nèi)存需求減少 32 倍簸州。問(wèn)題在于鉴竭，這種直接的方法會(huì)大大降低收斂速度，沒(méi)什么實(shí)用價(jià)值岸浑。最近的研究表明搏存，通過(guò)使用誤差補(bǔ)償壓縮，我們有望在通信壓縮下保證幾乎相同的收斂率矢洲。

誤差補(bǔ)償?shù)乃枷肟梢愿爬椋?）進(jìn)行壓縮璧眠，2）記憶壓縮誤差，然后3）在下一次迭代中把壓縮誤差加回來(lái)读虏。對(duì)于 SGD责静，誤差壓縮相當(dāng)于：

image.png

其中(C(?))是1比特壓縮算子。這種誤差壓縮的優(yōu)點(diǎn)在于壓縮誤差的歷史值(e_t)和(e_t-1)最終會(huì)相互抵消, 這使得：

image.png

了解將誤差補(bǔ)償應(yīng)用于 Adam 的挑戰(zhàn)

我們?cè)谙旅嫣峁┝?Adam 算法的概述。更新規(guī)則如下：

image.png

如上圖的公式所示葱轩，方差項(xiàng) (v_t) 和梯度 (g_t) 呈非線程關(guān)系睦焕。如果我們對(duì) Adam 進(jìn)行普通的誤差補(bǔ)償，我們會(huì)發(fā)現(xiàn)（見(jiàn)圖 13）Adam 將無(wú)法收斂靴拱。

由于對(duì)梯度的非線性依賴垃喊，誤差補(bǔ)償壓縮不適用于 Adam

用 1 比特 Adam 壓縮通信

為了在使用 Adam 優(yōu)化器時(shí)壓縮通信，我們開(kāi)發(fā)了 1 比特 Adam袜炕，它通過(guò)預(yù)處理解決了梯度中的非線性依賴問(wèn)題本谜。我們觀察到非線性項(xiàng)方差（(v_t)）的變化幅度在幾個(gè)訓(xùn)練周期后顯著降低，之后將 (v_t) 設(shè)置為常數(shù)不會(huì)改變收斂速度偎窘。所以提出的 1 位 Adam 優(yōu)化器由兩部分組成（如圖 14 所示）：預(yù)熱階段乌助，本質(zhì)上就是原始的 Adam 算法。壓縮階段陌知，使方差項(xiàng)保持恒定他托，并將剩余的線性項(xiàng)（即動(dòng)量）壓縮為 1 位表示形式。

該算法的壓縮階段由閾值參數(shù)控制（如圖 14 所示）仆葡。當(dāng)我們檢測(cè)到“方差”的變化降至某個(gè)閾值以下時(shí)赏参，就切換到壓縮階段。我們的研究表明，熱身階段只需要全部訓(xùn)練步驟的 15-20％把篓。

進(jìn)一步了解 1 比特 Adam 的底層機(jī)制

1 比特 Adam 的權(quán)重按以下公式進(jìn)行更新纫溃。對(duì)于第 i 個(gè) worker，在壓縮階段：

image.png

image.png

應(yīng)對(duì) 1 比特 Adam 的系統(tǒng)挑戰(zhàn)

除了算法上的挑戰(zhàn)外韧掩，在訓(xùn)練系統(tǒng)中應(yīng)用 1 比特 Adam 還有兩個(gè)系統(tǒng)挑戰(zhàn)紊浩。首先，我們需要具備將動(dòng)量轉(zhuǎn)換為 1 比特表示形式的功能的高效 kernel疗锐。其次坊谁，我們需要高效的通信方案來(lái)在不同的 GPU 之間傳輸壓縮后的動(dòng)量。壓縮的目的是減少總體訓(xùn)練時(shí)間窒悔，以使帶寬受限的商品系統(tǒng)可以用來(lái)訓(xùn)練大型模型呜袁。我們?cè)?DeepSpeed 中解決了這些具有挑戰(zhàn)性的問(wèn)題敌买，并針對(duì)在通信效率受限的系統(tǒng)上進(jìn)行訓(xùn)練的場(chǎng)景简珠，對(duì) 1 比特 Adam 實(shí)現(xiàn)進(jìn)行了全面的優(yōu)化。

1 比特 Adam 在通信受限系統(tǒng)上的優(yōu)勢(shì)

1 比特 Adam 提供了和 Adam 相同的收斂能力虹钮，并且最多可以減少 5 倍的通信量聋庵，用來(lái)進(jìn)行 BERT-Large 預(yù)訓(xùn)練任務(wù)時(shí)，可達(dá)最高 3.5 倍的吞吐量芙粱，用于 SQuAD fine-tuning 任務(wù)時(shí)祭玉，可達(dá) 2.7 倍的高吞吐量。端到端吞吐量的提高來(lái)源于在壓縮階段觀察到的 6.6 倍（圖 15 左）和 6.2 倍（圖 15 右）速度提升春畔。值得一提的是脱货，我們的 1 位 Adam 優(yōu)化器在 40 Gb 以太網(wǎng)系統(tǒng)上的擴(kuò)展性非常好，其性能可與 Adam 在 40 Gb InfiniBand QDR 系統(tǒng)上的擴(kuò)展性相媲美律姨。我們注意到振峻，基于 iPerf 基準(zhǔn)，40 Gb 以太網(wǎng)上的有效帶寬為 4.1 Gbps择份，而基于 InfiniBand perftest 微基準(zhǔn)扣孟，InfiniBand 提供了 32 Gbps 的近峰帶寬。

NVIDIA V100 GPU 上的 BERT-Large 預(yù)訓(xùn)練（左）和 SQuAD fine-tuning（右）的 1 比特 Adam 擴(kuò)展性荣赶。BERT 預(yù)訓(xùn)練的 batch 大小為 16/GPU凤价，SQuAD fine-tuning 為 3/GPU。

Reference
https://zhuanlan.zhihu.com/p/106783111
https://zhuanlan.zhihu.com/p/343570325

二拔创、Megatron

三利诺、horovod

四、Ray

五剩燥、paddlepaddle

六慢逾、oneflow

Data Parallelism (數(shù)據(jù)并行)

• Naive：每個(gè)worker存儲(chǔ)一份model和optimizer，每輪迭代時(shí)，將樣本分為若干份分發(fā)給各個(gè)worker氛改，實(shí)現(xiàn)并行計(jì)算
• ZeRO: Zero Redundancy Optimizer帐萎，微軟提出的數(shù)據(jù)并行內(nèi)存優(yōu)化技術(shù)，核心思想是保持Naive數(shù)據(jù)并行通信效率的同時(shí)胜卤，盡可能降低內(nèi)存占用

Model/Pipeline Parallelism (模型并行)

• Naive: 縱向切割模型疆导，將不同的layers放到不同的device上，按順序進(jìn)行正/反向傳播
• GPipe：小批量流水線方式的縱向切割模型并行
• Megatron-LM：Tensor-slicing方式的模型并行加速

Non-parallelism approach (非并行技術(shù))

• Gradient Accumulation: 通過(guò)梯度累加的方式解決顯存不足的問(wèn)題葛躏，常用于模型較大澈段，單卡只能塞下很小的batch的并行訓(xùn)練中
• CPU Offload: 同時(shí)利用 CPU 和 GPU 內(nèi)存來(lái)訓(xùn)練大型模型，即存在GPU-CPU-GPU的 transfers操作
• etc.：還有很多不一一羅列(如Checkpointing, Memory Efficient Optimizer等)