BLOOM 訓練背後的技術

近年來，訓練越來越大的語言模型已成為常態。雖然這些模型未釋出以供進一步研究的問題經常被討論，但關於如何訓練這些模型的隱藏知識卻很少受到關注。本文旨在改變這一現狀，以 1760 億引數的語言模型 BLOOM 為例，從硬體和軟體兩方面闡述訓練此類模型背後的技術和工程。

但首先，我們要感謝那些公司、關鍵人物和團隊，是他們使得一小群敬業的人能夠完成訓練一個 1760 億引數模型的驚人壯舉。

然後將討論硬體設定和主要技術元件。

以下是專案概覽

人員

該專案由 Thomas Wolf (Hugging Face 聯合創始人兼 CSO) 構思，他敢於與大公司競爭，不僅要訓練出最大的多語言模型之一，還要讓最終成果惠及所有人，從而將大多數人的夢想變為現實。

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

主要有 6 組人員需要感謝：

1. Hugging Face 的 BigScience 團隊，他們投入了超過六名全職員工，從專案開始到結束負責規劃和執行訓練，並提供和支付了 Jean Zay 計算資源之外的所有基礎設施費用。
2. 微軟 DeepSpeed 團隊，他們開發了 DeepSpeed 並後來將其與 Megatron-LM 整合。其開發人員花費數週時間滿足專案需求，並在訓練前後提供了大量寶貴的實踐經驗建議。
3. NVIDIA Megatron-LM 團隊，他們開發了 Megatron-LM，非常熱心地回答了我們的大量問題，並提供了一流的經驗建議。
4. 管理 Jean Zay 超級計算機的 IDRIS / GENCI 團隊，他們為專案捐贈了海量的計算資源和出色的系統管理支援。
5. PyTorch 團隊，他們建立了一個功能強大的框架，其他軟體都基於此框架。他們在我們準備訓練期間給予了大力支援，修復了多個錯誤，並改善了我們訓練中依賴的 PyTorch 元件的可用性。
6. BigScience 工程工作組的志願者們。

很難一一列舉所有為該專案工程方面做出貢獻的傑出人士，因此我只提及幾位在過去 14 個月中作為該專案工程基石的 Hugging Face 之外的關鍵人物：

Olatunji Ruwase, Deepak Narayanan, Jeff Rasley, Jared Casper, Samyam Rajbhandari 以及 Rémi Lacroix。

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

概述

BLOOM 的架構與 GPT3 非常相似，並增加了一些改進，本文稍後會討論。

該模型在 Jean Zay 上進行訓練，這是一臺由法國政府資助的超級計算機，由 GENCI 管理，並安裝在法國國家科學研究中心 (CNRS) 的國家計算中心 IDRIS。計算資源由 GENCI 慷慨捐贈給該專案 (撥款號 2021-A0101012475)。

訓練期間使用了以下硬體：

• GPU：384 塊 NVIDIA A100 80GB GPU (48 個節點) + 32 塊備用 GPU
• 每節點 8 塊 GPU，使用 NVLink 4 進行 GPU 間互聯，4 個 OmniPath 連結
• CPU：AMD EPYC 7543 32 核處理器
• CPU 記憶體：每節點 512GB
• GPU 記憶體：每節點 640GB
• 節點間連線：Omni-Path 架構 (OPA)，採用無阻塞胖樹結構
• NCCL 通訊網路：一個完全專用的子網
• 磁碟 IO 網路：GPFS 與其他節點和使用者共享

檢查點

• 主要檢查點
• 每個包含 fp32 最佳化器狀態和 bf16+fp32 權重的檢查點大小為 2.3TB——僅 bf16 權重就為 329GB。

資料集

• 46 種語言，1.5TB 經過去重和大規模清洗的文字，轉換為 3500 億個唯一的詞元
• 模型的詞彙表大小為 250,680 個詞元
• 完整詳情請參見 The BigScience Corpus A 1.6TB Composite Multilingual Dataset

1760 億引數的 BLOOM 模型訓練發生在 2022 年 3 月至 7 月，耗時約 3.5 個月 (約 100 萬計算小時)。

Megatron-DeepSpeed

1760 億引數的 BLOOM 模型是使用 Megatron-DeepSpeed 進行訓練的，它結合了兩種主要技術：

