千億參數開源大模型 BLOOM 背後的技術

假設你現在有了數據，也搞到了預算，一切就緒，準備開始訓練一個大模型，一顯身手了，“一朝看盡長安花”似乎近在眼前 …… 且慢！訓練可不僅僅像這兩個字的發音那麼簡單，看看 BLOOM 的訓練或許對你有幫助。

近年來，語言模型越訓越大已成為常態。大家通常會詬病這些大模型本身的信息未被公開以供研究，但很少關注大模型訓練技術這種背後的知識。本文旨在以 1760 億參數的語言模型 BLOOM 為例，闡明訓練此類模型背後的軟硬件工程和技術要點，以促進大家對大模型訓練技術的討論。

首先，我們要感謝促成或贊助我們這個小組最終完成了訓練 1760 億參數模型這一驚人壯舉的公司、個人和團體。

然後，我們開始討論硬件配置和主要技術組件。

以下是對本項目的簡要總結:

硬件	384 張 80GB A100 GPU
軟件	Megatron-DeepSpeed
模型架構	基於 GPT3
數據集	含 59 種語言，共 3500 億詞元
訓練時長	3.5 個月

人員組成

該項目由 Thomas Wolf (Hugging Face 聯合創始人兼 CSO) 發想，他敢於與大公司競爭，提出不僅要訓練出立於世界上最大的多語言模型之林的模型，還要讓所有人都可以公開訪問訓練結果，圓了大多數人的夢想。

本文主要關注模型訓練的工程方面。 BLOOM 背後的技術中最重要的部分是分享專業知識並幫助我們進行編碼和訓練的人員和公司。

我們主要需要感謝 6 個羣體:

1. HuggingFace 的 BigScience 團隊投入了六名以上的全職員工全程參與了訓練的研究和運行，他們還提供或報銷了 Jean Zay 計算機之外的所有基礎設施。
2. Microsoft DeepSpeed 團隊，開發了 DeepSpeed，後來將其與 Megatron-LM 集成，其開發人員花費數週時間研究項目需求，並在訓練前和訓練期間提供了許多很棒的實用經驗建議。
3. NVIDIA Megatron-LM 團隊開發了 Megatron-LM，他們非常樂於回答我們的大量問題並提供一流的使用建議。
4. IDRIS / GENCI 團隊管理着 Jean Zay 超級計算機，他們為該項目捐贈了大量的算力和強大的系統管理支持。
5. PyTorch 團隊創建了一個超強的框架，其餘軟件都基於該框架，並且在準備訓練期間非常支持我們，修復了多個 bug 並提高了我們所依賴的 PyTorch 組件的訓練可用性。
6. BigScience 工程工作組志願者

很難説出所有為該項目的工程方面做出貢獻的傑出人物的名字，所以我只列舉 Hugging Face 之外的幾個關鍵人物，他們在過去 14 個月中為該項目奠定了工程基礎:

Olatunji Ruwase、Deepak Narayanan、Jeff Rasley、Jared Casper、Samyam Rajbhandari 和 Rémi Lacroix

我們也感謝所有允許其員工為該項目做出貢獻的公司。

概述

BLOOM 的模型架構與 GPT3 非常相似，只是增加了一些改進，本文稍後將對此進行討論。

該模型是在 Jean Zay 上訓練的，Jean Zay 是由 GENCI 管理的法國政府資助的超級計算機，安裝在法國國家科學研究中心 (CNRS) 的國家計算中心 IDRIS。訓練所需的算力由 GENCI 慷慨捐贈給本項目 (捐贈號 2021-A0101012475)。

訓練硬件:

• GPU: 384 張 NVIDIA A100 80GB GPU (48 個節點) + 32 張備用 GPU
• 每個節點 8 張 GPU，4 條 NVLink 卡間互聯，4 條 OmniPath 鏈路
• CPU: AMD EPYC 7543 32 核處理器
• CPU 內存: 每個節點 512GB
• GPU 顯存: 每個節點 640GB
• 節點間連接: 使用 Omni-Path Architecture (OPA) 網卡，網絡拓撲為無阻塞胖樹
• NCCL - 通信網絡: 一個完全專用的子網
• 磁盤 IO 網絡: GPFS 與其他節點和用户共享

Checkpoints:

