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

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

@(Engineering Practice)

假設(shè)你現(xiàn)在有了數(shù)據(jù)，也搞到了預(yù)算，一切就緒，準(zhǔn)備開始訓(xùn)練一個(gè)大模型，一顯身手了，“一朝看盡長(zhǎng)安花”似乎近在眼前...... 且慢！訓(xùn)練可不僅僅像這兩個(gè)字的發(fā)音那么簡(jiǎn)單，看看 BLOOM 的訓(xùn)練或許對(duì)你有幫助。

近年來(lái)，語(yǔ)言模型越訓(xùn)越大已成為常態(tài)。大家通常會(huì)詬病這些大模型本身的信息未被公開以供研究，但很少關(guān)注大模型訓(xùn)練技術(shù)這種背后的知識(shí)。本文旨在以 1760 億參數(shù)的語(yǔ)言模型 BLOOM 為例，闡明訓(xùn)練此類模型背后的軟硬件工程和技術(shù)要點(diǎn)，以促進(jìn)大家對(duì)大模型訓(xùn)練技術(shù)的討論。

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

然后，我們開始討論硬件配置和主要技術(shù)組件。

BLOOM

以下是對(duì)本項(xiàng)目的簡(jiǎn)要總結(jié)：

項(xiàng) 明細(xì)
硬件 384 張 80GB A100 GPU
軟件 Megatron-DeepSpeed
模型架構(gòu) 基于 GPT3
數(shù)據(jù)集 含 59 種語(yǔ)言，共 3500 億詞元
訓(xùn)練時(shí)長(zhǎng) 3.5 個(gè)月

人員組成

該項(xiàng)目由 Thomas Wolf（Hugging Face 聯(lián)合創(chuàng)始人兼 CSO）發(fā)想，他敢于與大公司競(jìng)爭(zhēng)，提出不僅要訓(xùn)練出立于世界上最大的多語(yǔ)言模型之林的模型，還要讓所有人都可以公開訪問(wèn)訓(xùn)練結(jié)果，圓了大多數(shù)人的夢(mèng)想。

本文主要關(guān)注模型訓(xùn)練的工程方面。 BLOOM 背后的技術(shù)中最重要的部分是分享專業(yè)知識(shí)并幫助我們進(jìn)行編碼和訓(xùn)練的人員和公司。

我們主要需要感謝 6 個(gè)群體：

  1. HuggingFace 的 BigScience 團(tuán)隊(duì)投入了六名以上的全職員工全程參與了訓(xùn)練的研究和運(yùn)行，他們還提供或報(bào)銷了 Jean Zay 計(jì)算機(jī)之外的所有基礎(chǔ)設(shè)施。
  2. Microsoft DeepSpeed 團(tuán)隊(duì)，開發(fā)了 DeepSpeed，后來(lái)將其與 Megatron-LM 集成，其開發(fā)人員花費(fèi)數(shù)周時(shí)間研究項(xiàng)目需求，并在訓(xùn)練前和訓(xùn)練期間提供了許多很棒的實(shí)用經(jīng)驗(yàn)建議。
  3. NVIDIA Megatron-LM 團(tuán)隊(duì)開發(fā)了 Megatron-LM，他們非常樂(lè)于回答我們的大量問(wèn)題并提供一流的使用建議。
  4. IDRIS / GENCI 團(tuán)隊(duì)管理著 Jean Zay 超級(jí)計(jì)算機(jī)，他們?yōu)樵擁?xiàng)目捐贈(zèng)了大量的算力和強(qiáng)大的系統(tǒng)管理支持。
  5. PyTorch 團(tuán)隊(duì)創(chuàng)建了一個(gè)超強(qiáng)的框架，其余軟件都基于該框架，并且在準(zhǔn)備訓(xùn)練期間非常支持我們，修復(fù)了多個(gè) bug 并提高了我們所依賴的 PyTorch 組件的訓(xùn)練可用性。
  6. BigScience 工程工作組志愿者

很難說(shuō)出所有為該項(xiàng)目的工程方面做出貢獻(xiàn)的杰出人物的名字，所以我只列舉 Hugging Face 之外的幾個(gè)關(guān)鍵人物，他們?cè)谶^(guò)去 14 個(gè)月中為該項(xiàng)目奠定了工程基礎(chǔ)：

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

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

概述

BLOOM 的模型架構(gòu)與 GPT3 非常相似，只是增加了一些改進(jìn)，本文稍后將對(duì)此進(jìn)行討論。