• DeepSpeed 是一個深度學習最佳化庫，使分散式訓練變得簡單、高效和有效。
• Megatron-LM 是一個由 NVIDIA 應用深度學習研究團隊開發的大型、強大的 Transformer 模型框架。

DeepSpeed 團隊透過將 DeepSpeed 庫中的 ZeRO 分片和流水線並行與 Megatron-LM 中的張量並行相結合，開發了一種基於 3D 並行的實現。下表中可以看到每個元件的更多細節。

請注意，BigScience 的 Megatron-DeepSpeed 是原始 Megatron-DeepSpeed 倉庫的一個分支，我們在此基礎上添加了多項功能。

下表列出了訓練 BLOOM 時各個元件由哪個框架提供：

請注意，Megatron-LM 和 DeepSpeed 都有流水線並行和 BF16 最佳化器的實現，但我們使用了 DeepSpeed 的版本，因為它們與 ZeRO 整合在一起。

Megatron-DeepSpeed 實現了 3D 並行，以非常高效的方式訓練大型模型。讓我們簡要討論一下這 3 個維度。

1. 資料並行 (DP) - 相同的設定被複制多次，每個副本處理一部分資料。處理過程並行進行，在每個訓練步驟結束時所有設定進行同步。
2. 張量並行 (TP) - 每個張量被分割成多個塊，因此整個張量不再駐留在單個 GPU 上，而是每個分片駐留在其指定的 GPU 上。在處理過程中，每個分片在不同的 GPU 上被獨立且並行地處理，結果在步驟結束時同步。這可以稱為水平並行，因為切分是在水平層面上發生的。
3. 流水線並行 (PP) - 模型在多個 GPU 上垂直 (層級) 拆分，因此只有一個或幾個模型層被放置在單個 GPU 上。每個 GPU 並行處理流水線的不同階段，並處理一小塊批次資料。
4. 零冗餘最佳化器 (ZeRO) - 也執行張量的分片，有點類似於 TP，不同之處在於整個張量會在前向或後向計算時被及時重構，因此模型不需要被修改。它還支援各種解除安裝技術以彌補有限的 GPU 記憶體。

資料並行

大多數只有幾塊 GPU 的使用者可能都熟悉 `DistributedDataParallel` (DDP) PyTorch 文件。在這種方法中，模型被完全複製到每個 GPU 上，然後在每次迭代後，所有模型相互同步它們的狀態。這種方法透過投入更多資源來加快訓練速度，但它只有在模型能裝入單個 GPU 的情況下才有效。

ZeRO 資料並行

基於 ZeRO 的資料並行 (ZeRO-DP) 在這篇 博文 中的下圖有描述

這可能很難理解，但實際上，這個概念相當簡單。這只是普通的 DDP，不同之處在於，不是複製完整的模型引數、梯度和最佳化器狀態，每個 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，除非你有非常快的網路。在我們的案例中，節點間通訊比 PCIe 慢得多。實際上，如果一個節點有 4 個 GPU，那麼最高的 TP 度就是 4。如果你需要 8 的 TP 度，你需要使用至少有 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 層，那裡有標籤 (或者將標籤傳送到最後一層)。現在可以計算損失，最佳化器可以開始工作。

問題

• 主要缺陷以及為什麼這種方法被稱為“樸素”PP，是在任何給定時刻只有一個 GPU 在工作。因此，如果使用 4 個 GPU，這幾乎等同於將單個 GPU 的記憶體量增加四倍，而忽略了其餘的硬體。此外，還存在裝置之間複製資料的開銷。所以，4 張 6GB 的顯示卡可以容納與 1 張 24GB 顯示卡相同大小的模型，但後者會更快地完成訓練，因為它沒有資料複製的開銷。但是，比如說，如果你有 40GB 的顯示卡，需要容納一個 45GB 的模型，你可以用 4 張 40GB 的顯示卡 (但由於梯度和最佳化器狀態的原因，會很勉強)。
• 共享的嵌入層可能需要在 GPU 之間來回複製。

流水線並行 (PP) 與上面描述的樸素 PP 幾乎相同，但它透過將輸入批次分塊為微批次並人為地建立一個流水線來解決 GPU 空閒問題，這使得不同的 GPU 可以同時參與計算過程。

以下來自 GPipe 論文 的插圖顯示了頂部的樸素 PP 和底部的 PP。

從底部圖表中可以很容易地看出 PP 如何減少了 GPU 空閒的死區。這些空閒部分被稱為“氣泡”。

