BLOOM 訓(xùn)練背后的技術(shù)
@(Engineering Practice)
假設(shè)你現(xiàn)在有了數(shù)據(jù)，也搞到了預(yù)算，一切就緒，準備開始訓(xùn)練一個大模型，一顯身手了，“一朝看盡長安花”似乎近在眼前...... 且慢！訓(xùn)練可不僅僅像這兩個字的發(fā)音那么簡單，看看 BLOOM 的訓(xùn)練或許對你有幫助。
近年來，語言模型越訓(xùn)越大已成為常態(tài)。大家通常會詬病這些大模型本身的信息未被公開以供研究，但很少關(guān)注大模型訓(xùn)練技術(shù)這種背后的知識。本文旨在以 1760 億參數(shù)的語言模型 BLOOM 為例，闡明訓(xùn)練此類模型背后的軟硬件工程和技術(shù)要點，以促進大家對大模型訓(xùn)練技術(shù)的討論。
首先，我們要感謝促成或贊助我們這個小組最終完成了訓(xùn)練 1760 億參數(shù)模型這一驚人壯舉的公司、個人和團體。
然后，我們開始討論硬件配置和主要技術(shù)組件。
以下是對本項目的簡要總結(jié)：
| 項 | 明細 |
|---|---|
| 硬件 | 384 張 80GB A100 GPU |
| 軟件 | Megatron-DeepSpeed |
| 模型架構(gòu) | 基于 GPT3 |
| 數(shù)據(jù)集 | 含 59 種語言，共 3500 億詞元 |
| 訓(xùn)練時長 | 3.5 個月 |
人員組成
該項目由 Thomas Wolf（Hugging Face 聯(lián)合創(chuàng)始人兼 CSO）發(fā)想，他敢于與大公司競爭，提出不僅要訓(xùn)練出立于世界上最大的多語言模型之林的模型，還要讓所有人都可以公開訪問訓(xùn)練結(jié)果，圓了大多數(shù)人的夢想。
本文主要關(guān)注模型訓(xùn)練的工程方面。 BLOOM 背后的技術(shù)中最重要的部分是分享專業(yè)知識并幫助我們進行編碼和訓(xùn)練的人員和公司。
我們主要需要感謝 6 個群體：
- HuggingFace 的 BigScience 團隊投入了六名以上的全職員工全程參與了訓(xùn)練的研究和運行，他們還提供或報銷了 Jean Zay 計算機之外的所有基礎(chǔ)設(shè)施。
- Microsoft DeepSpeed 團隊，開發(fā)了 DeepSpeed，后來將其與 Megatron-LM 集成，其開發(fā)人員花費數(shù)周時間研究項目需求，并在訓(xùn)練前和訓(xùn)練期間提供了許多很棒的實用經(jīng)驗建議。
- NVIDIA Megatron-LM 團隊開發(fā)了 Megatron-LM，他們非常樂于回答我們的大量問題并提供一流的使用建議。
- IDRIS / GENCI 團隊管理著 Jean Zay 超級計算機，他們?yōu)樵擁椖烤栀浟舜罅康乃懔蛷姶蟮南到y(tǒng)管理支持。
- PyTorch 團隊創(chuàng)建了一個超強的框架，其余軟件都基于該框架，并且在準備訓(xùn)練期間非常支持我們，修復(fù)了多個 bug 并提高了我們所依賴的 PyTorch 組件的訓(xùn)練可用性。
- BigScience 工程工作組志愿者
很難說出所有為該項目的工程方面做出貢獻的杰出人物的名字，所以我只列舉 Hugging Face 之外的幾個關(guān)鍵人物，他們在過去 14 個月中為該項目奠定了工程基礎(chǔ)：
Olatunji Ruwase、Deepak Narayanan、Jeff Rasley、Jared Casper、Samyam Rajbhandari 和 Rémi Lacroix
我們也感謝所有允許其員工為該項目做出貢獻的公司。
概述
BLOOM 的模型架構(gòu)與 GPT3 非常相似，只是增加了一些改進，本文稍后將對此進行討論。
該模型是在 Jean Zay 上訓(xùn)練的，Jean Zay 是由 GENCI 管理的法國政府資助的超級計算機，安裝在法國國家科學(xué)研究中心 (CNRS) 的國家計算中心 IDRIS。訓(xùn)練所需的算力由 GENCI 慷慨捐贈給本項目（捐贈號 2021-A0101012475）。
訓(xùn)練硬件：
- GPU：384 張 NVIDIA A100 80GB GPU（48 個節(jié)點）+ 32 張備用 GPU
- 每個節(jié)點 8 張 GPU，4 條 NVLink 卡間互聯(lián)，4 條 OmniPath 鏈路
- CPU：AMD EPYC 7543 32 核處理器
- CPU 內(nèi)存：每個節(jié)點 512GB
- GPU 顯存：每個節(jié)點 640GB
- 節(jié)點間連接：使用 Omni-Path Architecture (OPA) 網(wǎng)卡，網(wǎng)絡(luò)拓撲為無阻塞胖樹
- NCCL - 通信網(wǎng)絡(luò)：一個完全專用的子網(wǎng)
- 磁盤 IO 網(wǎng)絡(luò)：GPFS 與其他節(jié)點和用戶共享
Checkpoints:
- 主 checkpoints
- 每個 checkpoint 含精度為 fp32 的優(yōu)化器狀態(tài)和精度為 bf16+fp32 的權(quán)重，占用存儲空間為 2.3TB。如只保存 bf16 的權(quán)重，則僅占用 329GB 的存儲空間。
數(shù)據(jù)集:
- 1.5TB 經(jīng)過大量去重和清洗的文本，包含 46 種語言，最終轉(zhuǎn)換為 350B 個詞元
- 模型的詞匯表含 250,680 個詞元
- 更詳細信息，請參閱 The BigScience Corpus A 1.6TB Composite Multilingual Dataset
176B BLOOM 模型的訓(xùn)練于 2022 年 3 月至 7 月期間，耗時約 3.5 個月完成（約 100 萬計算時）。
Megatron-DeepSpeed
176B BLOOM 模型使用 Megatron-DeepSpeed 進行訓(xùn)練，它結(jié)合了兩種主要技術(shù)：
- DeepSpeed 是一個深度學(xué)習(xí)優(yōu)化庫，讓分布式訓(xùn)練變得簡單、高效且有效。
- Megatron-LM 是由 NVIDIA 的應(yīng)用深度學(xué)習(xí)研究團隊開發(fā)的大型、強大的 transformer 模型框架。
DeepSpeed 團隊通過將 DeepSpeed 庫中的 ZeRO 分片和流水線并行（Pipeline Parallelism）與 Megatron-LM 中的張量并行（Tensor Parallelism）相結(jié)合，開發(fā)了一種基于 3D 并行的方案。有關(guān)每個組件的更多詳細信息，請參見下表。
請注意，BigScience 的 Megatron-DeepSpeed 是基于原始 Megatron-DeepSpeed 代碼庫，我們還在其上添加了不少代碼。
下表列出了我們在訓(xùn)練 BLOOM 時各采用了兩個框架的哪些組件：
| 組件 | DeepSpeed | Megatron-LM |
|---|---|---|
| ZeRO 數(shù)據(jù)并行 | 是 | |
| 張量并行 | 是 | |
| 流水線并行 | 是 | |
| BF16 優(yōu)化器 | 是 | |
| CUDA 融合核函數(shù) | 是 | |
| 數(shù)據(jù)加載器 | 是 |
請注意，Megatron-LM 和 DeepSpeed 都有流水線并行和 BF16 優(yōu)化器實現(xiàn)，但我們使用 DeepSpeed 的實現(xiàn)，因為它們集成進了 ZeRO。
Megatron-DeepSpeed 實現(xiàn)了 3D 并行以允許大模型以非常有效的方式進行訓(xùn)練。我們簡要討論一下有哪些 3D 組件。
- 數(shù)據(jù)并行 (Data Parallelism，DP) - 相同的設(shè)置和模型被復(fù)制多份，每份每次都被饋送不同的一份數(shù)據(jù)。處理是并行完成的，所有份在每個訓(xùn)練步結(jié)束時同步。
- 張量并行 (Tensor Parallelism，TP) - 每個張量都被分成多個塊，因此張量的每個分片都位于其指定的 GPU 上，而不是讓整個張量駐留在單個 GPU 上。在處理過程中，每個分片在不同的 GPU 上分別并行處理，結(jié)果在步驟結(jié)束時同步。這就是所謂的水平并行，因為是做的水平拆分。
- 流水線并行 (Pipeline Parallelism，PP) - 模型在多個 GPU 上垂直（即按層）拆分，因此只有一個或多個模型層放置在單個 GPU 上。每個 GPU 并行處理流水線的不同階段，并處理 batch 的一部分數(shù)據(jù)。
- 零冗余優(yōu)化器 (Zero Redundancy Optimizer，ZeRO) - 也執(zhí)行與 TP 相類似的張量分片，但整個張量會及時重建以進行前向或反向計算，因此不需要修改模型。它還支持各種卸載技術(shù)以補償有限的 GPU 內(nèi)存。
數(shù)據(jù)并行
大多數(shù)只有幾張 GPU 的用戶可能比較熟悉 DistributedDataParallel(DDP)，這是相應(yīng)的 PyTorch 文檔。在該方法中，模型被完全復(fù)制到每個 GPU，然后在每次迭代后所有模型相互同步各自的狀態(tài)。這種方法可以通過投入更多 GPU 資源的方式加快訓(xùn)練速度，解決問題。但它有個限制，即只有當模型能夠放進單個 GPU 時才有效。
ZeRO 數(shù)據(jù)并行
下圖很好地描述了 ZeRO 數(shù)據(jù)并行（來自 此博文）。
看上去比較高大上，可能讓你很難專心去理解，但實際上，這個概念非常簡單。這只是通常的 DDP，只是沒有每個 GPU 都復(fù)制完整的模型參數(shù)、梯度和優(yōu)化器狀態(tài)，而是每個 GPU 只存儲其中的一部分。在隨后的運行過程中，當需要給定層的完整層參數(shù)時，所有 GPU 同步以相互提供它們?nèi)笔У牟糠?—— 僅此而已。
該組件由 DeepSpeed 實現(xiàn)。
張量并行
在張量并行 (TP) 中，每個 GPU 僅處理張量的一部分，并且僅當某些算子需要完整的張量時才觸發(fā)聚合操作。
在本節(jié)中，我們使用 Megatron-LM 論文 Efficient Large-Scale Language Model Training on GPU Clusters 中的概念和圖表。
Transformer 類模型的主要模塊為：一個全連接層 nn.Linear，后面跟一個非線性激活層 GeLU。
沿用 Megatron 論文的符號，我們可以將其點積部分寫為 Y = GeLU (XA)，其中 X 和 Y 是輸入和輸出向量，A 是權(quán)重矩陣。
如果以矩陣形式表示的話，很容易看出矩陣乘法可以如何在多個 GPU 之間拆分：
如果我們將權(quán)重矩陣 A 按列拆分到 N 個 GPU 上，然后并行執(zhí)行矩陣乘法 XA_1 到 XA_n，那么我們最終將得到 N 個輸出向量 Y_1、Y_2、…… 、 Y_n ，它們可以獨立輸入 GeLU：
注意因為 Y 矩陣是按列拆分的，因此隨后的 GEMM 我們可以選擇按行拆分方案，這樣它就可以直接獲取前面層的 GeLU 的輸出，而無需任何額外的通信。
使用該原理，我們可以更新任意深度的 MLP，只需在每個 拆列 - 拆行 序列之后同步 GPU。 Megatron-LM 論文作者為此提供了一個不錯的圖示：
這里 f 是前向傳播中的恒等運算符，后向傳播中的 all reduce，而 g 是前向傳播中的 all reduce 和后向傳播中的恒等式。
并行化多頭注意力層甚至更簡單，因為它們本來就是并行的，因為有多個獨立的頭！
需要特別考慮的是：由于前向和后向傳播中每層都有兩個 all reduce，因此 TP 需要設(shè)備間有非常快速的互聯(lián)。因此，除非你有一個非常快的網(wǎng)絡(luò)，否則不建議跨多個節(jié)點進行 TP。我們訓(xùn)練 BLOOM 的硬件配置中，節(jié)點間的速度比 PCIe 慢很多。實際上，如果節(jié)點有 4 個 GPU，則最高 TP 度設(shè)為 4 比較好。如果需要 TP 度為 8，則需要使用至少有 8 個 GPU 的節(jié)點。
該組件由 Megatron-LM 實現(xiàn)。 Megatron-LM 最近擴展了張量并行能力，新增了序列并行的能力，用于難以使用前述切分算法的算子，如 LayerNorm。Reducing Activation Recomputation in Large Transformer Models 論文提供了此技術(shù)的詳細信息。序列并行是在訓(xùn)練 BLOOM 之后開發(fā)的，所以 BLOOM 訓(xùn)練時并未采用此技術(shù)。
流水線并行
樸素流水線并行（naive PP）是將模型各層分組分布在多個 GPU 上，并簡單地將數(shù)據(jù)從 GPU 移動到 GPU，就好像它是一個大型復(fù)合 GPU 一樣。該機制相對簡單 - 將所需層用 .to () 方法綁到相應(yīng)設(shè)備，現(xiàn)在只要數(shù)據(jù)進出這些層，這些層就會將數(shù)據(jù)切換到與該層相同的設(shè)備，其余部分保持不變。
這其實就是垂直模型并行，因為如果你還記得我們是怎么畫大多數(shù)模型的拓撲圖的，我們其實是垂直切分模型各層的。例如，如果下圖顯示一個 8 層模型：
=================== ===================
| 0 | 1 | 2 | 3 | | 4 | 5 | 6 | 7 |
=================== ===================
GPU0 GPU1
我們將它垂直切成 2 部分，將層 0-3 放置在 GPU0 上，將層 4-7 放置在 GPU1 上。
現(xiàn)在，當數(shù)據(jù)從第 0 層傳到第 1 層、第 1 層傳到第 2 層以及第 2 層傳到第 3 層時，這就跟單 GPU 上的普通前向傳播一樣。但是當數(shù)據(jù)需要從第 3 層傳到第 4 層時，它需要從 GPU0 傳輸?shù)?GPU1，這會引入通信開銷。如果參與的 GPU 位于同一計算節(jié)點（例如同一臺物理機器）上，則傳輸非常快，但如果 GPU 位于不同的計算節(jié)點（例如多臺機器）上，通信開銷可能會大得多。
然后第 4 到 5 到 6 到 7 層又像普通模型一樣，當?shù)?7 層完成時，我們通常需要將數(shù)據(jù)發(fā)送回標簽所在的第 0 層（或者將標簽發(fā)送到最后一層）。現(xiàn)在可以計算損失，然后使用優(yōu)化器來進行更新參數(shù)了。
問題：
- 該方法為什么被稱為 樸素 流水線并行呢，它又有什么缺陷呢？主要是因為該方案在任意給定時刻除了一個 GPU 之外的其他所有 GPU 都是空閑的。因此，如果使用 4 個 GPU，則幾乎等同于將單個 GPU 的內(nèi)存量翻兩番，而其他資源（如計算）相當于沒用上。另外還需要加上在設(shè)備之間復(fù)制數(shù)據(jù)的開銷。所以 4 張 使用樸素流水線并行的 6GB 卡將能夠容納與 1 張 24GB 卡相同大小的模型，而后者訓(xùn)練得更快，因為它沒有數(shù)據(jù)傳輸開銷。但是，比如說，如果你有 40GB 卡，但需要跑 45GB 模型，你可以使用 4x 40GB 卡（也就剛剛夠用，因為還有梯度和優(yōu)化器狀態(tài)需要顯存）。
- 共享嵌入可能需要在 GPU 之間來回復(fù)制。
我們使用的流水線并行 (PP) 與上述樸素 PP 幾乎相同，但它解決了 GPU 閑置問題，方法是將傳入的 batch 分塊為 micros batch 并人工創(chuàng)建流水線，從而允許不同的 GPU 同時參與計算過程。
下圖來自于 GPipe 論文 ，其上半部分表示樸素 PP 方案，下半部分是 PP 方法：
從圖的下半部分很容易看出 PP 的死區(qū)（指 GPU 處于空閑狀態(tài)）更少，即 “氣泡” 更少。
圖上兩種方案的并行度均為 4 ，即由 4 張 GPU 組成流水線。于是就有了 F0、F1、F2、F3 這 4 個管級的前向路徑，然后是 B3、B2、B1、B0 的逆序后向路徑。
PP 引入了一個新的超參數(shù)來調(diào)整，稱為 塊（chunks）。它定義了通過同一管級按順序發(fā)送多少數(shù)據(jù)塊。例如，在圖的下半部分，你可以看到 chunks = 4。 GPU0 在 chunk 0、1、2 和 3（F0,0、F0,1、F0,2、F0,3）上執(zhí)行相同的前向路徑，然后等待，等其他 GPU 完成工作后，GPU0 會再次開始工作，為塊 3、2、1 和 0（B0,3、B0,2、B0,1、B0,0）執(zhí)行后向路徑。
請注意，從概念上講，這與梯度累積 (gradient accumulation steps，GAS) 的意思相同。 PyTorch 叫它 塊，而 DeepSpeed 叫它 GAS。
因為 塊，PP 引入了 micro-batches（MBS）的概念。 DP 將全局 batch size 拆分為小 batch size，因此如果 DP 度為 4，則全局 batch size 1024 將拆分為 4 個小 batch size，每個小 batch size 為 256 (1024/4)。而如果 塊（或 GAS）的數(shù)量為 32，我們最終得到的 micro batch size 為 8 (256/32)。每個管級一次處理一個 micro batch。
計算 DP + PP 設(shè)置的全局批量大小的公式為：mbs * chunks * dp_degree (8 * 32 * 4 = 1024)。
我們回過頭再看一下圖。
使用 chunks=1 你最終得到的是樸素 PP，這是非常低效的。而使用非常大的 塊 數(shù)，你最終會得到很小的微批量大小，這很可能也不是很有效。因此，必須通過實驗來找到能最有效地利用 GPU 的塊數(shù)。
該圖顯示存在無法并行化的 “死” 時間氣泡，因為最后一個 forward 階段必須等待 backward 完成流水。那么，找到最佳的 塊 數(shù)，從而使所有參與的 GPU 達到高的并發(fā)利用率，這一問題其實就轉(zhuǎn)化為最小化氣泡數(shù)了。
這種調(diào)度機制被稱為 全前全后。其他一些可選方案有 一前一后 和 交錯一前一后。
雖然 Megatron-LM 和 DeepSpeed 都有自己的 PP 協(xié)議實現(xiàn)，但 Megatron-DeepSpeed 使用的是 DeepSpeed 實現(xiàn)，因為它與 DeepSpeed 的其他功能集成在一起。
這里的另一個重要問題是詞嵌入矩陣的大小。雖然通常詞嵌入矩陣比 transfomer 塊所需的內(nèi)存更少，但在 BLOOM 有 250k 詞匯表的情況下，嵌入層需要 7.2GB 的 bf16 權(quán)重，而變換器塊僅為 4.9GB。因此，我們不得不讓 Megatron-Deepspeed 將嵌入層視為一個轉(zhuǎn)換器塊。所以我們有一個 72 級的流水線，其中 2 個是專門用于嵌入的（第一個和最后一個）。這使得我們可以平衡 GPU 的內(nèi)存消耗。如果我們不這樣做，我們就會讓第一級和最后一級消耗很大的 GPU 內(nèi)存，而 95% 的 GPU 內(nèi)存使用會很少，因此訓(xùn)練將很不高效。
DP+PP
DeepSpeed 流水線并行教程 中有一張圖演示了如何將 DP 與 PP 結(jié)合起來，如下所示。
這里重要的是要了解 DP rank 0 是看不見 GPU2 的， DP rank 1 是看不到 GPU3 的。對于 DP 而言，只有 GPU 0 和 1，并向它們饋送數(shù)據(jù)。 GPU0 使用 PP “秘密地” 將它的一些負載卸載到 GPU2。同樣地， GPU1 也會得到 GPU3 的幫助。
由于每個維度至少需要 2 個 GPU，因此這兒至少需要 4 個 GPU。
DP+PP+TP
為了更高效地訓(xùn)練，可以將 PP、TP 和 DP 相結(jié)合，稱為 3D 并行，如下圖所示。
此圖來自博文 3D 并行：擴展到萬億參數(shù)模型), 這也是一篇好文章。
由于每個維度至少需要 2 個 GPU，因此在這里你至少需要 8 個 GPU 才能實現(xiàn)完整的 3D 并行。
ZeRO DP+PP+TP
DeepSpeed 的主要功能之一是 ZeRO，它是 DP 的超級可伸縮增強版，我們在 ZeRO 數(shù)據(jù)并行 一節(jié)中已經(jīng)討論過了。通常它是一個獨立的功能，不需要 PP 或 TP。但它也可以與 PP、TP 結(jié)合使用。
當 ZeRO-DP 與 PP（以及 TP）結(jié)合時，它通常只啟用 ZeRO 階段 1，它只對優(yōu)化器狀態(tài)進行分片。 ZeRO 階段 2 還會對梯度進行分片，階段 3 也對模型權(quán)重進行分片。
雖然理論上可以將 ZeRO 階段 2 與 流水線并行 一起使用，但它會對性能產(chǎn)生不良影響。每個 micro batch 都需要一個額外的 reduce-scatter 通信來在分片之前聚合梯度，這會增加潛在的顯著通信開銷。根據(jù)流水線并行的性質(zhì)，我們會使用小的 micro batch ，并把重點放在算術(shù)強度（micro batch size）與最小化流水線氣泡（micro batch 的數(shù)量）兩者間折衷。因此，增加的通信開銷會損害流水線并行。
此外，由于 PP，層數(shù)已經(jīng)比正常情況下少，因此并不會節(jié)省很多內(nèi)存。 PP 已經(jīng)將梯度大小減少了 1/PP，因此在此基礎(chǔ)之上的梯度分片和純 DP 相比節(jié)省不了多少內(nèi)存。
ZeRO 階段 3 也可用于訓(xùn)練這種規(guī)模的模型，但是，它需要的通信量比 DeepSpeed 3D 并行更多。一年前，在對我們的環(huán)境進行仔細評估后，我們發(fā)現(xiàn) Megatron-DeepSpeed 3D 并行性表現(xiàn)最佳。此后，ZeRO 階段 3 的性能有了顯著提高，如果我們今天要對其進行重新評估，也許我們會選擇階段 3。
BF16 優(yōu)化器
用 FP16 訓(xùn)練巨型 LLM 模型是一個禁忌。
我們已經(jīng)通過花費幾個月的時間 訓(xùn)練 104B 模型 自證了這一點，你可以從 tensorboard 發(fā)現(xiàn)，徹頭徹尾地失敗了。在與不斷發(fā)散的 lm-loss 作斗爭的過程中，我們學(xué)到了很多：
我們也從 Megatron-LM 和 DeepSpeed 團隊那里得到了相同的建議，在他們訓(xùn)得 530B 模型 后。最近發(fā)布的 OPT-175B 也報告說他們在 FP16 上訓(xùn)練得非常艱難。
所以早在一月份，我們就知道我們要在支持 BF16 格式的 A100 上進行訓(xùn)練。 Olatunji Ruwase 開發(fā)了一個用來訓(xùn)練 BLOOM 的 “BF16Optimizer”。
如果您不熟悉這種數(shù)據(jù)格式，請查看 它的位布局。 BF16 格式的關(guān)鍵是它的指數(shù)位數(shù)與 FP32 相同，因此不會溢出，但 FP16 經(jīng)常溢出！FP16 的最大數(shù)值范圍為 64k，您只能進行較小數(shù)的乘法。例如你可以做 250*250=62500，但如果你嘗試 255*255=65025，你就會溢出，這是導(dǎo)致訓(xùn)練出現(xiàn)問題的主要原因。這意味著你的權(quán)重必須保持很小。一種稱為損失縮放（loss scaling）的技術(shù)有助于緩解這個問題，但是當模型變得非常大時，F(xiàn)P16 較小的數(shù)值范圍仍然是一個問題。
BF16 沒有這個問題，你可以很容易地做 10_000*10_000=100_000_000, 完全沒問題。
當然，由于 BF16 和 FP16 的大小相同，均為 2 個字節(jié)，因此，沒有免費的午餐，當使用 BF16 時，代價就是它的精度非常差。然而，你應(yīng)該還記得我們在訓(xùn)練時采用的隨機梯度下降法及其變體，該方法有點像蹣跚而行，如果你這步?jīng)]有找到完美的方向其實沒關(guān)系，你會在接下來的步驟中糾正自己。
無論使用 BF16 還是 FP16，都有一個權(quán)重副本始終在 FP32 中 —— 這是由優(yōu)化器更新的內(nèi)容。因此 16 位格式僅用于計算，優(yōu)化器以全精度更新 FP32 權(quán)重，然后將它們轉(zhuǎn)換為 16 位格式以用于下一次迭代。
所有 PyTorch 組件都已更新，以確保它們在 FP32 中執(zhí)行任何累加，因此不會發(fā)生精度損失。
一個關(guān)鍵問題是梯度累積，它是流水線并行的主要特征之一，因為每個 micro batch 處理的梯度都會累積。在 FP32 中實現(xiàn)梯度累積以保證訓(xùn)練的精確性至關(guān)重要，這正是 BF16Optimizer 所做的。
除了其他改進之外，我們認為使用 BF16 混合精度訓(xùn)練將潛在的噩夢變成了一個相對平穩(wěn)的過程，這可以從以下 lm 損失圖中看出：
CUDA 融合核函數(shù)
GPU 主要做兩件事。它可以將數(shù)據(jù)寫到顯存或從顯存讀數(shù)據(jù)，并對這些數(shù)據(jù)執(zhí)行計算。當 GPU 忙于讀寫數(shù)據(jù)時， GPU 的計算單元就會空閑。如果我們想有效地利用 GPU，我們希望將空閑時間降至最低。
核函數(shù)是一組實現(xiàn)特定 PyTorch 操作的指令。例如，當你調(diào)用 torch.add 時，它會通過一個 PyTorch 調(diào)度器，它會根據(jù)輸入張量及其他變量的取值來決定它應(yīng)該運行哪些代碼，最后運行它。 CUDA 核函數(shù)使用 CUDA 來實現(xiàn)這些代碼，因此只能在 NVIDIA GPU 上運行。
現(xiàn)在，當使用 GPU 計算 c = torch.add (a, b); e = torch.max ([c,d]) 時，一般情況下，PyTorch 將執(zhí)行的操作是啟動兩個單獨的核函數(shù)，一個執(zhí)行 a 和 b 的加法，另一個執(zhí)行取 c 和 d 兩者的最大值。在這種情況下，GPU 從其顯存中獲取 a 和 b，執(zhí)行加法運算，然后將結(jié)果寫回顯存。然后它獲取 c 和 d 并執(zhí)行 max 操作，然后再次將結(jié)果寫回顯存。
如果我們要融合這兩個操作，即將它們放入一個 “融合核函數(shù)” 中，然后啟動那個內(nèi)核，我們不會將中間結(jié)果 c 寫到顯存中，而是將其保留在 GPU 寄存器中，并且僅需要獲取 d 來完成最后的計算。這節(jié)省了大量開銷并防止 GPU 空閑，因此整個操作會更加高效。
融合核函數(shù)就是這樣。它們主要將多個離散的計算和進出顯存的數(shù)據(jù)移動替換為有很少數(shù)據(jù)移動的融合計算。此外，一些融合核函數(shù)會對操作進行數(shù)學(xué)變換，以便可以更快地執(zhí)行某些計算組合。
為了快速高效地訓(xùn)練 BLOOM，有必要使用 Megatron-LM 提供的幾個自定義 CUDA 融合核函數(shù)。特別地，有一個 LayerNorm 的融合核函數(shù)以及用于融合縮放、掩碼和 softmax 這些操作的各種組合的核函數(shù)。Bias Add 也通過 PyTorch 的 JIT 功能與 GeLU 融合。這些操作都是瓶頸在內(nèi)存的，因此將它們?nèi)诤显谝黄鹨赃_到最大化每次顯存讀取后的計算量非常重要。因此，例如，在執(zhí)行瓶頸在內(nèi)存的 GeLU 操作時同時執(zhí)行 Bias Add，運行時間并不會增加。這些核函數(shù)都可以在 Megatron-LM 代碼庫 中找到。
數(shù)據(jù)集
Megatron-LM 的另一個重要特性是高效的數(shù)據(jù)加載器。在首次訓(xùn)練啟動前，每個數(shù)據(jù)集中的每個樣本都被分成固定序列長度（BLOOM 為 2048）的樣本，并創(chuàng)建索引以對每個樣本進行編號。基于訓(xùn)練超參，我們會確定每個數(shù)據(jù)集所需要參與的 epoch 數(shù)，并基于此創(chuàng)建一個有序的樣本索引列表，然后打亂它。舉個例子，如果一個數(shù)據(jù)集中有 10 個樣本并應(yīng)參與 2 個 epoch 的訓(xùn)練，則系統(tǒng)首先按 [0, ..., 9, 0, ..., 9] 順序排好樣本索引，然后打亂該順序為數(shù)據(jù)集創(chuàng)建最終的全局順序。請注意，這意味著訓(xùn)練不會簡單地遍歷整個數(shù)據(jù)集然后重復(fù)，你有可能在看到另一個樣本之前看到同一個樣本兩次，但在訓(xùn)練結(jié)束時模型將只看到每個樣本兩次。這有助于確保整個訓(xùn)練過程中的訓(xùn)練曲線平滑。這些索引，包括每個樣本在原始數(shù)據(jù)集中的偏移量，被保存到一個文件中，以避免每次開始訓(xùn)練時都重新計算它們。最后，可以將其中幾個數(shù)據(jù)集以不同的權(quán)重混合到訓(xùn)練最終使用的數(shù)據(jù)中。
嵌入 LayerNorm
在我們努力阻止 104B 模型發(fā)散的過程中，我們發(fā)現(xiàn)在第一個層詞嵌入層之后添加一個額外的 LayerNorm 可以使訓(xùn)練更加穩(wěn)定。
該洞察來自對 bitsandbytes 的實驗，bitsandbytes 有一個 StableEmbedding 操作，它是一個帶有 LayerNorm 的普通嵌入，其使用均勻 xavier 函數(shù)來初始化。
位置編碼
基于論文 Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation，我們還用 AliBi 替換了普通的位置嵌入，它允許外推比訓(xùn)練模型的輸入序列更長的輸入序列。因此，即使我們訓(xùn)練時使用長度為 2048 的序列，模型也可以在推理過程中處理更長的序列。
訓(xùn)練中的困難
隨著架構(gòu)、硬件和軟件的就位，我們得以在 2022 年 3 月上旬開始訓(xùn)練。然而，從那時起，事情其實并非一帆風(fēng)順。在本節(jié)中，我們將討論我們遇到的一些主要障礙。
在訓(xùn)練開始之前，有很多問題需要弄清楚。特別是，我們發(fā)現(xiàn)了幾個問題，這些問題只有在我們開始在 48 個節(jié)點上進行訓(xùn)練后才會出現(xiàn)，而不會在小規(guī)模時出現(xiàn)。例如，需要設(shè) CUDA_LAUNCH_BLOCKING=1 來防止框架掛起，我們需要將優(yōu)化器組分成更小的組，否則框架會再次掛起。你可以在 訓(xùn)前編年史 中詳細了解這些內(nèi)容。
訓(xùn)練期間遇到的主要問題類型是硬件故障。由于這是一個擁有大約 400 個 GPU 的新集群，平均每周我們會遇到 1-2 個 GPU 故障。我們每 3 小時（100 次迭代）保存一個檢查點。因此，我們每周因硬件崩潰平均損失 1.5 小時的訓(xùn)練成果。 Jean Zay 系統(tǒng)管理員隨后將更換有故障的 GPU 并恢復(fù)節(jié)點。與此同時，我們有備用節(jié)點可供使用。
我們還遇到過多次導(dǎo)致 5-10 小時停機的各種其他問題，其中一些與 PyTorch 中的死鎖錯誤有關(guān)，另一些則是由于磁盤空間不足。如果您對具體細節(jié)有興趣，請參閱 訓(xùn)練編年史。
在對訓(xùn)練這個模型進行可行性分析時，所有這些停機時間都被計劃在內(nèi)了，我們也據(jù)此選擇了合適的模型大小和我們希望模型消耗的數(shù)據(jù)量。因此，即使存在這些停機問題，我們還是成功地在預(yù)計時間內(nèi)完成了訓(xùn)練。如前所述，它需要大約 100 萬個計算時才能完成。
另一個問題是 SLURM 并非設(shè)計為供一組人使用。 SLURM 作業(yè)由單個用戶擁有，如果他們不在身邊，則該組的其他成員無法對正在運行的作業(yè)執(zhí)行任何操作。我們制定了一個終止方案，允許組中的其他用戶終止當前進程，而不需要啟動該進程的用戶在場。這在 90% 的問題上都很有效。如果 SLURM 設(shè)計者讀到這篇文章，請?zhí)砑右粋€ Unix 組的概念，這樣一個 SLURM 作業(yè)就可以由一個組擁有。
由于訓(xùn)練是全天候 24/7 進行的，我們需要有人隨叫隨到 - 但由于我們在歐洲和加拿大西海岸都有人，因此不需要有人攜帶傳呼機，我們能很好地互相備份。當然，周末的訓(xùn)練也得有人看著。我們自動化了大部分事情，包括自動從硬件崩潰中恢復(fù)，但有時仍需要人工干預(yù)。
結(jié)論
訓(xùn)練中最困難和最緊張的部分是訓(xùn)練開始前的 2 個月。我們承受著盡快開始訓(xùn)練的巨大壓力，因為資源分配的時間有限，我們直到最后一刻才接觸到 A100。所以這是一個非常困難的時期，考慮到 BF16Optimizer 是在最后一刻編寫出來的，我們需要調(diào)試它并修復(fù)各種 bug。正如上一節(jié)所述，我們發(fā)現(xiàn)了新問題，這些問題只有在我們開始在 48 個節(jié)點上進行訓(xùn)練后才會出現(xiàn)，并且不會在小規(guī)模時出現(xiàn)。
但是一旦我們把這些整理完，訓(xùn)練本身出奇的順利，沒有出現(xiàn)大的問題。大多數(shù)時候，我們只有一個人看著，只有少數(shù)幾個人參與故障排除。我們得到了 Jean Zay 管理部門的大力支持，他們迅速解決了訓(xùn)練期間出現(xiàn)的大部分需求。
總的來說，這是一次超級緊張但回報頗豐的經(jīng)歷。
訓(xùn)練大型語言模型仍然是一項具有挑戰(zhàn)性的任務(wù)，但我們希望通過公開構(gòu)建和共享這項技術(shù)，其他人可以借鑒我們的經(jīng)驗。
資源
重要鏈接
論文與文章
我們不可能在本文中詳細解釋所有內(nèi)容，因此如果此處介紹的技術(shù)激起你的好奇心，使你想了解更多信息，請閱讀以下論文：
Megatron-LM:
- Efficient Large-Scale Language Model Training on GPU Clusters.
- Reducing Activation Recomputation in Large Transformer Models
DeepSpeed:
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
- ZeRO-Offload: Democratizing Billion-Scale Model Training
- ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
- DeepSpeed: Extreme-scale model training for everyone
Megatron-LM 和 Deepspeeed 聯(lián)合:
ALiBi:
- Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation
- What Language Model to Train if You Have One Million GPU Hours? - 你會在那里找到最終使得我們選擇 ALiBi 的實驗。
BitsNBytes:
- 8-bit Optimizers via Block-wise Quantization （我們使用了該論文中的嵌入 LaynerNorm，但是論文的其他部分及其技術(shù)也很妙，我們沒用 8 位優(yōu)化器的唯一原因是我們已經(jīng)使用 DeepSpeed-ZeRO 節(jié)省了優(yōu)化器內(nèi)存）。
博文致謝
非常感謝以下這些人，他們提出了很好的問題并幫助提高了文章的可讀性（按字母序）：
Britney Muller,
Douwe Kiela,
Jared Casper,
Jeff Rasley,
Julien Launay,
Leandro von Werra,
Omar Sanseviero,
Stefan Schweter and
Thomas Wang.
本文圖表主要由 Chunte Lee 創(chuàng)作。
英文原文: <url> https://huggingface.co/blog/bloom-megatron-deepspeed </url>
原文作者：Stas Bekman
譯者: Matrix Yao (姚偉峰)，英特爾深度學(xué)習(xí)工程師，工作方向為 transformer-family 模型在各模態(tài)數(shù)據(jù)上的應(yīng)用及大規(guī)模模型的訓(xùn)練推理。