該模型是在 Jean Zay 上訓(xùn)練的，Jean Zay 是由 GENCI 管理的法國(guó)政府資助的超級(jí)計(jì)算機(jī)，安裝在法國(guó)國(guó)家科學(xué)研究中心 (CNRS) 的國(guó)家計(jì)算中心 IDRIS。訓(xùn)練所需的算力由 GENCI 慷慨捐贈(zèng)給本項(xiàng)目（捐贈(zèng)號(hào) 2021-A0101012475）。

訓(xùn)練硬件：

  • GPU：384 張 NVIDIA A100 80GB GPU（48 個(gè)節(jié)點(diǎn)）+ 32 張備用 GPU
  • 每個(gè)節(jié)點(diǎn) 8 張 GPU，4 條 NVLink 卡間互聯(lián)，4 條 OmniPath 鏈路
  • CPU：AMD EPYC 7543 32 核處理器
  • CPU 內(nèi)存：每個(gè)節(jié)點(diǎn) 512GB
  • GPU 顯存：每個(gè)節(jié)點(diǎn) 640GB
  • 節(jié)點(diǎn)間連接：使用 Omni-Path Architecture (OPA) 網(wǎng)卡，網(wǎng)絡(luò)拓?fù)錇闊o(wú)阻塞胖樹
  • NCCL - 通信網(wǎng)絡(luò)：一個(gè)完全專用的子網(wǎng)
  • 磁盤 IO 網(wǎng)絡(luò)：GPFS 與其他節(jié)點(diǎn)和用戶共享

Checkpoints:

  • 主 checkpoints
  • 每個(gè) checkpoint 含精度為 fp32 的優(yōu)化器狀態(tài)和精度為 bf16+fp32 的權(quán)重，占用存儲(chǔ)空間為 2.3TB。如只保存 bf16 的權(quán)重，則僅占用 329GB 的存儲(chǔ)空間。

數(shù)據(jù)集:

176B BLOOM 模型的訓(xùn)練于 2022 年 3 月至 7 月期間，耗時(shí)約 3.5 個(gè)月完成（約 100 萬(wàn)計(jì)算時(shí)）。

Megatron-DeepSpeed

176B BLOOM 模型使用 Megatron-DeepSpeed 進(jìn)行訓(xùn)練，它結(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ì)通過(guò)將 DeepSpeed 庫(kù)中的 ZeRO 分片和流水線并行（Pipeline Parallelism）與 Megatron-LM 中的張量并行（Tensor Parallelism）相結(jié)合，開發(fā)了一種基于 3D 并行的方案。有關(guān)每個(gè)組件的更多詳細(xì)信息，請(qǐng)參見下表。

請(qǐng)注意，BigScience 的 Megatron-DeepSpeed 是基于原始 Megatron-DeepSpeed 代碼庫(kù)，我們還在其上添加了不少代碼。

下表列出了我們?cè)谟?xùn)練 BLOOM 時(shí)各采用了兩個(gè)框架的哪些組件：

組件 DeepSpeed Megatron-LM
ZeRO 數(shù)據(jù)并行 是
張量并行 是
流水線并行 是
BF16 優(yōu)化器 是
CUDA 融合核函數(shù) 是
數(shù)據(jù)加載器 是

請(qǐng)注意，Megatron-LM 和 DeepSpeed 都有流水線并行和 BF16 優(yōu)化器實(shí)現(xiàn)，但我們使用 DeepSpeed 的實(shí)現(xiàn)，因?yàn)樗鼈兗蛇M(jìn)了 ZeRO。

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

  1. 數(shù)據(jù)并行 (Data Parallelism，DP) - 相同的設(shè)置和模型被復(fù)制多份，每份每次都被饋送不同的一份數(shù)據(jù)。處理是并行完成的，所有份在每個(gè)訓(xùn)練步結(jié)束時(shí)同步。
  2. 張量并行 (Tensor Parallelism，TP) - 每個(gè)張量都被分成多個(gè)塊，因此張量的每個(gè)分片都位于其指定的 GPU 上，而不是讓整個(gè)張量駐留在單個(gè) GPU 上。在處理過(guò)程中，每個(gè)分片在不同的 GPU 上分別并行處理，結(jié)果在步驟結(jié)束時(shí)同步。這就是所謂的水平并行，因?yàn)槭亲龅乃讲鸱帧?/li>
  3. 流水線并行 (Pipeline Parallelism，PP) - 模型在多個(gè) GPU 上垂直（即按層）拆分，因此只有一個(gè)或多個(gè)模型層放置在單個(gè) GPU 上。每個(gè) GPU 并行處理流水線的不同階段，并處理 batch 的一部分?jǐn)?shù)據(jù)。
  4. 零冗余優(yōu)化器 (Zero Redundancy Optimizer，ZeRO) - 也執(zhí)行與 TP 相類似的張量分片，但整個(gè)張量會(huì)及時(shí)重建以進(jìn)行前向或反向計(jì)算，因此不需要修改模型。它還支持各種卸載技術(shù)以補(bǔ)償有限的 GPU 內(nèi)存。

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