圖表的兩個部分都顯示了並行度為 4 的情況。也就是說，有 4 個 GPU 參與流水線。因此，有 4 個流水線階段的前向路徑 F0、F1、F2 和 F3，然後是反向路徑 B3、B2、B1 和 B0。

PP 引入了一個需要調整的新超引數，稱為 `chunks`。它定義了透過同一個流水線階段連續傳送的資料塊數量。例如，在底部圖表中，你可以看到 `chunks=4`。GPU0 對塊 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)。

注意，從概念上講，這與梯度累積步數 (GAS) 是相同的概念。PyTorch 使用 `chunks`，而 DeepSpeed 將同一個超引數稱為 GAS。

由於 `chunks` 的存在，PP 引入了微批次 (MBS) 的概念。DP 將全域性資料批次大小劃分為小批次，所以如果你的 DP 度為 4，全域性批次大小為 1024 會被劃分為 4 個 256 的小批次 (1024/4)。如果 `chunks` (或 GAS) 的數量是 32，我們最終得到的微批次大小為 8 (256/32)。每個流水線階段一次只處理一個微批次。

要計算 DP + PP 設定的全域性批次大小，我們這樣做：`mbs * chunks * dp_degree` (`8 * 32 * 4 = 1024`)。

讓我們回到圖表。

當 `chunks=1` 時，你得到的是樸素 PP，效率非常低。當 `chunks` 值非常大時，你得到的微批次大小會非常小，這可能也不高效。因此，必須透過實驗找到能帶來最高效 GPU 利用率的值。

雖然圖表顯示存在一個無法並行的“死區”時間氣泡，因為最後一個`前向`階段必須等待`後向`完成流水線，但找到最佳 `chunks` 值的目的在於實現所有參與 GPU 的高併發利用率，這相當於最小化氣泡的大小。

這種排程機制被稱為 `all forward all backward`。其他一些替代方案是 one forward one backward 和 interleaved one forward one backward。

雖然 Megatron-LM 和 DeepSpeed 都有自己的 PP 協議實現，但 Megatron-DeepSpeed 使用 DeepSpeed 的實現，因為它與 DeepSpeed 的其他方面整合在一起。

另一個重要問題是詞嵌入矩陣的大小。雖然通常詞嵌入矩陣消耗的記憶體比 Transformer 塊少，但在我們擁有 25 萬詞彙表的巨大情況下，嵌入層在 bf16 權重中需要 7.2GB，而 Transformer 塊僅為 4.9GB。因此，我們必須指示 Megatron-Deepspeed 將嵌入層視為一個 Transformer 塊。所以我們有一個 72 層的流水線，其中 2 層專用於嵌入 (第一層和最後一層)。這使得 GPU 記憶體消耗得以平衡。如果我們不這樣做，第一和最後一個階段會消耗大部分 GPU 記憶體，而 95% 的 GPU 將使用少得多的記憶體，因此訓練效率會很低。

DP+PP

以下來自 DeepSpeed 流水線教程 的圖表演示瞭如何將 DP 與 PP 結合起來。

在這裡，重要的是要看到 DP rank 0 看不到 GPU2，DP rank 1 看不到 GPU3。對 DP 來說，只有 GPU 0 和 GPU 1，它向它們饋送資料，就好像只有 2 個 GPU 一樣。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 資料並行 中討論過。通常它是一個獨立的功能，不需要 PP 或 TP。但它可以與 PP 和 TP 結合使用。

當 ZeRO-DP 與 PP (以及可選的 TP) 結合時，通常只啟用 ZeRO stage 1，它只分片最佳化器狀態。ZeRO stage 2 額外分片梯度，而 stage 3 也分片模型權重。

雖然理論上可以將 ZeRO stage 2 與流水線並行一起使用，但這會對效能產生不良影響。每個微批次都需要一個額外的 reduce-scatter 集合操作來在分片前聚合梯度，這會增加潛在的顯著通訊開銷。根據流水線並行的特性，會使用小的微批次，而重點在於平衡計算強度 (微批次大小) 與最小化流水線氣泡 (微批次數量)。因此，這些通訊成本將會造成損害。

此外，由於 PP 的原因，層數已經比正常情況少，因此記憶體節省不會很大。PP 已經將梯度大小減少了 `1/PP`，因此在此基礎上的梯度分片節省比純 DP 的情況要小。