• 主 checkpoints
• 每個 checkpoint 含精度為 fp32 的優化器狀態和精度為 bf16+fp32 的權重，佔用存儲空間為 2.3TB。如只保存 bf16 的權重，則僅佔用 329GB 的存儲空間。

數據集:

• 1.5TB 經過大量去重和清洗的文本，包含 46 種語言，最終轉換為 350B 個詞元
• 模型的詞彙表含 250,680 個詞元
• 更詳細信息，請參閲 The BigScience Corpus A 1.6TB Composite Multilingual Dataset

176B BLOOM 模型的訓練於 2022 年 3 月至 7 月期間，耗時約 3.5 個月完成 (約 100 萬計算時)。

Megatron-DeepSpeed

176B BLOOM 模型使用 Megatron-DeepSpeed 進行訓練，它結合了兩種主要技術:

• Megatron-DeepSpeed:
• DeepSpeed 是一個深度學習優化庫，讓分佈式訓練變得簡單、高效且有效。
• Megatron-LM 是由 NVIDIA 的應用深度學習研究團隊開發的大型、強大的 transformer 模型框架。

DeepSpeed 團隊通過將 DeepSpeed 庫中的 ZeRO 分片和流水線並行 (Pipeline Parallelism) 與 Megatron-LM 中的張量並行 (Tensor Parallelism) 相結合，開發了一種基於 3D 並行的方案。有關每個組件的更多詳細信息，請參見下表。

請注意，BigScience 的 Megatron-DeepSpeed 是基於原始 Megatron-DeepSpeed 代碼庫，我們還在其上添加了不少代碼。

下表列出了我們在訓練 BLOOM 時各採用了兩個框架的哪些組件:

組件	DeepSpeed	Megatron-LM
ZeRO 數據並行	是
張量並行		是
流水線並行	是
BF16 優化器	是
CUDA 融合核函數		是
數據加載器		是

請注意，Megatron-LM 和 DeepSpeed 都有流水線並行和 BF16 優化器實現，但我們使用 DeepSpeed 的實現，因為它們集成進了 ZeRO。

Megatron-DeepSpeed 實現了 3D 並行以允許大模型以非常有效的方式進行訓練。我們簡要討論一下有哪些 3D 組件。

1. 數據並行 (Data Parallelism，DP) - 相同的設置和模型被複制多份，每份每次都被饋送不同的一份數據。處理是並行完成的，所有份在每個訓練步結束時同步。
2. 張量並行 (Tensor Parallelism，TP) - 每個張量都被分成多個塊，因此張量的每個分片都位於其指定的 GPU 上，而不是讓整個張量駐留在單個 GPU 上。在處理過程中，每個分片在不同的 GPU 上分別並行處理，結果在步驟結束時同步。這就是所謂的水平並行，因為是做的水平拆分。
3. 流水線並行 (Pipeline Parallelism，PP) - 模型在多個 GPU 上垂直 (即按層) 拆分，因此只有一個或多個模型層放置在單個 GPU 上。每個 GPU 並行處理流水線的不同階段，並處理 batch 的一部分數據。
4. 零宂餘優化器 (Zero Redundancy Optimizer，ZeRO) - 也執行與 TP 相類似的張量分片，但整個張量會及時重建以進行前向或反向計算，因此不需要修改模型。它還支持各種卸載技術以補償有限的 GPU 內存。

數據並行

大多數只有幾張 GPU 的用户可能比較熟悉 DistributedDataParallel(DDP)，這是相應的 PyTorch 文檔。在該方法中，模型被完全複製到每個 GPU，然後在每次迭代後所有模型相互同步各自的狀態。這種方法可以通過投入更多 GPU 資源的方式加快訓練速度，解決問題。但它有個限制，即只有當模型能夠放進單個 GPU 時才有效。

ZeRO 數據並行

下圖很好地描述了 ZeRO 數據並行 (來自 此博文)。

看上去比較高大上，可能讓你很難專心去理解，但實際上，這個概念非常簡單。這只是通常的 DDP，只是沒有每個 GPU 都複製完整的模型參數、梯度和優化器狀態，​​而是每個 GPU 只存儲其中的一部分。在隨後的運行過程中，當需要給定層的完整層參數時，所有 GPU 同步以相互提供它們缺失的部分 —— 僅此而已。

該組件由 DeepSpeed 實現。

張量並行