大多數(shù)只有幾張 GPU 的用戶可能比較熟悉 DistributedDataParallel(DDP)，這是相應(yīng)的 PyTorch 文檔。在該方法中，模型被完全復(fù)制到每個(gè) GPU，然后在每次迭代后所有模型相互同步各自的狀態(tài)。這種方法可以通過(guò)投入更多 GPU 資源的方式加快訓(xùn)練速度，解決問(wèn)題。但它有個(gè)限制，即只有當(dāng)模型能夠放進(jìn)單個(gè) GPU 時(shí)才有效。

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

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

ZeRO DP

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

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

張量并行

在張量并行 (TP) 中，每個(gè) GPU 僅處理張量的一部分，并且僅當(dāng)某些算子需要完整的張量時(shí)才觸發(fā)聚合操作。

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

Transformer 類模型的主要模塊為：一個(gè)全連接層 nn.Linear，后面跟一個(gè)非線性激活層 GeLU。

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

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


并行 GEMM

如果我們將權(quán)重矩陣 A 按列拆分到 N 個(gè) GPU 上，然后并行執(zhí)行矩陣乘法 XA_1 到 XA_n，那么我們最終將得到 N 個(gè)輸出向量 Y_1、Y_2、…… 、 Y_n ，它們可以獨(dú)立輸入 GeLU：[Y_1, Y_2]=[GeLU(XA_1), GeLU(XA_2)]

注意因?yàn)?Y 矩陣是按列拆分的，因此隨后的 GEMM 我們可以選擇按行拆分方案，這樣它就可以直接獲取前面層的 GeLU 的輸出，而無(wú)需任何額外的通信。

使用該原理，我們可以更新任意深度的 MLP，只需在每個(gè) 拆列 - 拆行 序列之后同步 GPU。 Megatron-LM 論文作者為此提供了一個(gè)不錯(cuò)的圖示：

MLP TP

這里 f 是前向傳播中的恒等運(yùn)算符，后向傳播中的 all reduce，而 g 是前向傳播中的 all reduce 和后向傳播中的恒等式。

并行化多頭注意力層甚至更簡(jiǎn)單，因?yàn)樗鼈儽緛?lái)就是并行的，因?yàn)橛卸鄠€(gè)獨(dú)立的頭！


Self Attention TP

需要特別考慮的是：由于前向和后向傳播中每層都有兩個(gè) all reduce，因此 TP 需要設(shè)備間有非?？焖俚幕ヂ?lián)。因此，除非你有一個(gè)非?？斓木W(wǎng)絡(luò)，否則不建議跨多個(gè)節(jié)點(diǎn)進(jìn)行 TP。我們訓(xùn)練 BLOOM 的硬件配置中，節(jié)點(diǎn)間的速度比 PCIe 慢很多。實(shí)際上，如果節(jié)點(diǎn)有 4 個(gè) GPU，則最高 TP 度設(shè)為 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ì)信息。序列并行是在訓(xùn)練 BLOOM 之后開發(fā)的，所以 BLOOM 訓(xùn)練時(shí)并未采用此技術(shù)。

流水線并行

樸素流水線并行（naive PP）是將模型各層分組分布在多個(gè) GPU 上，并簡(jiǎn)單地將數(shù)據(jù)從 GPU 移動(dòng)到 GPU，就好像它是一個(gè)大型復(fù)合 GPU 一樣。該機(jī)制相對(duì)簡(jiǎn)單 - 將所需層用 .to () 方法綁到相應(yīng)設(shè)備，現(xiàn)在只要數(shù)據(jù)進(jìn)出這些層，這些層就會(huì)將數(shù)據(jù)切換到與該層相同的設(shè)備，其余部分保持不變。

