千億參數(shù)開源大模型 BLOOM 背后的技術(shù)
假設(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ù)要點(diǎn)，以促進(jìn)大家對大模型訓(xùn)練技術(shù)的討論。
BLOOM 文檔鏈接:
https://hf.co/bigscience/bloom
首先，我們要感謝促成或贊助我們這個小組最終完成了訓(xùn)練 1760 億參數(shù)模型這一驚人壯舉的公司、個人和團(tuán)體。
然后，我們開始討論硬件配置和主要技術(shù)組件。
以下是對本項目的簡要總結(jié):
人員組成
該項目由 Thomas Wolf (Hugging Face 聯(lián)合創(chuàng)始人兼 CSO) 發(fā)想，他敢于與大公司競爭，提出不僅要訓(xùn)練出立于世界上最大的多語言模型之林的模型，還要讓所有人都可以公開訪問訓(xùn)練結(jié)果，圓了大多數(shù)人的夢想。
本文主要關(guān)注模型訓(xùn)練的工程方面。BLOOM 背后的技術(shù)中最重要的部分是分享專業(yè)知識并幫助我們進(jìn)行編碼和訓(xùn)練的人員和公司。
我們主要需要感謝 6 個群體:
HuggingFace 的 BigScience 團(tuán)隊投入了六名以上的全職員工全程參與了訓(xùn)練的研究和運(yùn)行，他們還提供或報銷了 Jean Zay 計算機(jī)之外的所有基礎(chǔ)設(shè)施。
Microsoft DeepSpeed 團(tuán)隊，開發(fā)了 DeepSpeed，后來將其與 Megatron-LM 集成，其開發(fā)人員花費(fèi)數(shù)周時間研究項目需求，并在訓(xùn)練前和訓(xùn)練期間提供了許多很棒的實用經(jīng)驗建議。
NVIDIA Megatron-LM 團(tuán)隊開發(fā)了 Megatron-LM，他們非常樂于回答我們的大量問題并提供一流的使用建議。
IDRIS / GENCI 團(tuán)隊管理著 Jean Zay 超級計算機(jī)，他們?yōu)樵擁椖烤栀浟舜罅康乃懔蛷?qiáng)大的系統(tǒng)管理支持。
PyTorch 團(tuán)隊創(chuàng)建了一個超強(qiáng)的框架，其余軟件都基于該框架，并且在準(zhǔn)備訓(xùn)練期間非常支持我們，修復(fù)了多個 bug 并提高了我們所依賴的 PyTorch 組件的訓(xùn)練可用性。
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 非常相似，只是增加了一些改進(jìn)，本文稍后將對此進(jìn)行討論。
該模型是在 Jean Zay 上訓(xùn)練的，Jean Zay 是由 GENCI 管理的法國政府資助的超級計算機(jī)，安裝在法國國家科學(xué)研究中心 (CNRS) 的國家計算中心 IDRIS。訓(xùn)練所需的算力由 GENCI 慷慨捐贈給本項目 (捐贈號 2021-A0101012475)。
訓(xùn)練硬件:
GPU: 384 張 NVIDIA A100 80GB GPU (48 個節(jié)點(diǎn)) + 32 張備用 GPU
每個節(jié)點(diǎn) 8 張 GPU，4 條 NVLink 卡間互聯(lián)，4 條 OmniPath 鏈路
CPU: AMD EPYC 7543 32 核處理器
CPU 內(nèi)存: 每個節(jié)點(diǎn) 512GB
GPU 顯存: 每個節(jié)點(diǎn) 640GB
節(jié)點(diǎn)間連接: 使用 Omni-Path Architecture (OPA) 網(wǎng)卡，網(wǎng)絡(luò)拓?fù)錇闊o阻塞胖樹
NCCL - 通信網(wǎng)絡(luò): 一個完全專用的子網(wǎng)
磁盤 IO 網(wǎng)絡(luò): GPFS 與其他節(jié)點(diǎn)和用戶共享
Checkpoints:
主 checkpoints
https://hf.co/bigscience/bloom每個 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
https://openreview.net/forum?id=UoEw6KigkUn
176B BLOOM 模型的訓(xùn)練于 2022 年 3 月至 7 月期間，耗時約 3.5 個月完成 (約 100 萬計算時)。
Megatron-DeepSpeed
176B BLOOM 模型使用 Megatron-DeepSpeed 進(jìn)行訓(xùn)練，它結(jié)合了兩種主要技術(shù):
Megatron-DeepSpeed:
https://github.com/bigscience-workshop/Megatron-DeepSpeedDeepSpeed 是一個深度學(xué)習(xí)優(yōu)化庫，讓分布式訓(xùn)練變得簡單、高效且有效。
https://github.com/microsoft/DeepSpeedMegatron-LM 是由 NVIDIA 的應(yīng)用深度學(xué)習(xí)研究團(tuán)隊開發(fā)的大型、強(qiáng)大的 transformer 模型框架。
https://github.com/NVIDIA/Megatron-LM
DeepSpeed 團(tuán)隊通過將 DeepSpeed 庫中的 ZeRO 分片和流水線并行 (Pipeline Parallelism) 與 Megatron-LM 中的張量并行 (Tensor Parallelism) 相結(jié)合，開發(fā)了一種基于 3D 并行的方案。有關(guān)每個組件的更多詳細(xì)信息，請參見下表。
請注意，BigScience 的 Megatron-DeepSpeed 是基于原始 Megatron-DeepSpeed 代碼庫，我們還在其上添加了不少代碼。
BigScience 的 Megatron-DeepSpeed:
https://github.com/bigscience-workshop/Megatron-DeepSpeed原版的 Megatron-DeepSpeed:
https://github.com/microsoft/Megatron-DeepSpeed
下表列出了我們在訓(xùn)練 BLOOM 時各采用了兩個框架的哪些組件:
請注意，Megatron-LM 和 DeepSpeed 都有流水線并行和 BF16 優(yōu)化器實現(xiàn)，但我們使用 DeepSpeed 的實現(xiàn)，因為它們集成進(jìn)了 ZeRO。
Megatron-DeepSpeed 實現(xiàn)了 3D 并行以允許大模型以非常有效的方式進(jìn)行訓(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 的一部分?jǐn)?shù)據(jù)。
零冗余優(yōu)化器 (Zero Redundancy Optimizer，ZeRO)?- 也執(zhí)行與 TP 相類似的張量分片，但整個張量會及時重建以進(jìn)行前向或反向計算，因此不需要修改模型。它還支持各種卸載技術(shù)以補(bǔ)償有限的 GPU 內(nèi)存。
數(shù)據(jù)并行
大多數(shù)只有幾張 GPU 的用戶可能比較熟悉?DistributedDataParallel(DDP)，這是相應(yīng)的 PyTorch 文檔。在該方法中，模型被完全復(fù)制到每個 GPU，然后在每次迭代后所有模型相互同步各自的狀態(tài)。這種方法可以通過投入更多 GPU 資源的方式加快訓(xùn)練速度，解決問題。但它有個限制，即只有當(dāng)模型能夠放進(jìn)單個 GPU 時才有效。
PyTorch 文檔:
https://pytorch.org/docs/master/generated/torch.nn.parallel.DistributedDataParallel.html
ZeRO 數(shù)據(jù)并行
下圖很好地描述了 ZeRO 數(shù)據(jù)并行 (來自下列博文)。
https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-1000-billion-parameters/
看上去比較高大上，可能讓你很難專心去理解，但實際上，這個概念非常簡單。這只是通常的 DDP，只是沒有每個 GPU 都復(fù)制完整的模型參數(shù)、梯度和優(yōu)化器狀態(tài)，而是每個 GPU 只存儲其中的一部分。在隨后的運(yùn)行過程中，當(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 中的概念和圖表。
Efficient Large-Scale Language Model Training on GPU Clusters 論文:
https://arxiv.org/abs/2104.04473
Transformer 類模型的主要模塊為: 一個全連接層?nn.Linear，后面跟一個非線性激活層?GeLU。
沿用 Megatron 論文的符號，我們可以將其點(diǎn)積部分寫為?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?，它們可以獨(dú)立輸入?GeLU:
注意因為?Y?矩陣是按列拆分的，因此隨后的 GEMM 我們可以選擇按行拆分方案，這樣它就可以直接獲取前面層的 GeLU 的輸出，而無需任何額外的通信。
使用該原理，我們可以更新任意深度的 MLP，只需在每個?拆列 - 拆行?序列之后同步 GPU。Megatron-LM 論文作者為此提供了一個不錯的圖示:
這里?f?是前向傳播中的恒等運(yùn)算符，后向傳播中的 all reduce，而?g?是前向傳播中的 all reduce 和后向傳播中的恒等式。
并行化多頭注意力層甚至更簡單，因為它們本來就是并行的，因為有多個獨(dú)立的頭！
需要特別考慮的是: 由于前向和后向傳播中每層都有兩個 all reduce，因此 TP 需要設(shè)備間有非?？焖俚幕ヂ?lián)。因此，除非你有一個非?？斓木W(wǎng)絡(luò)，否則不建議跨多個節(jié)點(diǎn)進(jìn)行 TP。我們訓(xùn)練 BLOOM 的硬件配置中，節(jié)點(diǎn)間的速度比 PCIe 慢很多。實際上，如果節(jié)點(diǎn)有 4 個 GPU，則最高 TP 度設(shè)為 4 比較好。如果需要 TP 度為 8，則需要使用至少有 8 個 GPU 的節(jié)點(diǎn)。
該組件由 Megatron-LM 實現(xiàn)。Megatron-LM 最近擴(kuò)展了張量并行能力，新增了序列并行的能力，用于難以使用前述切分算法的算子，如 LayerNorm。Reducing Activation Recomputation in Large Transformer Models 論文提供了此技術(shù)的詳細(xì)信息。序列并行是在訓(xùn)練 BLOOM 之后開發(fā)的，所以 BLOOM 訓(xùn)練時并未采用此技術(shù)。
Reducing Activation Recomputation in Large Transformer Models 論文:
https://arxiv.org/abs/2205.05198
流水線并行
樸素流水線并行 (naive PP) 是將模型各層分組分布在多個 GPU 上，并簡單地將數(shù)據(jù)從 GPU 移動到 GPU，就好像它是一個大型復(fù)合 GPU 一樣。該機(jī)制相對簡單 - 將所需層用?.to()?方法綁到相應(yīng)設(shè)備，現(xiàn)在只要數(shù)據(jù)進(jìn)出這些層，這些層就會將數(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é)點(diǎn) (例如同一臺物理機(jī)器) 上，則傳輸非常快，但如果 GPU 位于不同的計算節(jié)點(diǎn) (例如多臺機(jī)器) 上，通信開銷可能會大得多。
然后第 4 到 5 到 6 到 7 層又像普通模型一樣，當(dāng)?shù)?7 層完成時，我們通常需要將數(shù)據(jù)發(fā)送回標(biāo)簽所在的第 0 層 (或者將標(biāo)簽發(fā)送到最后一層)?，F(xiàn)在可以計算損失，然后使用優(yōu)化器來進(jìn)行更新參數(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 方法:
GPipe 論文:
https://ai.googleblog.com/2019/03/introducing-gpipe-open-source-library.html
從圖的下半部分很容易看出 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)度機(jī)制被稱為?全前全后。其他一些可選方案有 一前一后 和 交錯一前一后。
一前一后方案:
https://www.microsoft.com/en-us/research/publication/pipedream-generalized-pipeline-parallelism-for-dnn-training/交錯一前一后方案:
https://arxiv.org/abs/2104.04473
雖然 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é)合起來，如下所示。
流水線并行教程:
https://www.deepspeed.ai/tutorials/pipeline/
這里重要的是要了解 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 并行: 擴(kuò)展到萬億參數(shù)模型》), 這也是一篇好文章。
《3D 并行: 擴(kuò)展到萬億參數(shù)模型》博文:
https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/
由于每個維度至少需要 2 個 GPU，因此在這里你至少需要 8 個 GPU 才能實現(xiàn)完整的 3D 并行。
ZeRO DP+PP+TP
DeepSpeed 的主要功能之一是 ZeRO，它是 DP 的超級可伸縮增強(qiáng)版，我們在 [ZeRO 數(shù)據(jù)并行](#ZeRO- 數(shù)據(jù)并行) 一節(jié)中已經(jīng)討論過了。通常它是一個獨(dú)立的功能，不需要 PP 或 TP。但它也可以與 PP、TP 結(jié)合使用。
當(dāng) ZeRO-DP 與 PP (以及 TP) 結(jié)合時，它通常只啟用 ZeRO 階段 1，它只對優(yōu)化器狀態(tài)進(jìn)行分片。ZeRO 階段 2 還會對梯度進(jìn)行分片，階段 3 也對模型權(quán)重進(jìn)行分片。
雖然理論上可以將 ZeRO 階段 2 與 流水線并行 一起使用，但它會對性能產(chǎn)生不良影響。每個 micro batch 都需要一個額外的 reduce-scatter 通信來在分片之前聚合梯度，這會增加潛在的顯著通信開銷。根據(jù)流水線并行的性質(zhì)，我們會使用小的 micro batch ，并把重點(diǎn)放在算術(shù)強(qiáng)度 (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)境進(jìn)行仔細(xì)評估后，我們發(fā)現(xiàn) Megatron-DeepSpeed 3D 并行性表現(xiàn)最佳。此后，ZeRO 階段 3 的性能有了顯著提高，如果我們今天要對其進(jìn)行重新評估，也許我們會選擇階段 3。
BF16 優(yōu)化器
用 FP16 訓(xùn)練巨型 LLM 模型是一個禁忌。
我們已經(jīng)通過花費(fèi)幾個月的時間訓(xùn)練 104B 模型自證了這一點(diǎn)，你可以從 Tensorboard 發(fā)現(xiàn)，徹頭徹尾地失敗了。在與不斷發(fā)散的 lm-loss 作斗爭的過程中，我們學(xué)到了很多:
訓(xùn)練 104B 模型代碼:
https://github.com/bigscience-workshop/bigscience/tree/master/train/tr8-104B-wideTensorboard 文檔:
https://hf.co/bigscience/tr8-104B-logs/tensorboard
我們也從 Megatron-LM 和 DeepSpeed 團(tuán)隊那里得到了相同的建議，在他們訓(xùn)得 530B 模型 后。最近發(fā)布的 OPT-175B 也報告說他們在 FP16 上訓(xùn)練得非常艱難。
530B 模型論文: https://arxiv.org/abs/2201.11990
OPT-175B 論文: https://arxiv.org/abs/2205.01068
所以早在一月份，我們就知道我們要在支持 BF16 格式的 A100 上進(jìn)行訓(xùn)練。Olatunji Ruwase 開發(fā)了一個用來訓(xùn)練 BLOOM 的 ?“BF16Optimizer”。
如果您不熟悉這種數(shù)據(jù)格式，請查看它的位布局。BF16 格式的關(guān)鍵是它的指數(shù)位數(shù)與 FP32 相同，因此不會溢出，但 FP16 經(jīng)常溢出！FP16 的最大數(shù)值范圍為 64k，您只能進(jìn)行較小數(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 格式的位布局:
https://en.wikipedia.org/wiki/Bfloat16_floating-point_format
BF16 沒有這個問題，你可以很容易地做?10_000*10_000=100_000_000, 完全沒問題。
當(dāng)然，由于 BF16 和 FP16 的大小相同，均為 2 個字節(jié)，因此，沒有免費(fèi)的午餐，當(dāng)使用 BF16 時，代價就是它的精度非常差。然而，你應(yīng)該還記得我們在訓(xùn)練時采用的隨機(jī)梯度下降法及其變體，該方法有點(diǎ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?所做的。
除了其他改進(jìn)之外，我們認(rèn)為使用 BF16 混合精度訓(xùn)練將潛在的噩夢變成了一個相對平穩(wěn)的過程，這可以從以下 lm 損失圖中看出:
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)該運(yùn)行哪些代碼，最后運(yùn)行它。CUDA 核函數(shù)使用 CUDA 來實現(xiàn)這些代碼，因此只能在 NVIDIA GPU 上運(yùn)行。
PyTorch 調(diào)度器:
http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/
現(xiàn)在，當(dāng)使用 GPU 計算?c = torch.add (a, b); e = torch.max ([c,d])?時，一般情況下，PyTorch 將執(zhí)行的操作是啟動兩個單獨(dú)的核函數(shù)，一個執(zhí)行?a?和 ?b?的加法，另一個執(zhí)行取?c?和 ?d?兩者的最大值。在這種情況下，GPU 從其顯存中獲取?a?和 ?b，執(zhí)行加法運(yùn)算，然后將結(jié)果寫回顯存。然后它獲取?c?和 ?d?并執(zhí)行?max?操作，然后再次將結(jié)果寫回顯存。
如果我們要融合這兩個操作，即將它們放入一個 “融合核函數(shù)” 中，然后啟動那個內(nèi)核，我們不會將中間結(jié)果?c?寫到顯存中，而是將其保留在 GPU 寄存器中，并且僅需要獲取?d?來完成最后的計算。這節(jié)省了大量開銷并防止 GPU 空閑，因此整個操作會更加高效。
融合核函數(shù)就是這樣。它們主要將多個離散的計算和進(jìn)出顯存的數(shù)據(jù)移動替換為有很少數(shù)據(jù)移動的融合計算。此外，一些融合核函數(shù)會對操作進(jìn)行數(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，運(yùn)行時間并不會增加。這些核函數(shù)都可以在 Megatron-LM 代碼庫 中找到。
數(shù)據(jù)集
Megatron-LM 的另一個重要特性是高效的數(shù)據(jù)加載器。在首次訓(xùn)練啟動前，每個數(shù)據(jù)集中的每個樣本都被分成固定序列長度 (BLOOM 為 2048) 的樣本，并創(chuàng)建索引以對每個樣本進(jìn)行編號。基于訓(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ù)來初始化。
bitsandbytes 代碼倉庫:
https://github.com/TimDettmers/bitsandbytes
位置編碼
基于論文 Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation，我們還用 AliBi 替換了普通的位置嵌入，它允許外推比訓(xùn)練模型的輸入序列更長的輸入序列。因此，即使我們訓(xùn)練時使用長度為 2048 的序列，模型也可以在推理過程中處理更長的序列。
Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation:
https://arxiv.org/abs/2108.12409
訓(xùn)練中的困難
隨著架構(gòu)、硬件和軟件的就位，我們得以在 2022 年 3 月上旬開始訓(xùn)練。然而，從那時起，事情其實并非一帆風(fēng)順。在本節(jié)中，我們將討論我們遇到的一些主要障礙。
在訓(xùn)練開始之前，有很多問題需要弄清楚。特別是，我們發(fā)現(xiàn)了幾個問題，這些問題只有在我們開始在 48 個節(jié)點(diǎn)上進(jìn)行訓(xùn)練后才會出現(xiàn)，而不會在小規(guī)模時出現(xiàn)。例如，需要設(shè)?CUDA_LAUNCH_BLOCKING=1?來防止框架掛起，我們需要將優(yōu)化器組分成更小的組，否則框架會再次掛起。你可以在 訓(xùn)前編年史 中詳細(xì)了解這些內(nèi)容。
訓(xùn)前編年史:
https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles-prequel.md
訓(xùn)練期間遇到的主要問題類型是硬件故障。由于這是一個擁有大約 400 個 GPU 的新集群，平均每周我們會遇到 1-2 個 GPU 故障。我們每 3 小時 (100 次迭代) 保存一個檢查點(diǎn)。因此，我們每周因硬件崩潰平均損失 1.5 小時的訓(xùn)練成果。Jean Zay 系統(tǒng)管理員隨后將更換有故障的 GPU 并恢復(fù)節(jié)點(diǎn)。與此同時，我們有備用節(jié)點(diǎn)可供使用。
我們還遇到過多次導(dǎo)致 5-10 小時停機(jī)的各種其他問題，其中一些與 PyTorch 中的死鎖錯誤有關(guān)，另一些則是由于磁盤空間不足。如果您對具體細(xì)節(jié)有興趣，請參閱 訓(xùn)練編年史。
訓(xùn)練編年史:
https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md
在對訓(xùn)練這個模型進(jìn)行可行性分析時，所有這些停機(jī)時間都被計劃在內(nèi)了，我們也據(jù)此選擇了合適的模型大小和我們希望模型消耗的數(shù)據(jù)量。因此，即使存在這些停機(jī)問題，我們還是成功地在預(yù)計時間內(nèi)完成了訓(xùn)練。如前所述，它需要大約 100 萬個計算時才能完成。
另一個問題是 SLURM 并非設(shè)計為供一組人使用。SLURM 作業(yè)由單個用戶擁有，如果他們不在身邊，則該組的其他成員無法對正在運(yùn)行的作業(yè)執(zhí)行任何操作。我們制定了一個終止方案，允許組中的其他用戶終止當(dāng)前進(jìn)程，而不需要啟動該進(jìn)程的用戶在場。這在 90% 的問題上都很有效。如果 SLURM 設(shè)計者讀到這篇文章，請?zhí)砑右粋€ Unix 組的概念，這樣一個 SLURM 作業(yè)就可以由一個組擁有。
由于訓(xùn)練是全天候 24/7 進(jìn)行的，我們需要有人隨叫隨到 - 但由于我們在歐洲和加拿大西海岸都有人，因此不需要有人攜帶傳呼機(jī)，我們能很好地互相備份。當(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é)點(diǎn)上進(jìn)行訓(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)練文檔
https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/README.mdTensorboard
https://huggingface.co/bigscience/tr11-176B-ml-logs/tensorboard訓(xùn)練用的 slurm 腳本
https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/tr11-176B-ml.slurm
論文與文章
我們不可能在本文中詳細(xì)解釋所有內(nèi)容，因此如果此處介紹的技術(shù)激起你的好奇心，使你想了解更多信息，請閱讀以下論文:
Megatron-LM:
Efficient Large-Scale Language Model Training on GPU Clusters.
https://arxiv.org/abs/2104.04473Reducing Activation Recomputation in Large Transformer Models
https://arxiv.org/abs/2205.05198
DeepSpeed:
ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
https://arxiv.org/abs/1910.02054ZeRO-Offload: Democratizing Billion-Scale Model Training
https://arxiv.org/abs/2101.06840ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
https://arxiv.org/abs/2104.07857DeepSpeed: Extreme-scale model training for everyone
https://www.microsoft.com/en-us/research/blog/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.
https://arxiv.org/abs/2201.11990
ALiBi:
Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation
https://arxiv.org/abs/2108.12409What Language Model to Train if You Have One Million GPU Hours? - 你會在那里找到最終使得我們選擇 ALiBi 的實驗。
https://openreview.net/forum?id=rI7BL3fHIZq
BitsNBytes:
8-bit Optimizers via Block-wise Quantization (我們使用了該論文中的嵌入 LaynerNorm，但是論文的其他部分及其技術(shù)也很妙，我們沒用 8 位優(yōu)化器的唯一原因是我們已經(jīng)使用 DeepSpeed-ZeRO 節(jié)省了優(yōu)化器內(nèi)存)。
https://arxiv.org/abs/2110.02861
博文致謝
非常感謝以下這些人，他們提出了很好的問題并幫助提高了文章的可讀性 (按字母序):
Britney Muller,
Douwe Kiela,
Jared Casper,
Jeff Rasley,
Julien Launay,
Leandro von Werra,
Omar Sanseviero,
Stefan Schweter and
Thomas Wang.
本文圖表主要由 Chunte Lee 創(chuàng)作。
英文原文:?https://hf.co/blog/bloom-megatron-deepspeed
原文作者: Stas Bekman
譯者: Matrix Yao (姚偉峰)，英特爾深度學(xué)習(xí)工程師，工作方向為 transformer-family 模型在各模態(tài)數(shù)據(jù)上的應(yīng)用及大規(guī)模模型的訓(xùn)練推理。
審校、排版: zhongdongy (阿東)