在張量並行 (TP) 中，每個 GPU 僅處理張量的一部分，並且僅當某些算子需要完整的張量時才觸發聚合操作。

在本節中，我們使用 Megatron-LM 論文 Efficient Large-Scale Language Model Training on GPU Clusters 中的概念和圖表。

Transformer 類模型的主要模塊為: 一個全連接層 nn.Linear，後面跟一個非線性激活層 GeLU。

沿用 Megatron 論文的符號，我們可以將其點積部分寫為 Y = GeLU (XA)，其中 X 和 Y 是輸入和輸出向量， A 是權重矩陣。

如果以矩陣形式表示的話，很容易看出矩陣乘法可以如何在多個 GPU 之間拆分:

如果我們將權重矩陣 A 按列拆分到 N 個 GPU 上，然後並行執行矩陣乘法 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 需要設備間有非常快速的互聯。因此，除非你有一個非常快的網絡，否則不建議跨多個節點進行 TP。我們訓練 BLOOM 的硬件配置中，節點間的速度比 PCIe 慢很多。實際上，如果節點有 4 個 GPU，則最高 TP 度設為 4 比較好。如果需要 TP 度為 8，則需要使用至少有 8 個 GPU 的節點。

該組件由 Megatron-LM 實現。 Megatron-LM 最近擴展了張量並行能力，新增了序列並行的能力，用於難以使用前述切分算法的算子，如 LayerNorm。Reducing Activation Recomputation in Large Transformer Models 論文提供了此技術的詳細信息。序列並行是在訓練 BLOOM 之後開發的，所以 BLOOM 訓練時並未採用此技術。

流水線並行

樸素流水線並行 (naive PP) 是將模型各層分組分佈在多個 GPU 上，並簡單地將數據從 GPU 移動到 GPU，就好像它是一個大型複合 GPU 一樣。該機制相對簡單 - 將所需層用 .to() 方法綁到相應設備，現在只要數據進出這些層，這些層就會將數據切換到與該層相同的設備，其餘部分保持不變。

這其實就是垂直模型並行，因為如果你還記得我們是怎麼畫大多數模型的拓撲圖的，我們其實是垂直切分模型各層的。例如，如果下圖顯示一個 8 層模型:

=================== ===================
| 0 | 1 | 2 | 3  | |  4 | 5 | 6 | 7 |
=================== ===================
        GPU0 GPU1

我們將它垂直切成 2 部分，將層 0-3 放置在 GPU0 上，將層 4-7 放置在 GPU1 上。

現在，當數據從第 0 層傳到第 1 層、第 1 層傳到第 2 層以及第 2 層傳到第 3 層時，這就跟單 GPU 上的普通前向傳播一樣。但是當數據需要從第 3 層傳到第 4 層時，它需要從 GPU0 傳輸到 GPU1，這會引入通信開銷。如果參與的 GPU 位於同一計算節點 (例如同一台物理機器) 上，則傳輸非常快，但如果 GPU 位於不同的計算節點 (例如多台機器) 上，通信開銷可能會大得多。

然後第 4 到 5 到 6 到 7 層又像普通模型一樣，當第 7 層完成時，我們通常需要將數據發送回標籤所在的第 0 層 (或者將標籤發送到最後一層)。現在可以計算損失，然後使用優化器來進行更新參數了。

問題:

• 該方法為什麼被稱為 樸素 流水線並行呢，它又有什麼缺陷呢？主要是因為該方案在任意給定時刻除了一個 GPU 之外的其他所有 GPU 都是空閒的。因此，如果使用 4 個 GPU，則幾乎等同於將單個 GPU 的內存量翻兩番，而其他資源 (如計算) 相當於沒用上。另外還需要加上在設備之間複製數據的開銷。所以 4 張 使用樸素流水線並行的 6GB 卡將能夠容納與 1 張 24GB 卡相同大小的模型，而後者訓練得更快，因為它沒有數據傳輸開銷。但是，比如説，如果你有 40GB 卡，但需要跑 45GB 模型，你可以使用 4x 40GB 卡 (也就剛剛夠用，因為還有梯度和優化器狀態需要顯存)。
• 共享嵌入可能需要在 GPU 之間來回複製。 我們使用的流水線並行 (PP) 與上述樸素 PP 幾乎相同，但它解決了 GPU 閒置問題，方法是將傳入的 batch 分塊為 micros batch 並人工創建流水線，從而允許不同的 GPU 同時參與計算過程。