這其實(shí)就是垂直模型并行，因?yàn)槿绻氵€記得我們是怎么畫大多數(shù)模型的拓?fù)鋱D的，我們其實(shí)是垂直切分模型各層的。例如，如果下圖顯示一個(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í)，這就跟單 GPU 上的普通前向傳播一樣。但是當(dāng)數(shù)據(jù)需要從第 3 層傳到第 4 層時(shí)，它需要從 GPU0 傳輸?shù)?GPU1，這會(huì)引入通信開銷。如果參與的 GPU 位于同一計(jì)算節(jié)點(diǎn)（例如同一臺(tái)物理機(jī)器）上，則傳輸非常快，但如果 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)化器來(lái)進(jìn)行更新參數(shù)了。

問(wèn)題：

  • 該方法為什么被稱為 樸素 流水線并行呢，它又有什么缺陷呢？主要是因?yàn)樵摲桨冈谌我饨o定時(shí)刻除了一個(gè) GPU 之外的其他所有 GPU 都是空閑的。因此，如果使用 4 個(gè) GPU，則幾乎等同于將單個(gè) GPU 的內(nèi)存量翻兩番，而其他資源（如計(jì)算）相當(dāng)于沒用上。另外還需要加上在設(shè)備之間復(fù)制數(shù)據(jù)的開銷。所以 4 張 使用樸素流水線并行的 6GB 卡將能夠容納與 1 張 24GB 卡相同大小的模型，而后者訓(xùn)練得更快，因?yàn)樗鼪]有數(shù)據(jù)傳輸開銷。但是，比如說(shuō)，如果你有 40GB 卡，但需要跑 45GB 模型，你可以使用 4x 40GB 卡（也就剛剛夠用，因?yàn)檫€有梯度和優(yōu)化器狀態(tài)需要顯存）。
  • 共享嵌入可能需要在 GPU 之間來(lái)回復(fù)制。

我們使用的流水線并行 (PP) 與上述樸素 PP 幾乎相同，但它解決了 GPU 閑置問(wèn)題，方法是將傳入的 batch 分塊為 micros batch 并人工創(chuàng)建流水線，從而允許不同的 GPU 同時(shí)參與計(jì)算過(guò)程。

下圖來(lái)自于 GPipe 論文 ，其上半部分表示樸素 PP 方案，下半部分是 PP 方法：

從圖的下半部分很容易看出 PP 的死區(qū)（指 GPU 處于空閑狀態(tài)）更少，即 “氣泡” 更少。

圖上兩種方案的并行度均為 4 ，即由 4 張 GPU 組成流水線。于是就有了 F0、F1、F2、F3 這 4 個(gè)管級(jí)的前向路徑，然后是 B3、B2、B1、B0 的逆序后向路徑。

PP 引入了一個(gè)新的超參數(shù)來(lái)調(diào)整，稱為 塊（chunks）。它定義了通過(guò)同一管級(jí)按順序發(fā)送多少數(shù)據(jù)塊。例如，在圖的下半部分，你可以看到 chunks = 4。 GPU0 在 chunk 0、1、2 和 3（F0,0、F0,1、F0,2、F0,3）上執(zhí)行相同的前向路徑，然后等待，等其他 GPU 完成工作后，GPU0 會(huì)再次開始工作，為塊 3、2、1 和 0（B0,3、B0,2、B0,1、B0,0）執(zhí)行后向路徑。

請(qǐng)注意，從概念上講，這與梯度累積 (gradient accumulation steps，GAS) 的意思相同。 PyTorch 叫它 塊，而 DeepSpeed 叫它 GAS。

因?yàn)?塊，PP 引入了 micro-batches（MBS）的概念。 DP 將全局 batch size 拆分為小 batch size，因此如果 DP 度為 4，則全局 batch size 1024 將拆分為 4 個(gè)小 batch size，每個(gè)小 batch size 為 256 (1024/4)。而如果 塊（或 GAS）的數(shù)量為 32，我們最終得到的 micro batch size 為 8 (256/32)。每個(gè)管級(jí)一次處理一個(gè) micro batch。

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

我們回過(guò)頭再看一下圖。

使用 chunks=1 你最終得到的是樸素 PP，這是非常低效的。而使用非常大的 塊 數(shù)，你最終會(huì)得到很小的微批量大小，這很可能也不是很有效。因此，必須通過(guò)實(shí)驗(yàn)來(lái)找到能最有效地利用 GPU 的塊數(shù)。