ZeRO stage 3 也可以用來訓練這種規模的模型，但是，它比 DeepSpeed 3D 並行實現需要更多的通訊。在一年前我們對環境進行仔細評估後，我們發現 Megatron-DeepSpeed 3D 並行表現最好。從那時起，ZeRO stage 3 的效能得到了顯著提升，如果今天我們再次評估，也許我們會選擇 stage 3。

BF16Optimizer

用 FP16 訓練大型語言模型是不可行的。

我們透過花費數月時間 訓練一個 1040 億引數模型 證明了這一點，從 tensorboard 可以看出，那次訓練完全失敗了。在與不斷髮散的 lm-loss 作鬥爭的過程中，我們學到了很多東西。

在 Megatron-LM 和 DeepSpeed 團隊訓練了 5300 億引數模型 後，我們也從他們那裡得到了同樣的建議。最近釋出的 OPT-175B 也報告說他們在 FP16 訓練中遇到了很大的困難。

因此，早在 1 月份，我們就知道我們將在支援 BF16 格式的 A100 上進行訓練，Olatunji Ruwase 開發了一個 `BF16Optimizer`，我們用它來訓練 BLOOM。

如果你不熟悉這種資料格式，請檢視位佈局。BF16 格式的關鍵在於它與 FP32 有相同的指數，因此不會遭受 FP16 經常遇到的溢位問題！對於 FP16，其最大數值範圍為 64k，你只能乘以較小的數字。例如，你可以做 `250*250=62500`，但如果你嘗試 `255*255=65025`，就會發生溢位，這是訓練期間的主要問題。這意味著你的權重必須保持很小。一種稱為損失縮放的技術可以幫助解決這個問題，但當模型變得非常大時，FP16 的有限範圍仍然是一個問題。

BF16 沒有這樣的問題，你可以輕鬆地做 `10_000*10_000=100_000_000`，毫無問題。

當然，由於 BF16 和 FP16 的大小都是 2 位元組，天下沒有免費的午餐，使用 BF16 會付出精度非常差的代價。然而，如果你還記得使用隨機梯度下降及其變體的訓練過程是一種蹣跚前行的方式，所以如果你沒有立即得到完美的方向也沒關係，你會在接下來的步驟中自我修正。

無論使用 BF16 還是 FP16，總會有一個 FP32 的權重副本——這是由最佳化器更新的。所以 16 位格式只用於計算，最佳化器用全精度更新 FP32 權重，然後在下一次迭代中將它們轉換為 16 位格式。

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

一個關鍵問題是梯度累積，這也是流水線並行的主要特性之一，因為來自每個微批次處理的梯度會被累積。在 FP32 中實現梯度累積以保持訓練的精確性至關重要，而這正是 `BF16Optimizer` 所做的。

除了其他改進之外，我們相信使用 BF16 混合精度訓練將一個潛在的噩夢變成了一個相對順利的過程，這可以從以下 lm loss 圖中觀察到：

融合 CUDA 核心

GPU 執行兩件事。它可以將資料複製到記憶體或從記憶體中複製資料，並對這些資料進行計算。當 GPU 忙於複製時，GPU 的計算單元處於空閒狀態。如果我們想高效地利用 GPU，我們希望最小化空閒時間。

一個核心 (kernel) 是一組實現特定 PyTorch 操作的指令。例如，當你呼叫 `torch.add` 時，它會透過一個 PyTorch 排程器，該排程器會檢視輸入的張量和其他各種東西，然後決定應該執行哪個程式碼，並執行它。CUDA 核心是一種使用 CUDA API 庫的特定實現，只能在 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 操作的各種組合的核心。偏置項的新增也使用 PyTorch 的 JIT 功能與 GeLU 操作融合。這些操作都是記憶體密集型的，因此將它們融合以最大化從記憶體中檢索到值後完成的計算量非常重要。例如，在已經進行記憶體密集型的 GeLU 操作時新增偏置項不會增加額外的時間。這些核心都可以在 Megatron-LM 倉庫中找到。

資料集

