BLOOM 訓(xùn)練背后的技術(shù)

@(Engineering Practice)

假設(shè)你現(xiàn)在有了數(shù)據(jù)，也搞到了預(yù)算矩屁，一切就緒辟宗，準(zhǔn)備開始訓(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ù)組件谎僻。

BLOOM

以下是對本項目的簡要總結(jié)：

人員組成

該項目由 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 個群體：

1. HuggingFace 的 BigScience 團隊投入了六名以上的全職員工全程參與了訓(xùn)練的研究和運行，他們還提供或報銷了 Jean Zay 計算機之外的所有基礎(chǔ)設(shè)施蕉朵。
2. Microsoft DeepSpeed 團隊崔涂，開發(fā)了 DeepSpeed，后來將其與 Megatron-LM 集成始衅，其開發(fā)人員花費數(shù)周時間研究項目需求冷蚂，并在訓(xùn)練前和訓(xùn)練期間提供了許多很棒的實用經(jīng)驗建議。
3. NVIDIA Megatron-LM 團隊開發(fā)了 Megatron-LM汛闸，他們非常樂于回答我們的大量問題并提供一流的使用建議蝙茶。
4. IDRIS / GENCI 團隊管理著 Jean Zay 超級計算機，他們?yōu)樵擁椖烤栀浟舜罅康乃懔蛷姶蟮南到y(tǒng)管理支持诸老。
5. PyTorch 團隊創(chuàng)建了一個超強的框架尸闸，其余軟件都基于該框架，并且在準(zhǔn)備訓(xùn)練期間非常支持我們孕锄，修復(fù)了多個 bug 并提高了我們所依賴的 PyTorch 組件的訓(xùn)練可用性。
6. BigScience 工程工作組志愿者

很難說出所有為該項目的工程方面做出貢獻(xiàn)的杰出人物的名字苞尝，所以我只列舉 Hugging Face 之外的幾個關(guān)鍵人物畸肆，他們在過去 14 個月中為該項目奠定了工程基礎(chǔ)：

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

我們也感謝所有允許其員工為該項目做出貢獻(xiàn)的公司抡砂。

概述

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ò)拓?fù)錇闊o阻塞胖樹
• 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 個詞元
• 更詳細(xì)信息杉武，請參閱 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)每個組件的更多詳細(xì)信息曹铃，請參見下表缰趋。

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

下表列出了我們在訓(xùn)練 BLOOM 時各采用了兩個框架的哪些組件：

請注意，Megatron-LM 和 DeepSpeed 都有流水線并行和 BF16 優(yōu)化器實現(xiàn)评甜，但我們使用 DeepSpeed 的實現(xiàn)灰粮，因為它們集成進了 ZeRO。

Megatron-DeepSpeed 實現(xiàn)了 3D 并行以允許大模型以非常有效的方式進行訓(xùn)練忍坷。我們簡要討論一下有哪些 3D 組件粘舟。

1. 數(shù)據(jù)并行 (Data Parallelism，DP) - 相同的設(shè)置和模型被復(fù)制多份佩研，每份每次都被饋送不同的一份數(shù)據(jù)柑肴。處理是并行完成的，所有份在每個訓(xùn)練步結(jié)束時同步旬薯。
2. 張量并行 (Tensor Parallelism晰骑，TP) - 每個張量都被分成多個塊，因此張量的每個分片都位于其指定的 GPU 上绊序，而不是讓整個張量駐留在單個 GPU 上些侍。在處理過程中，每個分片在不同的 GPU 上分別并行處理政模，結(jié)果在步驟結(jié)束時同步岗宣。這就是所謂的水平并行，因為是做的水平拆分淋样。
3. 流水線并行 (Pipeline Parallelism耗式，PP) - 模型在多個 GPU 上垂直（即按層）拆分，因此只有一個或多個模型層放置在單個 GPU 上。每個 GPU 并行處理流水線的不同階段刊咳，并處理 batch 的一部分?jǐn)?shù)據(jù)彪见。
4. 零冗余優(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)練速度贴届，解決問題靠粪。但它有個限制，即只有當(dāng)模型能夠放進單個 GPU 時才有效毫蚓。

