近年來(lái)，訓(xùn)練越來(lái)越大的語(yǔ)言模型已成為常態(tài)（悟道 2.0 模型參數(shù)量已經(jīng)到達(dá) 1.75T ??，為 GPT-3 的 10 倍）。但如何訓(xùn)練大型語(yǔ)言模型的信息卻很少查到 ??。
通過查找，這里整理了簡(jiǎn)單的訓(xùn)練指南??
以 BLOOM-175B 的訓(xùn)練為例
這里為 BLOOM 的訓(xùn)練使用的硬件設(shè)施，可以參考 ??
包括 fp32 優(yōu)化器狀態(tài)和 bf16+fp32 權(quán)重的 Checkpoints 為 2.3TB
只有 bf16 權(quán)重的 Checkpoints 為 329GB
176B BLOOM 模型使用 Megatron-DeepSpeed 進(jìn)行訓(xùn)練。
Megatron-DeepSpeed 結(jié)合了兩種主要技術(shù)：
DeepSpeed 是一個(gè)深度學(xué)習(xí)優(yōu)化庫(kù)，它使分布式訓(xùn)練變得簡(jiǎn)單、高效和有效。
Megatron-LM 是由 NVIDIA 的應(yīng)用深度學(xué)習(xí)研究團(tuán)隊(duì)開發(fā)的大型、強(qiáng)大的Transformer模型框架。
DeepSpeed 團(tuán)隊(duì)通過將 DeepSpeed 庫(kù)中的 ZeRO 分片（ZeRO sharding）和管道并行（pipeline parallelism）與 Megatron-LM 中的張量并行（Tensor Parallelism）相結(jié)合，開發(fā)了一種基于 3D 并行的實(shí)現(xiàn)。下文會(huì)更為詳細(xì)介紹這些技術(shù)。
Megatron-DeepSpeed 實(shí)施 3D 并行可以讓大型模型以非常有效的方式進(jìn)行訓(xùn)練。。
分布式訓(xùn)練最常見的就是 DistributedDataParallel (DDP) PyTorch 文檔。在這種方法中，模型被完全復(fù)制到每個(gè) GPU，然后在每次迭代后所有模型相互同步它們的狀態(tài)。這種方法可以加快訓(xùn)練速度，但只有當(dāng)模型可以適合單個(gè) GPU 時(shí)，它才有效。
我們以經(jīng)典的手寫數(shù)字識(shí)別為例：
數(shù)據(jù)并行通過在 N 臺(tái)機(jī)器上復(fù)制模型來(lái)實(shí)現(xiàn)。拆分 minibatch ，分成 N 個(gè)塊，讓每臺(tái)機(jī)器處理一個(gè)塊。
通過跨多個(gè)節(jié)點(diǎn)拆分，每個(gè)節(jié)點(diǎn)要做的工作更少，而且，如果忽略通信開銷，上圖的訓(xùn)練速度應(yīng)該快 2 倍。Batch 中的樣本可以獨(dú)立處理。（但值得注意 Batchnorm 等其他算子）因此，在前向傳播（計(jì)算每個(gè)樣本的輸出）和反向傳播（計(jì)算單個(gè)樣本損失權(quán)重的梯度）期間不需要通信。
為了實(shí)現(xiàn)順序一致性（生成的梯度與在單臺(tái)機(jī)器上使用順序訓(xùn)練計(jì)算出的梯度相同）。需要在更新權(quán)重之前同步梯度。最常用的損失函數(shù)是單個(gè)樣本損失的均值：
為了計(jì)算更方便，最終梯度是每個(gè)項(xiàng)的梯度的總和。因此，可以在每臺(tái)機(jī)器上獨(dú)立計(jì)算樣本的梯度，并在執(zhí)行權(quán)重更新之前將它們相加。
值得注意的是如果使用隨機(jī)梯度下降（SGD），同步權(quán)重和同步梯度是一樣的：
但是，這不適用于像 Adam 這樣的有狀態(tài)優(yōu)化器，因?yàn)楦聽顟B(tài)是梯度的非線性函數(shù)。如果使用 Adam 并同步權(quán)重而不是梯度，則每個(gè)節(jié)點(diǎn)上的優(yōu)化器狀態(tài)都會(huì)發(fā)生分歧，并且會(huì)失去順序一致性。
詳細(xì)內(nèi)容可以參考：https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/
ZeRO 具有三個(gè)主要的優(yōu)化階段，它們對(duì)應(yīng)于優(yōu)化器狀態(tài)（optimizer states）、梯度（gradients）和參數(shù)（parameters）的劃分。累積啟用時(shí)：
優(yōu)化器狀態(tài)分區(qū) (
) – 內(nèi)存減少 4 倍，通信量與數(shù)據(jù)并行性相同添加梯度分區(qū) (
) – 內(nèi)存減少 8 倍，通信量與數(shù)據(jù)并行性相同添加參數(shù)分區(qū) (
) – 內(nèi)存減少與數(shù)據(jù)并行度 成線性關(guān)系。例如，拆分為 64 個(gè) GPU () 內(nèi)存將減少到 1/64 。GPU 通信量略有增加 50%。ZeRO 消除了顯存冗余并使集群顯存容量可用。啟用所有三個(gè)階段后，ZeRO 可以僅在 1024 個(gè) NVIDIA GPU 上訓(xùn)練萬(wàn)億參數(shù)模型。例如像 Adam 這樣的 16 位精度優(yōu)化器的萬(wàn)億參數(shù)模型需要大約 16 TB 的內(nèi)存來(lái)保存優(yōu)化器狀態(tài)、梯度和參數(shù)。16TB 除以 1024 就是 16GB，這對(duì)于 GPU 來(lái)說是一個(gè)合理的范圍。
Megatron-LM 論文：https://arxiv.org/abs/2104.04473
在 Tensor Parallelism (TP) 中，每個(gè) GPU 僅處理張量的一部分，并且僅聚合完整的張量以用于需要整個(gè)事物的操作。這里使用 Megatron-LM 論文中的實(shí)現(xiàn)：GPU 集群上的高效大規(guī)模語(yǔ)言模型訓(xùn)練。任何 Transformer 的主要構(gòu)建塊都是一個(gè)完全連接的 nn.Linear，然后是一個(gè)非線性激活 GeLU。按照 Megatron 論文的符號(hào)，可以將其點(diǎn)積部分寫為
X 和 Y 是輸入和輸出向量，A 是權(quán)重矩陣。如果以矩陣形式查看計(jì)算，很容易看出矩陣乘法如何在多個(gè) GPU 之間拆分：
如果我們將權(quán)重矩陣 A 按列拆分到 N 個(gè) GPU 上并并行執(zhí)行矩陣乘法
注意 Y 矩陣沿列拆分，我們可以沿其行拆分第二個(gè) GEMM，這樣它就可以直接獲取 GeLU 的輸出，而無(wú)需任何額外的通信。
使用這個(gè)原理，可以更新任意深度的 MLP，同時(shí)在每個(gè)行列序列之后同步 GPU。Megatron-LM 論文為此提供了一個(gè)有用的說明：
其中 f 是前向傳播中的恒等運(yùn)算符，在反向傳播中是 all-reduce，而 g 是前向傳播中的 all-reduce 和反向傳播中的恒等運(yùn)算符。并行化多頭注意力層甚至更簡(jiǎn)單，因?yàn)樗鼈儽緛?lái)就是并行的
特殊考慮：由于前向和后向傳播中每層都有兩個(gè) all-reduce，因此 TP 需要在設(shè)備之間進(jìn)行非常快速的互連。因此，除非有一個(gè)非?？斓木W(wǎng)絡(luò)，否則不建議跨多個(gè)節(jié)點(diǎn)進(jìn)行 TP。
在 BLOOM 的例子中，節(jié)點(diǎn)間通信比 PCIe 慢得多。實(shí)際上，如果節(jié)點(diǎn)有 4 個(gè) GPU，則最高 TP 度因此為 4。如果需要 TP 度為 8，則需要使用至少有 8 個(gè) GPU 的節(jié)點(diǎn)。該組件由 Megatron-LM 實(shí)現(xiàn)。Megatron-LM 最近擴(kuò)展了張量并行性，以包括序列并行性，它沿著序列維度拆分不能像上面那樣拆分的操作，例如 LayerNorm。論文 Reducing Activation Recomputation in Large Transformer Models 提供了此技術(shù)的詳細(xì)信息
樸素流水線并行（Naive Pipeline Parallelism）是將一組模型層分布在多個(gè) GPU 上，并簡(jiǎn)單地將數(shù)據(jù)從 GPU 移動(dòng)到 GPU，就好像它是一個(gè)大型復(fù)合 GPU 一樣。該機(jī)制相對(duì)簡(jiǎn)單 - 切換所需的層 .to() 所需的設(shè)備，現(xiàn)在只要數(shù)據(jù)進(jìn)出這些層，就會(huì)將數(shù)據(jù)切換到與該層相同的設(shè)備，其余部分保持不變。這顯示了縱向模型并行性，會(huì)垂直切片 layers。例如，如果下圖顯示一個(gè) 8 層模型：
=================== ===================
| 0 | 1 | 2 | 3 | | 4 | 5 | 6 | 7 |
=================== ===================
GPU0 GPU1
只是將它垂直分成 2 部分，將層 0-3 放置在 GPU0 上，將層 4-7 放置在 GPU1 上。
現(xiàn)在，當(dāng)數(shù)據(jù)從第 0 層到第 1 層、第 1 層到第 2 層和第 2 層到第 3 層傳輸時(shí)，這就像單個(gè) GPU 上普通模型的前向傳遞。但是當(dāng)數(shù)據(jù)需要從第 3 層傳遞到第 4 層時(shí)，它需要從 GPU0 傳遞到 GPU1，這會(huì)引入通信開銷。如果參與的 GPU 位于同一計(jì)算節(jié)點(diǎn)（例如同一臺(tái)物理機(jī)器）上，則此復(fù)制非常快，但如果 GPU 位于不同的計(jì)算節(jié)點(diǎn)（例如多臺(tái)機(jī)器）上，通信開銷可能會(huì)大得多。然后第 4 到 5 到 6 到 7 層就像普通模型一樣，當(dāng)?shù)?7 層完成時(shí)，我們通常需要將數(shù)據(jù)發(fā)送回標(biāo)簽所在的第 0 層（或者將標(biāo)簽發(fā)送到最后一層）?，F(xiàn)在可以計(jì)算損失并且優(yōu)化器可以完成它的工作。
問題：
之所以這個(gè)被稱為樸素流水線并行，因?yàn)槠浯嬖谌毕荩?strong>是除了一個(gè) GPU 之外的所有 GPU 在任何給定時(shí)刻都是空閑的。因此，如果使用 4 個(gè) GPU，則幾乎等同于將單個(gè) GPU 的內(nèi)存量翻兩番，而忽略其余硬件。另外還有在設(shè)備之間復(fù)制數(shù)據(jù)的開銷。所以 4x 6GB 卡將能夠容納與使用樸素流水線并行的 1x 24GB 卡相同大小的模型訓(xùn)練，但后者將更快地完成訓(xùn)練，因?yàn)樗鼪]有數(shù)據(jù)復(fù)制開銷。
共享嵌入可能需要在 GPU 之間來(lái)回復(fù)制。
流水線并行 (PP) 與上述樸素流水線并行幾乎相同，但它解決了 GPU 閑置問題，方法是將傳入的 batch 為 micro-batches 并人工創(chuàng)建流水線，從而允許不同的 GPU 同時(shí)參與計(jì)算過程。GPipe 論文（https://ai.googleblog.com/2019/03/introducing-gpipe-open-source-library.html）中的下圖顯示了兩者差別：
從圖表中很容易看出第二種方式的空白區(qū)域（GPU 處于空閑狀態(tài)）更少?？瞻撞糠址Q為“氣泡”。該圖的兩個(gè)部分都顯示了 4 度的并行性。即 4 個(gè) GPU 參與管道。于是就有了 F0、F1、F2、F3 這 4 個(gè)管道的正向路徑，然后是 B3、B2、B1、B0 的返回逆序反向路徑。PP 引入了一個(gè)新的超參數(shù)來(lái)調(diào)整，稱為塊。它定義了通過同一管道階段按順序發(fā)送多少數(shù)據(jù)塊。例如，在底部圖表中，可以看到 chunks=4。GPU0 在塊 0、1、2 和 3（F0,0、F0,1、F0,2、F0,3）上執(zhí)行相同的前向路徑，然后等待其他 GPU 完成它們的工作，只有當(dāng)他們的工作開始完成時(shí)，GPU0 才開始重新工作，做 3、2、1 和 0 塊的后向路徑（B0,3, B0,2, B0,1, B0,0）
請(qǐng)注意，從概念上講，這與梯度累積步驟 (GAS) 的概念相同。PyTorch 使用 chunks，而 DeepSpeed 指的是與 GAS 相同的超參數(shù)。
因?yàn)閴K，PP 引入了 micro-batches（MBS）的概念。DP 將全局?jǐn)?shù)據(jù)批量大小拆分為小批量，因此如果 DP 度為 4，則 global data batch size 1024 將拆分為 4 個(gè) mini-batches，每個(gè) mini-batches 256 (1024/4)。如果塊（或 GAS）的數(shù)量為 32，最終得到的 micro-batches 大小為 8（256/32）。每個(gè)流水線階段一次處理一個(gè) micro-batches。
為了計(jì)算 DP + PP 設(shè)置的 global data batch size，執(zhí)行：$$mbschunksdp_{degree} (8324=1024)$$。
使用 chunks=1 你最終會(huì)得到非常低效的樸素管道并行。如果塊值非常大，您最終會(huì)得到很小的 micro-batches 大小，這也可能不是很有效。因此，必須通過實(shí)驗(yàn)來(lái)找到導(dǎo)致 GPU 最有效利用的值。雖然該圖顯示存在無(wú)法并行化的空白時(shí)間氣泡，因?yàn)樽詈笠粋€(gè)前向階段必須等待后向完成管道，但為塊找到最佳值的目的是實(shí)現(xiàn)高并發(fā)所有參與 GPU 的 GPU 利用率，這轉(zhuǎn)化為最小化氣泡的大小。這種調(diào)度機(jī)制被稱為 all forward all backward。雖然 Megatron-LM 和 DeepSpeed 都有自己的 PP 協(xié)議實(shí)現(xiàn)，但 Megatron-DeepSpeed 使用 DeepSpeed 實(shí)現(xiàn)，因?yàn)樗c DeepSpeed 的其他方面集成在一起。
在 bloom 實(shí)踐中:
這里的另一個(gè)重要問題是 word embedding 矩陣的大小。雖然通常 word embedding 矩陣比 transformer block 消耗更少的內(nèi)存，但在有 250k 詞匯表的情況下， word embedding 層需要 7.2GB 的 bf16 權(quán)重，而 transformer block 僅為 4.9GB。因此，不得不指示 Megatron-Deepspeed 將 word embedding 層視為一個(gè) transformer block。所以有一個(gè) 72 層的管道，其中 2 個(gè)專門用于 embedding（第一個(gè)和最后一個(gè)）。這允許平衡 GPU 內(nèi)存消耗。如果不這樣做，我們就會(huì)讓第一階段和最后階段消耗大部分 GPU 內(nèi)存，而 95% 的 GPU 將使用更少的內(nèi)存，因此訓(xùn)練將遠(yuǎn)非高效。
DeepSpeed 教程（https://www.deepspeed.ai/tutorials/pipeline/）中的下圖演示了如何將 DP 與 PP 結(jié)合起來(lái)。
這里重要的是要了解 DP Rank 0 如何看不到 GPU2 (和普通 DP 一樣)以及 DP Rank 1 如何看不到 GPU3。對(duì)于 DP，只有 GPU 0 和 1，它在其中提供數(shù)據(jù)，就像只有 2 個(gè) GPU 一樣。GPU0 使用 PP ”透明地“將它的一些負(fù)載卸載到 GPU2。GPU1 通過使用 GPU3 的幫助來(lái)做同樣的事情。篇幅限制就不展示了。由于每個(gè)維度至少需要 2 個(gè) GPU，因此在這里至少需要 4 個(gè) GPU。
為了獲得更高效的訓(xùn)練，PP 與 TP 和 DP 相結(jié)合，稱為 3D 并行性。這可以在下圖中看到。
可以參考：https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/
由于每個(gè)維度至少需要 2 個(gè) GPU，因此在這里至少需要 8 個(gè) GPU 才能實(shí)現(xiàn)完整的 3D 并行性。
DeepSpeed 的主要功能之一是 ZeRO，它是 DP 的超級(jí)可擴(kuò)展擴(kuò)展。ZeRO Data Parallelism 中已經(jīng)在前文討論過了。通常它是一個(gè)獨(dú)立的功能，不需要 PP 或 TP。但可以與PP、TP結(jié)合使用。當(dāng) ZeRO-DP 與 PP（和可選的 TP）結(jié)合時(shí)，它通常只啟用 ZeRO 階段 1，它只對(duì)優(yōu)化器狀態(tài)進(jìn)行分片。ZeRO 第 2 階段還對(duì)梯度進(jìn)行分片，第 3 階段也對(duì)模型權(quán)重進(jìn)行分片。雖然理論上可以將 ZeRO 第 2 階段與 Pipeline Parallelism 一起使用，但它會(huì)對(duì)性能產(chǎn)生不良影響。每個(gè) micro-batches 都需要一個(gè)額外的 reduce-scatter 集合來(lái)在分片之前聚合梯度，這會(huì)增加潛在的顯著通信開銷。根據(jù)流水線并行性的性質(zhì)，使用小的 micro-batches，重點(diǎn)是嘗試平衡計(jì)算強(qiáng)度（ micro-batches 大?。┡c最小化流水線氣泡（micro-batches 的數(shù)量）。因此，這些通信成本將會(huì)顯著增加。此外，由于 PP，層數(shù)已經(jīng)比正常情況下少，因此內(nèi)存節(jié)省不會(huì)很大。PP 已經(jīng)將梯度大小減少了 1/PP，因此在此之上的梯度分片節(jié)省不如純 DP 顯著。ZeRO 第 3 階段也可用于訓(xùn)練這種規(guī)模的模型，但是，它需要比 DeepSpeed 3D 并行實(shí)現(xiàn)更多的通信。
在 FP16 中訓(xùn)練巨大的 LLM 模型是一個(gè)禁忌（在 FP16 訓(xùn)練會(huì)導(dǎo)致數(shù)值不穩(wěn)定，或者不能產(chǎn)生足夠的精度使模型正確收斂 ??）。
BLOOM 訓(xùn)練報(bào)告中也指出了 FP16 loss 不穩(wěn)定的問題
BF16 格式的關(guān)鍵是具有與 FP32 相同的指數(shù)位，因此不會(huì)與 FP16 一樣容易溢出，使用最大數(shù)值范圍為 64k 的 FP16，只能乘以小范圍的數(shù)。例如可以做 250*250=62500，但如果你嘗試 255*255=65025，結(jié)果就會(huì)溢出，這是導(dǎo)致訓(xùn)練期間出現(xiàn)主要問題的原因。這意味著你的權(quán)重必須保持很小。一種稱為損失縮放的技術(shù)可以幫助解決這個(gè)問題，但是當(dāng)模型變得非常大時(shí)，F(xiàn)P16 的有限范圍仍然是一個(gè)問題。BF16 就沒有這個(gè)問題，可以輕松做到 10 000*10 000=100 000 000當(dāng)然，由于 BF16 和 FP16 的大小相同，均為 2 個(gè)字節(jié)，因此，當(dāng)使用 BF16 時(shí)，它的劣勢(shì)也會(huì)暴露：精度非常差。無(wú)論使用 BF16 還是 FP16，都有一個(gè)權(quán)重副本始終在 FP32 中——這是由優(yōu)化器更新的內(nèi)容。因此 16 位格式僅用于計(jì)算，優(yōu)化器以全精度 FP32 更新權(quán)重，然后將它們轉(zhuǎn)換為 16 位格式以用于下一次迭代。
一個(gè)關(guān)鍵問題是梯度累積，它是管道并行性的主要特征之一，因?yàn)槊總€(gè) micro-batches 的梯度都會(huì)累積。在 FP32 中實(shí)現(xiàn)梯度累積以保持訓(xùn)練的精確性至關(guān)重要，這就是 BF16Optimizer 所做的。
實(shí)踐中，除了其他改進(jìn)之外，BLOOM 團(tuán)隊(duì)認(rèn)為使用 BF16 混合精度訓(xùn)練將潛在的噩夢(mèng)變成了一個(gè)相對(duì)平穩(wěn)的過程，這可以從以下 lm 損失圖中觀察到：
NCCL 全稱 Nvidia Collective multi-GPU Communication Library ，是一個(gè)實(shí)現(xiàn)多 GPU 的collective communication 通信（all-gather, reduce, broadcast)庫(kù)，Nvidia 做了很多優(yōu)化，可以在 PCIe、Nvlink、InfiniBand 上實(shí)現(xiàn)較高的通信速度。
NCCL 具有以下技術(shù)特性：
下面分別從以下幾個(gè)方面來(lái)介紹 NCCL 的特點(diǎn)，包括基本的 communication primitive、ring-base collectives、NCCL 在單機(jī)多卡上以及多機(jī)多卡實(shí)現(xiàn)
并行任務(wù)的通信一般可以分為 Point-to-point communication 和 Collective communication 。P2P 通信這種模式只有一個(gè) sender 和一個(gè) receiver，實(shí)現(xiàn)起來(lái)比較簡(jiǎn)單。第二種 Collective communication 包含多個(gè) sender 多個(gè) receiver，一般的通信原語(yǔ)包括 broadcast，gather,all-gather,scatter,reduce,all-reduce,reduce-scatter,all-to-all 等。簡(jiǎn)單介紹幾個(gè)常用的操作：
ring-base collectives 將所有的通信節(jié)點(diǎn)通過首尾連接形成一個(gè)單向環(huán)，數(shù)據(jù)在環(huán)上依次傳輸。以 broadcast 為例， 假設(shè)有 4 個(gè) GPU，GPU0 為 sender 將信息發(fā)送給剩下的 GPU，按照環(huán)的方式依次傳輸，GPU0-->GPU1-->GPU2-->GPU3，若數(shù)據(jù)量為 N，帶寬為 B，整個(gè)傳輸時(shí)間為
。時(shí)間隨著節(jié)點(diǎn)數(shù)線性增長(zhǎng)，不是很高效。，條件是 S 遠(yuǎn)大于 K，即數(shù)據(jù)的份數(shù)大于節(jié)點(diǎn)數(shù)，這個(gè)很容易滿足。所以通信時(shí)間不隨節(jié)點(diǎn)數(shù)的增加而增加，只和數(shù)據(jù)總量以及帶寬有關(guān)。其它通信操作比如 reduce、gather 以此類推。那么在以GPU為通信節(jié)點(diǎn)的場(chǎng)景下，怎么構(gòu)建通信環(huán)呢？如下圖所示：?jiǎn)螜C(jī) 4 卡通過同一個(gè) PCIe switch 掛載在一棵CPU的場(chǎng)景：
NCCL 實(shí)現(xiàn)成 CUDA C++ kernels，包含 3 種 primitive operations：Copy，Reduce，ReduceAndCopy。NCCL 1.0 版本只支持單機(jī)多卡，卡之間通過 PCIe、NVlink、GPU Direct P2P來(lái)通信。NCCL 2.0 會(huì)支持多機(jī)多卡，多機(jī)間通過 Sockets (Ethernet) 或者 InfiniBand with GPU Direct RDMA 通信。下圖所示，單機(jī)內(nèi)多卡通過 PCIe 以及 CPU socket 通信，多機(jī)通過 InfiniBand 通信。
同樣，在多機(jī)多卡內(nèi)部，也要構(gòu)成一個(gè)通信環(huán)
聯(lián)系客服