下圖來自於 GPipe 論文，其上半部分表示樸素 PP 方案，下半部分是 PP 方法:

從圖的下半部分很容易看出 PP 的死區 (指 GPU 處於空閒狀態) 更少，即 “氣泡” 更少。

圖上兩種方案的並行度均為 4 ，即由 4 張 GPU 組成流水線。於是就有了 F0、F1、F2、F3 這 4 個管級的前向路徑，然後是 B3、B2、B1、B0 的逆序後向路徑。

PP 引入了一個新的超參數來調整，稱為 塊 (chunks)。它定義了通過同一管級按順序發送多少數據塊。例如，在圖的下半部分，你可以看到 chunks = 4。 GPU0 在 chunk 0、1、2 和 3 (F0,0、F0,1、F0,2、F0,3) 上執行相同的前向路徑，然後等待，等其他 GPU 完成工作後，GPU0 會再次開始工作，為塊 3、2、1 和 0 (B0,3、B0,2、B0,1、B0,0) 執行後向路徑。

請注意，從概念上講，這與梯度累積 (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) 的數量為 32，我們最終得到的 micro batch size 為 8 (256/32)。每個管級一次處理一個 micro batch。

計算 DP + PP 設置的全局批量大小的公式為: mbs * chunks * dp_degree (8 * 32 * 4 = 1024)。

我們回過頭再看一下圖。

使用 chunks=1 你最終得到的是樸素 PP，這是非常低效的。而使用非常大的 塊 數，你最終會得到很小的微批量大小，這很可能也不是很有效。因此，必須通過實驗來找到能最有效地利用 GPU 的 塊數。

該圖顯示存在無法並行化的 “死” 時間氣泡，因為最後一個 forward 階段必須等待 backward 完成流水。那麼，找到最佳的 塊 數，從而使所有參與的 GPU 達到高的併發利用率，這一問題其實就轉化為最小化氣泡數了。

這種調度機制被稱為 全前全後。其他一些可選方案有 一前一後 和 交錯一前一後。

雖然 Megatron-LM 和 DeepSpeed 都有自己的 PP 協議實現，但 Megatron-DeepSpeed 使用的是 DeepSpeed 實現，因為它與 DeepSpeed 的其他功能集成在一起。

這裏的另一個重要問題是詞嵌入矩陣的大小。雖然通常詞嵌入矩陣比 transfomer 塊所需的內存更少，但在 BLOOM 有 250k 詞彙表的情況下，嵌入層需要 7.2GB 的 bf16 權重，而變換器塊僅為 4.9GB。因此，我們不得不讓 Megatron-Deepspeed 將嵌入層視為一個轉換器塊。所以我們有一個 72 級的流水線，其中 2 個是專門用於嵌入的 (第一個和最後一個)。這使得我們可以平衡 GPU 的內存消耗。如果我們不這樣做，我們就會讓第一級和最後一級消耗很大的 GPU 內存，而 95% 的 GPU 內存使用會很少，因此訓練將很不高效。

DP+PP

DeepSpeed 流水線並行教程 中有一張圖演示瞭如何將 DP 與 PP 結合起來，如下所示。

這裏重要的是要了解 DP rank 0 是看不見 GPU2 的， DP rank 1 是看不到 GPU3 的。對於 DP 而言，只有 GPU 0 和 1，並向它們饋送數據。 GPU0 使用 PP “祕密地” 將它的一些負載卸載到 GPU2。同樣地， GPU1 也會得到 GPU3 的幫助。

由於每個維度至少需要 2 個 GPU，因此這兒至少需要 4 個 GPU。

DP+PP+TP

為了更高效地訓練，可以將 PP、TP 和 DP 相結合，稱為 3D 並行，如下圖所示。

此圖來自博文 3D 並行: 擴展到萬億參數模型), 這也是一篇好文章。

由於每個維度至少需要 2 個 GPU，因此在這裏你至少需要 8 個 GPU 才能實現完整的 3D 並行。

ZeRO DP+PP+TP