ZeRO 數(shù)據(jù)并行

下圖很好地描述了 ZeRO 數(shù)據(jù)并行（來自 此博文）占键。

ZeRO DP

看上去比較高大上，可能讓你很難專心去理解元潘，但實際上畔乙，這個概念非常簡單。這只是通常的 DDP柬批，只是沒有每個 GPU 都復(fù)制完整的模型參數(shù)、梯度和優(yōu)化器狀態(tài)袖订，而是每個 GPU 只存儲其中的一部分氮帐。在隨后的運行過程中，當(dāng)需要給定層的完整層參數(shù)時洛姑，所有 GPU 同步以相互提供它們?nèi)笔У牟糠?—— 僅此而已上沐。

該組件由 DeepSpeed 實現(xiàn)。

張量并行

在張量并行 (TP) 中楞艾，每個 GPU 僅處理張量的一部分参咙，并且僅當(dāng)某些算子需要完整的張量時才觸發(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 之間拆分：

并行 GEMM

如果我們將權(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 論文作者為此提供了一個不錯的圖示：

MLP TP

這里 f 是前向傳播中的恒等運算符返奉，后向傳播中的 all reduce贝搁，而 g 是前向傳播中的 all reduce 和后向傳播中的恒等式。

并行化多頭注意力層甚至更簡單芽偏，因為它們本來就是并行的雷逆，因為有多個獨立的頭！

Self Attention TP

需要特別考慮的是：由于前向和后向傳播中每層都有兩個 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ì)信息呵俏。序列并行是在訓(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ù)模型的拓?fù)鋱D的，我們其實是垂直切分模型各層的赁咙。例如钮莲，如果下圖顯示一個 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 層時链瓦，這就跟單 GPU 上的普通前向傳播一樣。但是當(dāng)數(shù)據(jù)需要從第 3 層傳到第 4 層時盯桦，它需要從 GPU0 傳輸?shù)?GPU1慈俯，這會引入通信開銷。如果參與的 GPU 位于同一計算節(jié)點（例如同一臺物理機器）上拥峦，則傳輸非程欤快埠帕，但如果 GPU 位于不同的計算節(jié)點（例如多臺機器）上线衫，通信開銷可能會大得多。

然后第 4 到 5 到 6 到 7 層又像普通模型一樣，當(dāng)?shù)?7 層完成時璃哟，我們通常需要將數(shù)據(jù)發(fā)送回標(biāo)簽所在的第 0 層（或者將標(biāo)簽發(fā)送到最后一層）。現(xiàn)在可以計算損失喊递，然后使用優(yōu)化器來進行更新參數(shù)了随闪。

問題：

• 該方法為什么被稱為 樸素 流水線并行呢，它又有什么缺陷呢骚勘？主要是因為該方案在任意給定時刻除了一個 GPU 之外的其他所有 GPU 都是空閑的铐伴。因此撮奏，如果使用 4 個 GPU，則幾乎等同于將單個 GPU 的內(nèi)存量翻兩番当宴，而其他資源（如計算）相當(dāng)于沒用上畜吊。另外還需要加上在設(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 達(dá)到高的并發(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é)合起來踢械，如下所示。

2D Parallelism

這里重要的是要了解 DP rank 0 是看不見 GPU2 的魄藕， DP rank 1 是看不到 GPU3 的内列。對于 DP 而言，只有 GPU 0 和 1背率，并向它們饋送數(shù)據(jù)话瞧。 GPU0 使用 PP “秘密地” 將它的一些負(fù)載卸載到 GPU2。同樣地寝姿， GPU1 也會得到 GPU3 的幫助交排。

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

DP+PP+TP

為了更高效地訓(xùn)練埃篓，可以將 PP、TP 和 DP 相結(jié)合根资，稱為 3D 并行架专，如下圖所示同窘。

3D Parallelism

此圖來自博文 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é)合使用。