該圖顯示存在無(wú)法并行化的 “死” 時(shí)間氣泡，因?yàn)樽詈笠粋€(gè) forward 階段必須等待 backward 完成流水。那么，找到最佳的 塊 數(shù)，從而使所有參與的 GPU 達(dá)到高的并發(fā)利用率，這一問(wèn)題其實(shí)就轉(zhuǎn)化為最小化氣泡數(shù)了。

這種調(diào)度機(jī)制被稱為 全前全后。其他一些可選方案有 一前一后 和 交錯(cuò)一前一后。

雖然 Megatron-LM 和 DeepSpeed 都有自己的 PP 協(xié)議實(shí)現(xiàn)，但 Megatron-DeepSpeed 使用的是 DeepSpeed 實(shí)現(xiàn)，因?yàn)樗c DeepSpeed 的其他功能集成在一起。

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

DP+PP

DeepSpeed 流水線并行教程 中有一張圖演示了如何將 DP 與 PP 結(jié)合起來(lái)，如下所示。

2D Parallelism

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

由于每個(gè)維度至少需要 2 個(gè) GPU，因此這兒至少需要 4 個(gè) GPU。

DP+PP+TP

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

3D Parallelism

此圖來(lái)自博文 3D 并行：擴(kuò)展到萬(wàn)億參數(shù)模型), 這也是一篇好文章。

由于每個(gè)維度至少需要 2 個(gè) GPU，因此在這里你至少需要 8 個(gè) GPU 才能實(shí)現(xiàn)完整的 3D 并行。

ZeRO DP+PP+TP

DeepSpeed 的主要功能之一是 ZeRO，它是 DP 的超級(jí)可伸縮增強(qiáng)版，我們?cè)?ZeRO 數(shù)據(jù)并行 一節(jié)中已經(jīng)討論過(guò)了。通常它是一個(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 還會(huì)對(duì)梯度進(jìn)行分片，階段 3 也對(duì)模型權(quán)重進(jìn)行分片。

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

此外，由于 PP，層數(shù)已經(jīng)比正常情況下少，因此并不會(huì)節(jié)省很多內(nèi)存。 PP 已經(jīng)將梯度大小減少了 1/PP，因此在此基礎(chǔ)之上的梯度分片和純 DP 相比節(jié)省不了多少內(nèi)存。

ZeRO 階段 3 也可用于訓(xùn)練這種規(guī)模的模型，但是，它需要的通信量比 DeepSpeed 3D 并行更多。一年前，在對(duì)我們的環(huán)境進(jìn)行仔細(xì)評(píng)估后，我們發(fā)現(xiàn) Megatron-DeepSpeed 3D 并行性表現(xiàn)最佳。此后，ZeRO 階段 3 的性能有了顯著提高，如果我們今天要對(duì)其進(jìn)行重新評(píng)估，也許我們會(huì)選擇階段 3。

BF16 優(yōu)化器

用 FP16 訓(xùn)練巨型 LLM 模型是一個(gè)禁忌。

我們已經(jīng)通過(guò)花費(fèi)幾個(gè)月的時(shí)間 訓(xùn)練 104B 模型 自證了這一點(diǎn)，你可以從 tensorboard 發(fā)現(xiàn)，徹頭徹尾地失敗了。在與不斷發(fā)散的 lm-loss 作斗爭(zhēng)的過(guò)程中，我們學(xué)到了很多：

FP16 lm-los

我們也從 Megatron-LM 和 DeepSpeed 團(tuán)隊(duì)那里得到了相同的建議，在他們訓(xùn)得 530B 模型 后。最近發(fā)布的 OPT-175B 也報(bào)告說(shuō)他們?cè)?FP16 上訓(xùn)練得非常艱難。

所以早在一月份，我們就知道我們要在支持 BF16 格式的 A100 上進(jìn)行訓(xùn)練。 Olatunji Ruwase 開發(fā)了一個(gè)用來(lái)訓(xùn)練 BLOOM 的 “BF16Optimizer”。