DeepSpeed 的主要功能之一是 ZeRO，它是 DP 的超級可伸縮增強版，我們在 [ZeRO 數據並行](#ZeRO- 數據並行) 一節中已經討論過了。通常它是一個獨立的功能，不需要 PP 或 TP。但它也可以與 PP、TP 結合使用。

當 ZeRO-DP 與 PP (以及 TP) 結合時，它通常只啟用 ZeRO 階段 1，它只對優化器狀態進行分片。 ZeRO 階段 2 還會對梯度進行分片，階段 3 也對模型權重進行分片。

雖然理論上可以將 ZeRO 階段 2 與 流水線並行 一起使用，但它會對性能產生不良影響。每個 micro batch 都需要一個額外的 reduce-scatter 通信來在分片之前聚合梯度，這會增加潛在的顯著通信開銷。根據流水線並行的性質，我們會使用小的 micro batch ，並把重點放在算術強度 (micro batch size) 與最小化流水線氣泡 (micro batch 的數量) 兩者間折衷。因此，增加的通信開銷會損害流水線並行。

此外，由於 PP，層數已經比正常情況下少，因此並不會節省很多內存。 PP 已經將梯度大小減少了 1/PP，因此在此基礎之上的梯度分片和純 DP 相比節省不了多少內存。

ZeRO 階段 3 也可用於訓練這種規模的模型，但是，它需要的通信量比 DeepSpeed 3D 並行更多。一年前，在對我們的環境進行仔細評估後，我們發現 Megatron-DeepSpeed 3D 並行性表現最佳。此後，ZeRO 階段 3 的性能有了顯著提高，如果我們今天要對其進行重新評估，也許我們會選擇階段 3。

BF16 優化器

用 FP16 訓練巨型 LLM 模型是一個禁忌。

我們已經通過花費幾個月的時間 訓練 104B 模型 自證了這一點，你可以從 Tensorboard 發現，徹頭徹尾地失敗了。在與不斷髮散的 lm-loss 作鬥爭的過程中，我們學到了很多:

我們也從 Megatron-LM 和 DeepSpeed 團隊那裏得到了相同的建議，在他們訓得 530B 模型 後。最近發佈的 OPT-175B 也報告説他們在 FP16 上訓練得非常艱難。

所以早在一月份，我們就知道我們要在支持 BF16 格式的 A100 上進行訓練。 Olatunji Ruwase 開發了一個用來訓練 BLOOM 的 “BF16Optimizer”。

如果您不熟悉這種數據格式，請查看 它的位佈局。 BF16 格式的關鍵是它的指數位數與 FP32 相同，因此不會溢出，但 FP16 經常溢出！FP16 的最大數值範圍為 64k，您只能進行較小數的乘法。例如你可以做 250*250=62500，但如果你嘗試 255*255=65025，你就會溢出，這是導致訓練出現問題的主要原因。這意味着你的權重必須保持很小。一種稱為損失縮放 (loss scaling) 的技術有助於緩解這個問題，但是當模型變得非常大時，FP16 較小的數值範圍仍然是一個問題。

BF16 沒有這個問題，你可以很容易地做 10_000*10_000=100_000_000, 完全沒問題。

當然，由於 BF16 和 FP16 的大小相同，均為 2 個字節，因此，沒有免費的午餐，當使用 BF16 時，代價就是它的精度非常差。然而，你應該還記得我們在訓練時採用的隨機梯度下降法及其變體，該方法有點像蹣跚而行，如果你這步沒有找到完美的方向其實沒關係，你會在接下來的步驟中糾正自己。

無論使用 BF16 還是 FP16，都有一個權重副本始終在 FP32 中 —— 這是由優化器更新的內容。因此 16 位格式僅用於計算，優化器以全精度更新 FP32 權重，然後將它們轉換為 16 位格式以用於下一次迭代。

所有 PyTorch 組件都已更新，以確保它們在 FP32 中執行任何累加，因此不會發生精度損失。

一個關鍵問題是梯度累積，它是流水線並行的主要特徵之一，因為每個 micro batch 處理的梯度都會累積。在 FP32 中實現梯度累積以保證訓練的精確性至關重要，這正是 BF16Optimizer 所做的。

除了其他改進之外，我們認為使用 BF16 混合精度訓練將潛在的噩夢變成了一個相對平穩的過程，這可以從以下 lm 損失圖中看出:

CUDA 融合核函數

GPU 主要做兩件事。它可以將數據寫到顯存或從顯存讀數據，並對這些數據執行計算。當 GPU 忙於讀寫數據時， GPU 的計算單元就會空閒。如果我們想有效地利用 GPU，我們希望將空閒時間降至最低。

核函數是一組實現特定 PyTorch 操作的指令。例如，當你調用 torch.add 時，它會通過一個 PyTorch 調度器，它會根據輸入張量及其他變量的取值來決定它應該運行哪些代碼，最後運行它。 CUDA 核函數使用 CUDA 來實現這些代碼，因此只能在 NVIDIA GPU 上運行。

現在，當使用 GPU 計算 c = torch.add (a, b); e = torch.max ([c,d]) 時，一般情況下，PyTorch 將執行的操作是啟動兩個單獨的核函數，一個執行 a 和 b 的加法，另一個執行取 c 和 d 兩者的最大值。在這種情況下，GPU 從其顯存中獲取 a 和 b，執行加法運算，然後將結果寫回顯存。然後它獲取 c 和 d 並執行 max 操作，然後再次將結果寫回顯存。

如果我們要融合這兩個操作，即將它們放入一個 “融合核函數” 中，然後啟動那個內核，我們不會將中間結果 c 寫到顯存中，而是將其保留在 GPU 寄存器中，並且僅需要獲取 d 來完成最後的計算。這節省了大量開銷並防止 GPU 空閒，因此整個操作會更加高效。

融合核函數就是這樣。它們主要將多個離散的計算和進出顯存的數據移動替換為有很少數據移動的融合計算。此外，一些融合核函數會對操作進行數學變換，以便可以更快地執行某些計算組合。

為了快速高效地訓練 BLOOM，有必要使用 Megatron-LM 提供的幾個自定義 CUDA 融合核函數。特別地，有一個 LayerNorm 的融合核函數以及用於融合縮放、掩碼和 softmax 這些操作的各種組合的核函數。Bias Add 也通過 PyTorch 的 JIT 功能與 GeLU 融合。這些操作都是瓶頸在內存的，因此將它們融合在一起以達到最大化每次顯存讀取後的計算量非常重要。因此，例如，在執行瓶頸在內存的 GeLU 操作時同時執行 Bias Add，運行時間並不會增加。這些核函數都可以在 Megatron-LM 代碼庫 中找到。

數據集

Megatron-LM 的另一個重要特性是高效的數據加載器。在首次訓練啟動前，每個數據集中的每個樣本都被分成固定序列長度 (BLOOM 為 2048) 的樣本，並創建索引以對每個樣本進行編號。基於訓練超參，我們會確定每個數據集所需要參與的 epoch 數，並基於此創建一個有序的樣本索引列表，然後打亂它。舉個例子，如果一個數據集中有 10 個樣本並應參與 2 個 epoch 的訓練，則系統首先按 [0, ..., 9, 0, ..., 9] 順序排好樣本索引，然後打亂該順序為數據集創建最終的全局順序。請注意，這意味着訓練不會簡單地遍歷整個數據集然後重複，你有可能在看到另一個樣本之前看到同一個樣本兩次，但在訓練結束時模型將只看到每個樣本兩次。這有助於確保整個訓練過程中的訓練曲線平滑。這些索引，包括每個樣本在原始數據集中的偏移量，被保存到一個文件中，以避免每次開始訓練時都重新計算它們。最後，可以將其中幾個數據集以不同的權重混合到訓練最終使用的數據中。

嵌入 LayerNorm

在我們努力阻止 104B 模型發散的過程中，我們發現在第一個層詞嵌入層之後添加一個額外的 LayerNorm 可以使訓練更加穩定。

該洞察來自對 bitsandbytes 的實驗，bitsandbytes 有一個 StableEmbedding 操作，它是一個帶有 LayerNorm 的普通嵌入，其使用均勻 xavier 函數來初始化。

位置編碼

基於論文 Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation，我們還用 AliBi 替換了普通的位置嵌入，它允許外推比訓練模型的輸入序列更長的輸入序列。因此，即使我們訓練時使用長度為 2048 的序列，模型也可以在推理過程中處理更長的序列。

訓練中的困難

隨着架構、硬件和軟件的就位，我們得以在 2022 年 3 月上旬開始訓練。然而，從那時起，事情其實並非一帆風順。在本節中，我們將討論我們遇到的一些主要障礙。

在訓練開始之前，有很多問題需要弄清楚。特別是，我們發現了幾個問題，這些問題只有在我們開始在 48 個節點上進行訓練後才會出現，而不會在小規模時出現。例如，需要設 CUDA_LAUNCH_BLOCKING=1 來防止框架掛起，我們需要將優化器組分成更小的組，否則框架會再次掛起。你可以在 訓前編年史 中詳細瞭解這些內容。

訓練期間遇到的主要問題類型是硬件故障。由於這是一個擁有大約 400 個 GPU 的新集羣，平均每週我們會遇到 1-2 個 GPU 故障。我們每 3 小時 (100 次迭代) 保存一個檢查點。因此，我們每週因硬件崩潰平均損失 1.5 小時的訓練成果。 Jean Zay 系統管理員隨後將更換有故障的 GPU 並恢復節點。與此同時，我們有備用節點可供使用。

我們還遇到過多次導致 5-10 小時停機的各種其他問題，其中一些與 PyTorch 中的死鎖錯誤有關，另一些則是由於磁盤空間不足。如果您對具體細節有興趣，請參閲 訓練編年史。

在對訓練這個模型進行可行性分析時，所有這些停機時間都被計劃在內了，我們也據此選擇了合適的模型大小和我們希望模型消耗的數據量。因此，即使存在這些停機問題，我們還是成功地在預計時間內完成了訓練。如前所述，它需要大約 100 萬個計算時才能完成。

另一個問題是 SLURM 並非設計為供一組人使用。 SLURM 作業由單個用户擁有，如果他們不在身邊，則該組的其他成員無法對正在運行的作業執行任何操作。我們制定了一個終止方案，允許組中的其他用户終止當前進程，而不需要啟動該進程的用户在場。這在 90% 的問題上都很有效。如果 SLURM 設計者讀到這篇文章，請添加一個 Unix 組的概念，這樣一個 SLURM 作業就可以由一個組擁有。

由於訓練是全天候 24/7 進行的，我們需要有人隨叫隨到 - 但由於我們在歐洲和加拿大西海岸都有人，因此不需要有人攜帶傳呼機，我們能很好地互相備份。當然，週末的訓練也得有人看着。我們自動化了大部分事情，包括自動從硬件崩潰中恢復，但有時仍需要人工干預。

結論

訓練中最困難和最緊張的部分是訓練開始前的 2 個月。我們承受着儘快開始訓練的巨大壓力，因為資源分配的時間有限，我們直到最後一刻才接觸到 A100。所以這是一個非常困難的時期，考慮到 BF16Optimizer 是在最後一刻編寫出來的，我們需要調試它並修復各種 bug。正如上一節所述，我們發現了新問題，這些問題只有在我們開始在 48 個節點上進行訓練後才會出現，並且不會在小規模時出現。

但是一旦我們把這些整理完，訓練本身出奇的順利，沒有出現大的問題。大多數時候，我們只有一個人看着，只有少數幾個人參與故障排除。我們得到了 Jean Zay 管理部門的大力支持，他們迅速解決了訓練期間出現的大部分需求。

總的來説，這是一次超級緊張但回報頗豐的經歷。

訓練大型語言模型仍然是一項具有挑戰性的任務，但我們希望通過公開構建和共享這項技術，其他人可以借鑑我們的經驗。

資源

重要鏈接

• 主訓練文檔
• Tensorboard
• 訓練用的 slurm 腳本
• 訓練編年史

論文與文章

我們不可能在本文中詳細解釋所有內容，因此如果此處介紹的技術激起你的好奇心，使你想了解更多信息，請閲讀以下論文:

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 聯合:

• Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model.

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，但是論文的其他部分及其技術也很妙，我們沒用 8 位優化器的唯一原因是我們已經使用 DeepSpeed-ZeRO 節省了優化器內存)。

博文致謝

非常感謝以下這些人，他們提出了很好的問題並幫助提高了文章的可讀性 (按字母序):

• Britney Muller,
• Douwe Kiela,
• Jared Casper,
• Jeff Rasley,
• Julien Launay,
• Leandro von Werra,
• Omar Sanseviero,
• Stefan Schweter and
• Thomas Wang.

本文圖表主要由 Chunte Lee 創作。

---

英文原文: http://hf.co/blog/bloom-megatron-deepspeed

原文作者: Stas Bekman

譯者: Matrix Yao (姚偉峯)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態數據上的應用及大規模模型的訓練推理。

審校、排版: zhongdongy (阿東)