Megatron-LM 的另一個重要特性是高效的資料載入器。在初次訓練啟動時，每個資料集被分割成請求序列長度 (BLOOM 為 2048) 的樣本，並建立一個索引來編號每個樣本。根據訓練引數計算資料集的 epoch 數，併為那麼多 epoch 建立一個順序，然後進行洗牌。例如，如果一個數據集有 10 個樣本，應該遍歷兩次，系統首先按順序排列樣本索引 `[0, ..., 9, 0, ..., 9]`，然後打亂這個順序以建立該資料集的最終全域性順序。請注意，這意味著訓練不會簡單地遍歷整個資料集然後重複，有可能在看到另一個樣本之前看到同一個樣本兩次，但在訓練結束時，模型將看到每個樣本兩次。這有助於確保整個訓練過程中的平滑訓練曲線。這些索引，包括每個樣本在基礎資料集中的偏移量，都會儲存到一個檔案中，以避免每次啟動訓練過程時重新計算。然後可以將這些資料集中的幾個以不同的權重混合到訓練過程最終看到的資料中。

嵌入層歸一化

在我們努力阻止 1040 億引數模型發散時，我們發現在第一個詞嵌入層之後立即新增一個額外的層歸一化 (LayerNorm) 會使訓練更加穩定。

這一見解來自於對 bitsandbytes 的實驗，該庫包含一個 `StableEmbedding`，它是一個帶有層歸一化的普通嵌入層，並使用均勻 xavier 初始化。

位置編碼

我們還將通常的位置嵌入替換為 AliBi——基於論文：Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation，這使得模型能夠外推到比訓練時更長的輸入序列。因此，即使我們在長度為 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 全天候進行的，我們需要有人值班——但由於我們在歐洲和加拿大西海岸都有人，總體上不需要有人攜帶傳呼機，我們的工作時間可以很好地重疊。當然，週末也需要有人照看訓練。我們自動化了大部分事情，包括從硬體崩潰中恢復，但有時也需要人工干預。

結論

訓練中最困難和最緊張的部分是開始訓練前的兩個月。我們面臨著儘快開始訓練的巨大壓力，因為資源分配時間有限，而且直到最後一刻我們才接觸到 A100。所以那是一段非常艱難的時期，考慮到 `BF16Optimizer` 是在最後一刻編寫的，我們需要除錯它並修復各種錯誤。正如上一節所解釋的，我們發現了只有在 48 個節點上開始訓練時才會出現的新問題，這些問題在小規模下不會出現。

但是一旦我們解決了這些問題，訓練本身就出奇地順利，沒有出現重大問題。大多數時候我們只有一個人監控訓練，只有幾次需要幾個人參與故障排除。我們得到了 Jean Zay 管理部門的大力支援，他們迅速解決了訓練期間出現的大多數需求。

總的來說，這是一次超級緊張但非常有益的經歷。

訓練大型語言模型仍然是一項具有挑戰性的任務，但我們希望透過開放地構建和分享這項技術，其他人可以在我們的經驗基礎上繼續發展。

資源

重要連結

• 主要訓練文件
• tensorboard
• 訓練 slurm 指令碼
• 訓練紀事

論文和文章

我們不可能在本文中詳細解釋所有內容，所以如果這裡介紹的技術激起了您的好奇心，並且您想了解更多，這裡有一些值得閱讀的論文：

Megatron-LM

• 在 GPU 叢集上高效訓練大規模語言模型 (Efficient Large-Scale Language Model Training on GPU Clusters).
• 減少大型 Transformer 模型中的啟用重計算 (Reducing Activation Recomputation in Large Transformer Models)

DeepSpeed

• ZeRO：記憶體最佳化以訓練萬億引數模型
• ZeRO-Offload：民主化百億級模型訓練
• ZeRO-Infinity：突破 GPU 記憶體牆，實現超大規模深度學習
• DeepSpeed：面向所有人的超大規模模型訓練 (DeepSpeed: Extreme-scale model training for everyone)

Megatron-LM 和 DeepSpeed 聯合 (Joint Megatron-LM and Deepspeeed)

• 使用 DeepSpeed 和 Megatron 訓練 Megatron-Turing NLG 530B，一個大規模生成語言模型 (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)
• 如果你有一百萬個 GPU 小時，應該訓練什麼樣的語言模型？(What Language Model to Train if You Have One Million GPU Hours?) - 在那裡你會找到導致我們選擇 ALiBi 的實驗。

BitsNBytes

• 透過塊級量化實現 8 位最佳化器 (8-bit Optimizers via Block-wise Quantization) (在嵌入層歸一化的背景下，但論文的其餘部分和技術都非常出色——我們沒有使用 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 創作。