如果您不熟悉這種數(shù)據(jù)格式，請(qǐng)查看 它的位布局。 BF16 格式的關(guān)鍵是它的指數(shù)位數(shù)與 FP32 相同，因此不會(huì)溢出，但 FP16 經(jīng)常溢出！FP16 的最大數(shù)值范圍為 64k，您只能進(jìn)行較小數(shù)的乘法。例如你可以做 250*250=62500，但如果你嘗試 255*255=65025，你就會(huì)溢出，這是導(dǎo)致訓(xùn)練出現(xiàn)問(wèn)題的主要原因。這意味著你的權(quán)重必須保持很小。一種稱為損失縮放（loss scaling）的技術(shù)有助于緩解這個(gè)問(wèn)題，但是當(dāng)模型變得非常大時(shí)，F(xiàn)P16 較小的數(shù)值范圍仍然是一個(gè)問(wèn)題。

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

當(dāng)然，由于 BF16 和 FP16 的大小相同，均為 2 個(gè)字節(jié)，因此，沒有免費(fèi)的午餐，當(dāng)使用 BF16 時(shí)，代價(jià)就是它的精度非常差。然而，你應(yīng)該還記得我們?cè)谟?xùn)練時(shí)采用的隨機(jī)梯度下降法及其變體，該方法有點(diǎn)像蹣跚而行，如果你這步?jīng)]有找到完美的方向其實(shí)沒關(guān)系，你會(huì)在接下來(lái)的步驟中糾正自己。

無(wú)論使用 BF16 還是 FP16，都有一個(gè)權(quán)重副本始終在 FP32 中 —— 這是由優(yōu)化器更新的內(nèi)容。因此 16 位格式僅用于計(jì)算，優(yōu)化器以全精度更新 FP32 權(quán)重，然后將它們轉(zhuǎn)換為 16 位格式以用于下一次迭代。

所有 PyTorch 組件都已更新，以確保它們?cè)?FP32 中執(zhí)行任何累加，因此不會(huì)發(fā)生精度損失。

一個(gè)關(guān)鍵問(wèn)題是梯度累積，它是流水線并行的主要特征之一，因?yàn)槊總€(gè) micro batch 處理的梯度都會(huì)累積。在 FP32 中實(shí)現(xiàn)梯度累積以保證訓(xùn)練的精確性至關(guān)重要，這正是 BF16Optimizer 所做的。

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

BF16 lm-loss

CUDA 融合核函數(shù)

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

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

現(xiàn)在，當(dāng)使用 GPU 計(jì)算 c = torch.add (a, b); e = torch.max ([c,d]) 時(shí)，一般情況下，PyTorch 將執(zhí)行的操作是啟動(dòng)兩個(gè)單獨(dú)的核函數(shù)，一個(gè)執(zhí)行 a 和 b 的加法，另一個(gè)執(zhí)行取 c 和 d 兩者的最大值。在這種情況下，GPU 從其顯存中獲取 a 和 b，執(zhí)行加法運(yùn)算，然后將結(jié)果寫回顯存。然后它獲取 c 和 d 并執(zhí)行 max 操作，然后再次將結(jié)果寫回顯存。

如果我們要融合這兩個(gè)操作，即將它們放入一個(gè) “融合核函數(shù)” 中，然后啟動(dòng)那個(gè)內(nèi)核，我們不會(huì)將中間結(jié)果 c 寫到顯存中，而是將其保留在 GPU 寄存器中，并且僅需要獲取 d 來(lái)完成最后的計(jì)算。這節(jié)省了大量開銷并防止 GPU 空閑，因此整個(gè)操作會(huì)更加高效。

融合核函數(shù)就是這樣。它們主要將多個(gè)離散的計(jì)算和進(jìn)出顯存的數(shù)據(jù)移動(dòng)替換為有很少數(shù)據(jù)移動(dòng)的融合計(jì)算。此外，一些融合核函數(shù)會(huì)對(duì)操作進(jìn)行數(shù)學(xué)變換，以便可以更快地執(zhí)行某些計(jì)算組合。

為了快速高效地訓(xùn)練 BLOOM，有必要使用 Megatron-LM 提供的幾個(gè)自定義 CUDA 融合核函數(shù)。特別地，有一個(gè) LayerNorm 的融合核函數(shù)以及用于融合縮放、掩碼和 softmax 這些操作的各種組合的核函數(shù)。Bias Add 也通過(guò) PyTorch 的 JIT 功能與 GeLU 融合。這些操作都是瓶頸在內(nèi)存的，因此將它們?nèi)诤显谝黄鹨赃_(dá)到最大化每次顯存讀取后的計(jì)算量非常重要。因此，例如，在執(zhí)行瓶頸在內(nèi)存的 GeLU 操作時(shí)同時(shí)執(zhí)行 Bias Add，運(yùn)行時(shí)間并不會(huì)增加。這些核函數(shù)都可以在 Megatron-LM 代碼庫(kù) 中找到。