當(dāng) 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)境進行仔細(xì)評估后崭添，我們發(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é)到了很多：

FP16 lm-los

我們也從 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ù)有助于緩解這個問題，但是當(dāng)模型變得非常大時均牢，F(xiàn)P16 較小的數(shù)值范圍仍然是一個問題糠雨。

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

當(dāng)然见秤，由于 BF16 和 FP16 的大小相同，均為 2 個字節(jié)真椿，因此鹃答，沒有免費的午餐，當(dāng)使用 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 所做的吆玖。

除了其他改進之外筒溃，我們認(rèn)為使用 BF16 混合精度訓(xùn)練將潛在的噩夢變成了一個相對平穩(wěn)的過程，這可以從以下 lm 損失圖中看出：

BF16 lm-loss

CUDA 融合核函數(shù)

GPU 主要做兩件事沾乘。它可以將數(shù)據(jù)寫到顯存或從顯存讀數(shù)據(jù)铡羡，并對這些數(shù)據(jù)執(zhí)行計算。當(dāng) GPU 忙于讀寫數(shù)據(jù)時意鲸， GPU 的計算單元就會空閑烦周。如果我們想有效地利用 GPU，我們希望將空閑時間降至最低怎顾。

核函數(shù)是一組實現(xiàn)特定 PyTorch 操作的指令读慎。例如，當(dāng)你調(diào)用 torch.add 時槐雾，它會通過一個 PyTorch 調(diào)度器夭委，它會根據(jù)輸入張量及其他變量的取值來決定它應(yīng)該運行哪些代碼，最后運行它募强。 CUDA 核函數(shù)使用 CUDA 來實現(xiàn)這些代碼株灸，因此只能在 NVIDIA GPU 上運行。

現(xiàn)在擎值，當(dāng)使用 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)诤显谝黄鹨赃_(dá)到最大化每次顯存讀取后的計算量非常重要。因此碉克，例如凌唬，在執(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)前編年史 中詳細(xì)了解這些內(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)坊饶，另一些則是由于磁盤空間不足泄伪。如果您對具體細(xì)節(jié)有興趣，請參閱 訓(xùn)練編年史匿级。

在對訓(xùn)練這個模型進行可行性分析時蟋滴，所有這些停機時間都被計劃在內(nèi)了，我們也據(jù)此選擇了合適的模型大小和我們希望模型消耗的數(shù)據(jù)量痘绎。因此津函，即使存在這些停機問題，我們還是成功地在預(yù)計時間內(nèi)完成了訓(xùn)練孤页。如前所述尔苦，它需要大約 100 萬個計算時才能完成。

另一個問題是 SLURM 并非設(shè)計為供一組人使用行施。 SLURM 作業(yè)由單個用戶擁有允坚，如果他們不在身邊，則該組的其他成員無法對正在運行的作業(yè)執(zhí)行任何操作蛾号。我們制定了一個終止方案稠项，允許組中的其他用戶終止當(dāng)前進程，而不需要啟動該進程的用戶在場鲜结。這在 90% 的問題上都很有效展运。如果 SLURM 設(shè)計者讀到這篇文章活逆，請?zhí)砑右粋€ Unix 組的概念，這樣一個 SLURM 作業(yè)就可以由一個組擁有拗胜。

由于訓(xùn)練是全天候 24/7 進行的蔗候，我們需要有人隨叫隨到 - 但由于我們在歐洲和加拿大西海岸都有人，因此不需要有人攜帶傳呼機埂软，我們能很好地互相備份琴庵。當(dāng)然，周末的訓(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)驗共缕。

資源

重要鏈接

• 主訓(xùn)練文檔
• tensorboard
• 訓(xùn)練用的 slurm 腳本
• 訓(xùn)練編年史

論文與文章

我們不可能在本文中詳細(xì)解釋所有內(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)合:

• 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阱洪，但是論文的其他部分及其技術(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)練推理。