數(shù)據(jù)集

Megatron-LM 的另一個(gè)重要特性是高效的數(shù)據(jù)加載器。在首次訓(xùn)練啟動(dòng)前，每個(gè)數(shù)據(jù)集中的每個(gè)樣本都被分成固定序列長(zhǎng)度（BLOOM 為 2048）的樣本，并創(chuàng)建索引以對(duì)每個(gè)樣本進(jìn)行編號(hào)?；谟?xùn)練超參，我們會(huì)確定每個(gè)數(shù)據(jù)集所需要參與的 epoch 數(shù)，并基于此創(chuàng)建一個(gè)有序的樣本索引列表，然后打亂它。舉個(gè)例子，如果一個(gè)數(shù)據(jù)集中有 10 個(gè)樣本并應(yīng)參與 2 個(gè) epoch 的訓(xùn)練，則系統(tǒng)首先按 [0, ..., 9, 0, ..., 9] 順序排好樣本索引，然后打亂該順序?yàn)閿?shù)據(jù)集創(chuàng)建最終的全局順序。請(qǐng)注意，這意味著訓(xùn)練不會(huì)簡(jiǎn)單地遍歷整個(gè)數(shù)據(jù)集然后重復(fù)，你有可能在看到另一個(gè)樣本之前看到同一個(gè)樣本兩次，但在訓(xùn)練結(jié)束時(shí)模型將只看到每個(gè)樣本兩次。這有助于確保整個(gè)訓(xùn)練過(guò)程中的訓(xùn)練曲線平滑。這些索引，包括每個(gè)樣本在原始數(shù)據(jù)集中的偏移量，被保存到一個(gè)文件中，以避免每次開始訓(xùn)練時(shí)都重新計(jì)算它們。最后，可以將其中幾個(gè)數(shù)據(jù)集以不同的權(quán)重混合到訓(xùn)練最終使用的數(shù)據(jù)中。

嵌入 LayerNorm

在我們努力阻止 104B 模型發(fā)散的過(guò)程中，我們發(fā)現(xiàn)在第一個(gè)層詞嵌入層之后添加一個(gè)額外的 LayerNorm 可以使訓(xùn)練更加穩(wěn)定。

該洞察來(lái)自對(duì) bitsandbytes 的實(shí)驗(yàn)，bitsandbytes 有一個(gè) StableEmbedding 操作，它是一個(gè)帶有 LayerNorm 的普通嵌入，其使用均勻 xavier 函數(shù)來(lái)初始化。

位置編碼

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

訓(xùn)練中的困難

隨著架構(gòu)、硬件和軟件的就位，我們得以在 2022 年 3 月上旬開始訓(xùn)練。然而，從那時(shí)起，事情其實(shí)并非一帆風(fēng)順。在本節(jié)中，我們將討論我們遇到的一些主要障礙。

在訓(xùn)練開始之前，有很多問(wèn)題需要弄清楚。特別是，我們發(fā)現(xiàn)了幾個(gè)問(wèn)題，這些問(wèn)題只有在我們開始在 48 個(gè)節(jié)點(diǎn)上進(jìn)行訓(xùn)練后才會(huì)出現(xiàn)，而不會(huì)在小規(guī)模時(shí)出現(xiàn)。例如，需要設(shè) CUDA_LAUNCH_BLOCKING=1 來(lái)防止框架掛起，我們需要將優(yōu)化器組分成更小的組，否則框架會(huì)再次掛起。你可以在 訓(xùn)前編年史 中詳細(xì)了解這些內(nèi)容。

訓(xùn)練期間遇到的主要問(wèn)題類型是硬件故障。由于這是一個(gè)擁有大約 400 個(gè) GPU 的新集群，平均每周我們會(huì)遇到 1-2 個(gè) GPU 故障。我們每 3 小時(shí)（100 次迭代）保存一個(gè)檢查點(diǎn)。因此，我們每周因硬件崩潰平均損失 1.5 小時(shí)的訓(xùn)練成果。 Jean Zay 系統(tǒng)管理員隨后將更換有故障的 GPU 并恢復(fù)節(jié)點(diǎn)。與此同時(shí)，我們有備用節(jié)點(diǎn)可供使用。

我們還遇到過(guò)多次導(dǎo)致 5-10 小時(shí)停機(jī)的各種其他問(wèn)題，其中一些與 PyTorch 中的死鎖錯(cuò)誤有關(guān)，另一些則是由于磁盤空間不足。如果您對(duì)具體細(xì)節(jié)有興趣，請(qǐng)參閱 訓(xùn)練編年史。

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

另一個(gè)問(wèn)題是 SLURM 并非設(shè)計(jì)為供一組人使用。 SLURM 作業(yè)由單個(gè)用戶擁有，如果他們不在身邊，則該組的其他成員無(wú)法對(duì)正在運(yùn)行的作業(yè)執(zhí)行任何操作。我們制定了一個(gè)終止方案，允許組中的其他用戶終止當(dāng)前進(jìn)程，而不需要啟動(dòng)該進(jìn)程的用戶在場(chǎng)。這在 90% 的問(wèn)題上都很有效。如果 SLURM 設(shè)計(jì)者讀到這篇文章，請(qǐng)?zhí)砑右粋€(gè) Unix 組的概念，這樣一個(gè) SLURM 作業(yè)就可以由一個(gè)組擁有。

由于訓(xùn)練是全天候 24/7 進(jìn)行的，我們需要有人隨叫隨到 - 但由于我們?cè)跉W洲和加拿大西海岸都有人，因此不需要有人攜帶傳呼機(jī)，我們能很好地互相備份。當(dāng)然，周末的訓(xùn)練也得有人看著。我們自動(dòng)化了大部分事情，包括自動(dòng)從硬件崩潰中恢復(fù)，但有時(shí)仍需要人工干預(yù)。

結(jié)論

訓(xùn)練中最困難和最緊張的部分是訓(xùn)練開始前的 2 個(gè)月。我們承受著盡快開始訓(xùn)練的巨大壓力，因?yàn)橘Y源分配的時(shí)間有限，我們直到最后一刻才接觸到 A100。所以這是一個(gè)非常困難的時(shí)期，考慮到 BF16Optimizer 是在最后一刻編寫出來(lái)的，我們需要調(diào)試它并修復(fù)各種 bug。正如上一節(jié)所述，我們發(fā)現(xiàn)了新問(wèn)題，這些問(wèn)題只有在我們開始在 48 個(gè)節(jié)點(diǎn)上進(jìn)行訓(xùn)練后才會(huì)出現(xiàn)，并且不會(huì)在小規(guī)模時(shí)出現(xiàn)。

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

總的來(lái)說(shuō)，這是一次超級(jí)緊張但回報(bào)頗豐的經(jīng)歷。

訓(xùn)練大型語(yǔ)言模型仍然是一項(xiàng)具有挑戰(zhàn)性的任務(wù)，但我們希望通過(guò)公開構(gòu)建和共享這項(xiàng)技術(shù)，其他人可以借鑒我們的經(jīng)驗(yàn)。

資源

重要鏈接

論文與文章

我們不可能在本文中詳細(xì)解釋所有內(nèi)容，因此如果此處介紹的技術(shù)激起你的好奇心，使你想了解更多信息，請(qǐng)閱讀以下論文：

Megatron-LM:

DeepSpeed:

Megatron-LM 和 Deepspeeed 聯(lián)合:

ALiBi:

BitsNBytes:

  • 8-bit Optimizers via Block-wise Quantization （我們使用了該論文中的嵌入 LaynerNorm，但是論文的其他部分及其技術(shù)也很妙，我們沒用 8 位優(yōu)化器的唯一原因是我們已經(jīng)使用 DeepSpeed-ZeRO 節(jié)省了優(yōu)化器內(nèi)存）。

博文致謝

非常感謝以下這些人，他們提出了很好的問(wèn)題并幫助提高了文章的可讀性（按字母序）：
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í)工程師，工作方向?yàn)?transformer-family 模型在各模態(tài)數(shù)據(jù)上的應(yīng)用及大規(guī)模模型的訓(xùn)練推理。

?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請(qǐng)聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成，瀏覽時(shí)請(qǐng)結(jié)合常識(shí)與多方信息審慎甄別。
平臺(tái)聲明：文章內(nèi)容（如有圖片或視頻亦包括在內(nèi)）由作者上傳并發(fā)布，文章內(nèi)容僅代表作者本人觀點(diǎn)，簡(jiǎn)書系信息發(fā)布平臺(tái)，僅提供信息存儲(chǔ)服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容