原文：語(yǔ)言大模型的分布式訓(xùn)練與高效微調(diào)指南 - 知乎

目錄

收起

1 分布式訓(xùn)練

2 ZeRO驅(qū)動(dòng)的數(shù)據(jù)并行

3 全分片數(shù)據(jù)并行

4 實(shí)現(xiàn)

5 高效微調(diào)

6 實(shí)踐指南

7 關(guān)于DeepSpeed和FSDP的更多內(nèi)容

OneFlow編譯
翻譯｜楊婷、宛子琳

最近語(yǔ)言大模型（LLM）異?；鸨粋€(gè)非常特別的開源社區(qū)正在探索在消費(fèi)級(jí)硬件上微調(diào)、提供服務(wù)和進(jìn)行推理的最佳方式。為滿足上述需求，出現(xiàn)了許多出色的開源代碼庫(kù)，以HuggingFace生態(tài)系統(tǒng)為中心，這些代碼庫(kù)還包括FastChat、Axolotl和LLama.cpp。

本文專注于分布式訓(xùn)練策略的具體細(xì)節(jié)，特別是DeepSpeed和FSDP，并總結(jié)了以多GPU和多節(jié)點(diǎn)訓(xùn)練為主的不同高效微調(diào)方法。顯然，當(dāng)前的趨勢(shì)是，我們會(huì)使用越來(lái)越多的計(jì)算資源，因此將需要更多GPU來(lái)運(yùn)行更大的模型。

在這種情況下，理解這些主題尤為重要，尤其是當(dāng)你想要將幾個(gè)3090家庭服務(wù)器升級(jí)到具有8個(gè)A100 80GB的GCP容器時(shí)，另外，這對(duì)于試圖微調(diào)自己的語(yǔ)言模型的初創(chuàng)公司或企業(yè)來(lái)說(shuō)也很重要。大型科技公司進(jìn)行的實(shí)際大規(guī)模訓(xùn)練涉及大量資料，這些內(nèi)容大部分來(lái)自主導(dǎo)了BLOOM-176B訓(xùn)練的Stas Bekman，GPU匱乏的用戶（即GPU-poors）關(guān)注這些資料的意義并不大。

本文從多個(gè)優(yōu)秀的資源中整合了各種觀點(diǎn)，著重討論了HuggingFace生態(tài)系統(tǒng)的相關(guān)內(nèi)容，并考慮了一些來(lái)自在線資源以及作者本人在2023年暑期實(shí)習(xí)中所學(xué)到的實(shí)際情況。

綜上，本文希望能夠回答以下問題：

1. 在分布式訓(xùn)練和性能優(yōu)化方面，我們應(yīng)該關(guān)注什么？DeepSpeed和FSDP在背后是如何運(yùn)作的？2. 不同的分布式訓(xùn)練策略需要的硬件設(shè)置和注意事項(xiàng)？

3. 各種有效的微調(diào)優(yōu)化器以及可能存在的權(quán)衡？

4. 一些可以覆蓋所有重要微調(diào)優(yōu)化的實(shí)用指南，用以在多GPU和多節(jié)點(diǎn)設(shè)置中訓(xùn)練大模型。

5. 現(xiàn)在可以使用的開源代碼庫(kù)以及各自的優(yōu)缺點(diǎn)？

（本文作者為加州大學(xué)圣地亞哥分校計(jì)算機(jī)科學(xué)系的碩士研究生Sumanth R Hegde。以下內(nèi)容由OneFlow編譯發(fā)布，轉(zhuǎn)載請(qǐng)聯(lián)系授權(quán)。原文：https://sumanthrh.com/post/distributed-and-efficient-finetuning/?）

1 分布式訓(xùn)練

分布式訓(xùn)練涵蓋的范圍非常廣泛，因此本文無(wú)法覆蓋所有內(nèi)容。在訓(xùn)練/微調(diào)LLM時(shí)，我們通常會(huì)面對(duì)超10億參數(shù)的龐大模型和大規(guī)模數(shù)據(jù)集（超1萬(wàn)億詞元的預(yù)訓(xùn)練數(shù)據(jù)集，超100萬(wàn)詞元的監(jiān)督微調(diào)數(shù)據(jù)集）。我們最終的目標(biāo)是盡快完成訓(xùn)練，以最大化吞吐量，即希望能夠每秒處理盡可能多的樣本。LLM在訓(xùn)練過(guò)程中需要大量的GPU顯存，不僅僅因?yàn)槟Ｐ蛥?shù)數(shù)量龐大（例如，F(xiàn)alcon 40B擁有40億參數(shù)，在BF16格式下僅模型權(quán)重就需要約74GB顯存），還因?yàn)閮?yōu)化器狀態(tài)所需的顯存——例如，使用普通的AdamW優(yōu)化器，每個(gè)參數(shù)需要12個(gè)字節(jié)來(lái)存儲(chǔ)模型權(quán)重的副本、動(dòng)量和方差參數(shù)。因此，我們需要智能的分布式訓(xùn)練策略，以確保每個(gè)GPU worker只需處理部分訓(xùn)練狀態(tài)和數(shù)據(jù)。

主要的并行策略有：

1. 數(shù)據(jù)并行（Data Parallelism, DP）：每個(gè)GPU worker獲取整個(gè)小批量數(shù)據(jù)的一部分，并在該部分?jǐn)?shù)據(jù)上計(jì)算梯度。然后在所有worker上對(duì)梯度進(jìn)行平均，以更新模型權(quán)重。在其最基本的形式中，比如PyTorch中的DDP，每個(gè)GPU存儲(chǔ)模型權(quán)重的副本、優(yōu)化器狀態(tài)以及所處理數(shù)據(jù)部分的梯度。

2. 模型并行/垂直模型并行（MP）：在模型并行中，模型被分割成多個(gè)部分，每個(gè)部分被放置在不同的GPU上，這被稱為垂直模型并行。舉例來(lái)說(shuō)，如果有一個(gè)包含12層的單一模型，該模型的不同層會(huì)被分別放置在3個(gè)不同的GPU上。

---------------  ---------------  -----------------
1 | 2 | 3 | 4 |  5 | 6 | 7 | 8 |  9 | 10 | 11 | 12 |
---------------  ---------------  -----------------

在樸素模型并行（naive model parallelism）中，所有GPU在處理相同的數(shù)據(jù)批次時(shí)都需要等待前一個(gè)GPU完成計(jì)算，然后才能處理數(shù)據(jù)。這意味著在任何給定時(shí)刻，除了一個(gè)GPU外，其他GPU實(shí)際上都處于閑置狀態(tài)（因此稱為“樸素”）。為改善這種情況，可以使用流水線并行（PP），這種方式通過(guò)讓不同微批次數(shù)據(jù)的計(jì)算重疊，給你帶來(lái)并行的錯(cuò)覺。這類似于計(jì)算機(jī)架構(gòu)中的經(jīng)典流水線。參見以下關(guān)于GPipe的博文：

為了在多個(gè)加速器上實(shí)現(xiàn)高效訓(xùn)練，GPipe將模型劃分到不同加速器，并自動(dòng)將一個(gè)小批次的訓(xùn)練樣本分成更小的微批次。通過(guò)在這些微批次上進(jìn)行流水線式的執(zhí)行，各個(gè)加速器可以進(jìn)行并行計(jì)算。

3. 張量并行（TP)：在張量并行中，每個(gè)GPU通過(guò)在GPU worker之間對(duì)模型進(jìn)行水平切片，僅處理張量的一部分。每個(gè)worker處理相同的數(shù)據(jù)批次，計(jì)算他們所擁有權(quán)重部分的激活值，并交換彼此需要的部分，每個(gè)worker計(jì)算他們所擁有權(quán)重部分的梯度。

我們可以將上述各種并行策略結(jié)合起來(lái)，以實(shí)現(xiàn)更好的吞吐量增益。接下來(lái)，我們將更詳細(xì)地了解兩種用于數(shù)據(jù)并行訓(xùn)練的改進(jìn)方法：零冗余優(yōu)化器（Zero Redundancy Optimizer）和密切相關(guān)的全切片數(shù)據(jù)并行策略（Fully Sharded Data-Parallel strategies）。

注釋：我將使用術(shù)語(yǔ)“GPU worker”來(lái)指代在每個(gè)GPU上運(yùn)行的各個(gè)進(jìn)程。雖然這種表述并不十分準(zhǔn)確，但在數(shù)據(jù)并行設(shè)置中，這樣說(shuō)更方便，更易于理解。

拓展閱讀：

• 多個(gè)GPU上的高效訓(xùn)練：https://huggingface.co/docs/transformers/perf_train_gpu_many
• 如何在多個(gè)GPU上訓(xùn)練大模型？
  https://lilianweng.github.io/posts/2021-09-25-train-large/

2 ZeRO驅(qū)動(dòng)的數(shù)據(jù)并行

這是當(dāng)前最高效、最熱門的分布式訓(xùn)練策略之一。DeepSpeed的ZeRO是一種數(shù)據(jù)并行處理形式，它極大提高了內(nèi)存效率，其主要思想是利用數(shù)據(jù)并行訓(xùn)練中的內(nèi)存冗余和快速GPU間通信的最新改進(jìn)來(lái)提高吞吐量，根據(jù)不同階段，會(huì)增加一些通信量。實(shí)際上，ZeRO有兩個(gè)組成部分：ZeRO-DP（數(shù)據(jù)并行）和ZeRO-R（殘留內(nèi)存）。DeepSpeed團(tuán)隊(duì)還提出了一些后續(xù)優(yōu)化措施，這進(jìn)一步提升了ZeRO的吸引力，例如ZeRO-Offload/Infinity（將計(jì)算卸載到CPU/NVMe磁盤）和ZeRO++（實(shí)現(xiàn)了靈活的多節(jié)點(diǎn)訓(xùn)練和量化權(quán)重）。

ZeRO-DP的可視化圖表如下（來(lái)自DeepSpeed的博客文章）：

在64個(gè)GPU上訓(xùn)練7.5B參數(shù)模型時(shí)，模型表現(xiàn)如下：

1. 基準(zhǔn)：PyTorch DDP

2.?ZeRO 第一階段/?????：內(nèi)存減少4倍（特定示例），與基準(zhǔn)的通信量相同（沒有額外的GPU間通信）。

3.?ZeRO 第二階段/????+??：內(nèi)存減少8倍（特定示例），與基準(zhǔn)的通信量相同。

4.?ZeRO 第三階段/????+??：內(nèi)存減少64倍（特定示例），通信量為基準(zhǔn)的1.5倍（這里的1.5倍針對(duì)的是不同硬件設(shè)置和模型大?。?/p>

基準(zhǔn)

PyTorch DDP實(shí)現(xiàn)了簡(jiǎn)單的數(shù)據(jù)并行性。每個(gè)GPU worker都有模型權(quán)重、優(yōu)化器狀態(tài)和梯度的副本。后向傳播后，各個(gè)worker間的梯度會(huì)被平均化（即全局歸約（all-reduce）），并且模型權(quán)重會(huì)被更新。

關(guān)于通信量的注釋：為理解ZeRO的好處，我認(rèn)為，重要的是要明確通訊量的確切含義。典型DP中存在一個(gè)全局歸約步驟，其中每個(gè)worker都會(huì)發(fā)送出其所擁有的梯度數(shù)組，然后接收其他worker的梯度數(shù)組，以獲得平均值。下文摘自ZeRO論文：

當(dāng)前最先進(jìn)的全局歸約實(shí)現(xiàn)采用了兩步操作：第一步reduce-scatter，它在不同進(jìn)程上的數(shù)據(jù)的不同部分會(huì)被歸約。第二步是all-gather，其中每個(gè)進(jìn)程都在這一步中收集所有進(jìn)程上的歸約數(shù)據(jù)。這兩步的結(jié)果就是全局歸約。同時(shí)，reduce-scatter和all-gather都采用了流水線方式實(shí)現(xiàn)，這導(dǎo)致每個(gè)步驟需要Ψ個(gè)元素（對(duì)于具有Ψ個(gè)元素的數(shù)據(jù)）的總數(shù)據(jù)搬運(yùn)量。因此，在標(biāo)準(zhǔn)的數(shù)據(jù)并行性中，每個(gè)訓(xùn)練步驟會(huì)產(chǎn)生2Ψ的數(shù)據(jù)搬運(yùn)量。

在這種情況下，“數(shù)據(jù)”指的是我們的梯度，而進(jìn)程指的是每個(gè)GPU上運(yùn)行的單個(gè)worker?？偠灾?，我想表達(dá)的是：如果你有Ψ個(gè)參數(shù)，那么普通數(shù)據(jù)并行性將產(chǎn)生2Ψ通信成本。

ZeRO 第一階段/Pos（優(yōu)化器狀態(tài)分區(qū)）

在這種情況下，只有優(yōu)化器狀態(tài)在GPU worker之間進(jìn)行了分區(qū)/分片，而模型權(quán)重和梯度在所有worker之間進(jìn)行了復(fù)制。反向傳播后，需要進(jìn)行一次常規(guī)全局歸約，以便在所有worker間獲取平均梯度值。然后，每個(gè)worker更新其分區(qū)中的優(yōu)化器狀態(tài)。Adam方程如下。

(w、g、v 和m 分別對(duì)應(yīng)權(quán)重、梯度、速度和動(dòng)量)。需要注意的是，這些都是逐元素操作，在計(jì)算梯度后，各個(gè)權(quán)重分片之間并沒有依賴關(guān)系。

通信量：首先我們會(huì)執(zhí)行一次全局歸約操作，將更新后的梯度通信給所有GPU，然后，在更新各自分片的優(yōu)化器狀態(tài)之后，每個(gè)GPU仍需要從其他GPU獲取更新后的權(quán)重。ZeRO論文并沒有清楚表明這樣做是否會(huì)增加通信量。在我看來(lái)，對(duì)ZeRO的第1階段和第2階段來(lái)說(shuō)，這種實(shí)現(xiàn)實(shí)際上是一樣的：

• 梯度的全局歸約由兩部分組成：reduce-scatter和all-gather。
• 在DeepSpeed ZeRO1和2中，首先進(jìn)行reduce-scatter操作，在不同GPU上分別減少梯度的不同部分。接著，每個(gè)GPU計(jì)算其所管理的優(yōu)化器分區(qū)對(duì)應(yīng)的更新后權(quán)重，之后我們只需進(jìn)行一次all-gather操作，將更新后的模型參數(shù)傳播給所有其他GPU。這樣就能在通信量不變的情況下，減少內(nèi)存消耗。

ZeRO Stage 2 /????+?

（優(yōu)化器狀態(tài)+梯度分區(qū)）

在這種情況下，優(yōu)化器狀態(tài)和梯度都被分區(qū)/分片到不同的worker上。這意味著，兩個(gè)GPU worker在訓(xùn)練期間不僅要關(guān)注不同的微批次數(shù)據(jù)，還要維護(hù)模型參數(shù)子集的梯度。關(guān)鍵在于，每個(gè)worker都在更新其優(yōu)化器狀態(tài)的分區(qū)，因此對(duì)于一個(gè)worker而言，它所需的梯度（或者說(shuō)，經(jīng)過(guò)歸約/平均的梯度）只是對(duì)應(yīng)于該狀態(tài)分區(qū)的梯度。至于實(shí)現(xiàn)方式，正如上面提到的，DeepSpeed有效地執(zhí)行了reduce-scatter操作，其中每個(gè)worker對(duì)應(yīng)的梯度在該worker處被平均（而不是對(duì)所有參數(shù)進(jìn)行典型的all-reduce操作）。這意味著，在相同通信量下節(jié)省了更多內(nèi)存，也就是說(shuō)，與DDP相比，這里沒有額外的數(shù)據(jù)搬運(yùn)成本。

注意：在使用ZeRO Stage 1和2時(shí)，仍然需要整個(gè)模型適配單個(gè)GPU。此外，在使用RAM時(shí)，需要注意如下事項(xiàng)：隨著進(jìn)程/GPU數(shù)量的增加以及模型大小的擴(kuò)展（超過(guò)40億參數(shù)），模型初始化會(huì)占用大量RAM。ZeRO 3對(duì)此有所改進(jìn)。

ZeRO Stage 3 /????+?+?

（優(yōu)化器狀態(tài)+梯度+參數(shù)分區(qū)）

對(duì)我來(lái)說(shuō)，這是ZeRO最有趣的階段。除優(yōu)化器狀態(tài)和梯度外，第3階段還跨worker劃分了模型參數(shù)。來(lái)自DeepSpeed的可視化圖表如下：

使用4個(gè)GPU實(shí)現(xiàn)不同訓(xùn)練狀態(tài)的可視化

使用DeepSpeed ZeRO-3對(duì)訓(xùn)練狀態(tài)分片

借用Stas Bekman指南中的一個(gè)例子，假設(shè)有以下3層模型和4個(gè)GPU：

?La | Lb | Lc
---|----|---
a0 | b0 | c0
a1 | b1 | c1
a2 | b2 | c2
a3 | b3 | c3

使用DeepSpeed ZeRO 3，GPU的配置方式如下：

GPU 0:
La | Lb | Lc
---|----|---
a0 | b0 | c0
?
GPU 1:
La | Lb | Lc
---|----|---
a1 | b1 | c1
?
GPU 2:
La | Lb | Lc
---|----|---
a2 | b2 | c2
?
GPU 3:
La | Lb | Lc
---|----|---
a3 | b3 | c3

在ZeRO-3中，模型的每一層都被水平切片，每個(gè)worker存儲(chǔ)權(quán)重張量的一部分。在前向和后向傳播過(guò)程中（每個(gè)GPU worker 仍然看到不同的微批次數(shù)據(jù)），不同的GPU worker交換它們所擁有的每一層的部分（按需進(jìn)行參數(shù)通信），并計(jì)算激活/梯度。

其余部分類似ZeRO Stage 2。很容易看出，ZeRO-3的通信量是基準(zhǔn)DDP的1.5倍：在每個(gè)訓(xùn)練步驟中，我們需要在前向傳播中額外進(jìn)行一次模型參數(shù)的all-gather操作。在這個(gè)操作中移動(dòng)的數(shù)據(jù)量為Ψ（每個(gè)GPU），因此總通信量為Ψ（參數(shù)all-gather）+Ψ（梯度reduce-scatter）+Ψ（all-gather，用于更新的參數(shù)）= 3Ψ= 1.5倍DDP。考慮到內(nèi)存消耗被GPU workerN 削減，這給人留下了深刻印象。ZeRO論文的另一個(gè)關(guān)鍵洞察如下：

只要有足夠的設(shè)備來(lái)共享模型狀態(tài)，ZeRO就可以使DP適應(yīng)任意大小的模型。

也就是說(shuō)，只要有足夠多的GPU，在進(jìn)行數(shù)據(jù)并行（DP）訓(xùn)練時(shí)，我們就不會(huì)再受到每個(gè)GPU顯存的限制（說(shuō)起來(lái)容易做起來(lái)難）。

ZeRO-R

我并不想深入探討這一話題，但ZeRO-R在ZeRO-DP的基礎(chǔ)上，通過(guò)關(guān)注激活的內(nèi)存消耗和管理內(nèi)存碎片化做了改進(jìn)提升。ZeRO-R通過(guò)對(duì)激活進(jìn)行分區(qū)，減少了激活的內(nèi)存占用。此外，它還在管理臨時(shí)緩沖區(qū)方面進(jìn)行了一些改進(jìn)，你可以將之視為在worker間進(jìn)行梯度累積和歸約期間分配用于存儲(chǔ)中間結(jié)果的內(nèi)存。

ZeRO-Offload

ZeRO-Offload是一種優(yōu)化技術(shù)，可以將優(yōu)化器和計(jì)算從GPU卸載到主機(jī)CPU上。在2021年1月發(fā)布時(shí)，ZeRO-Offload可在1個(gè)NVIDIA V100 GPU上實(shí)現(xiàn)40 TFLOPS（V100 32 GB vRAM，最大吞吐量為130 TFLOPS），適用于10億參數(shù)模型。而采用PyTorch DDP時(shí)，最大值為30 TFLOPS，適用于14億參數(shù)模型，在不耗盡內(nèi)存的情況下，這是可以運(yùn)行的最大模型。將計(jì)算卸載到CPU的主要問題是，以吞吐量衡量，CPU要比GPU慢上多個(gè)數(shù)量級(jí)。ZeRO-Offload采用的這種策略是，只將較少的密集計(jì)算（<< O(MB)，其中M 是模型大小，B 是批次大?。┬遁d到CPU，以使總計(jì)算復(fù)雜度保持不變（O(MB)）。在實(shí)踐中，這意味著諸如范數(shù)計(jì)算（norm calculation）、權(quán)重更新等操作可以在CPU上完成，而前向和后向傳播的矩陣乘法需要在GPU上完成。ZeRO-Offload適用于ZeRO的所有階段（1、2和3）。

這里（https://docs.it4i.cz/dgx2/introduction/）給出了用于實(shí)驗(yàn)的DGX-2節(jié)點(diǎn)規(guī)格，該節(jié)點(diǎn)配備了16個(gè)V100 GPU。需要注意的是，如果處于ZeRO-2設(shè)置中，ZeRO-Offload仍然會(huì)受到每個(gè)GPU可用內(nèi)存的限制，即在每個(gè)GPU上容納整個(gè)模型可能成為瓶頸。

ZeRO-Infinity

ZeRO-Infinity是ZeRO-Offload的改進(jìn)版本，于2021年4月推出。它允許將計(jì)算卸載到磁盤（NVMe內(nèi)存），并對(duì)CPU卸載進(jìn)行了改進(jìn)。研究表明，ZeRO-Infinity在一個(gè)DGX-2節(jié)點(diǎn)上訓(xùn)練時(shí)，可以適應(yīng)10-100萬(wàn)億參數(shù)的模型（萬(wàn)億級(jí)別?。?。ZeRO-Infinity通過(guò)同時(shí)利用CPU和NVMe內(nèi)存來(lái)實(shí)現(xiàn)這一點(diǎn)。以下是論文中的可視化圖表：

以上是ZeRO Infinity在4個(gè)數(shù)據(jù)并行排布（GPU）上的截圖，描述了反向傳播期間的狀態(tài)。分區(qū)/分片的參數(shù)從慢速內(nèi)存（CPU+ NVMe）移動(dòng)到GPU，然后被收集起來(lái)形成完整的層。在梯度計(jì)算之后，它們被聚合、重新分區(qū)，然后卸載到慢速內(nèi)存。分區(qū)/分片的參數(shù)從慢速內(nèi)存（CPU+ NVMe）移動(dòng)到GPU，然后被收集起來(lái)形成完整的層。在梯度計(jì)算之后，它們被聚合、重新分區(qū)，然后卸載到慢速內(nèi)存。

與ZeRO-Offload不同，ZeRO-Infinity是基于ZeRO-3專門構(gòu)建的。作者在32個(gè)DGX-2節(jié)點(diǎn)、512個(gè)GPU上對(duì)模型速度進(jìn)行了評(píng)估。結(jié)果表明，ZeRO-Infinity可訓(xùn)練高達(dá)20萬(wàn)億參數(shù)的模型，每個(gè)GPU的吞吐量可達(dá)49 TFlops，而3D并行等替代并行策略可訓(xùn)練的模型規(guī)模要小40倍。ZeRO-Infinity要想成為有競(jìng)爭(zhēng)力的選擇，需要滿足一些帶寬要求，即NVMe-CPU和CPU-GPU通信。

關(guān)于ZeRO-Offload和ZeRO-Infinity之間的差異，以下是DeepSpeed團(tuán)隊(duì)的評(píng)論：

DeepSpeed首先包含了ZeRO-Offload的卸載功能，ZeRO-Offload是一個(gè)用于將優(yōu)化器和梯度狀態(tài)卸載到ZeRO-2內(nèi)的CPU內(nèi)存的系統(tǒng)。ZeRO-Infinity是ZeRO-3可以使用的下一代卸載功能。相比ZeRO-Offload，ZeRO-Infinity能夠處理更多數(shù)據(jù)，能更有效地利用帶寬，并且能更好地實(shí)現(xiàn)計(jì)算和通信重疊。

默認(rèn)情況下，使用ZeRO-3進(jìn)行卸載時(shí)，ZeRO-Infinity的優(yōu)化機(jī)制會(huì)自動(dòng)生效；使用Stage 1/2進(jìn)行卸載時(shí)，ZeRO-Offload會(huì)生效。

拓展閱讀：1. ZeRO-Offload/Infinity 教程:https://www.deepspeed.ai/tutorials/zero-offload/

2. ZeRO-Offload - 十億級(jí)模型訓(xùn)練大眾化 :https://arxiv.org/abs/2101.06840

3. ZeRO-Infinity - 打破GPU內(nèi)存墻，實(shí)現(xiàn)超大規(guī)模深度學(xué)習(xí):https://arxiv.org/abs/2104.07857

ZeRO++

ZeRO++是DeepSpeed團(tuán)隊(duì)對(duì)ZeRO-3的最新改進(jìn)。主要改進(jìn)如下：

1. 量化權(quán)重（qwZ）：通過(guò)將模型權(quán)重量化為int8，將all-gather參數(shù)通信量減少一半。

2. 層次劃分 (hpZ)：層次劃分是一種混合分區(qū)方案，適用于DeepSpeed ZeRO 3的多節(jié)點(diǎn)設(shè)置。在這種情況下，模型參數(shù)會(huì)在一個(gè)節(jié)點(diǎn)內(nèi)分片，然后在不同節(jié)點(diǎn)之間復(fù)制。這意味著，與典型的ZeRO-3全面運(yùn)行相比，我們無(wú)法節(jié)省相同數(shù)量的內(nèi)存，但可以避免昂貴的節(jié)點(diǎn)間參數(shù)通信開銷，從而提高整體吞吐量。我更傾向于FSDP中使用的 "混合分片（hybrid sharding）" ，而非 "層次劃分"，下文討論FSDP時(shí)，會(huì)再次深入討論這個(gè)問題。

3. 量化梯度（qgZ）：通過(guò)在梯度reduce-scatter操作中使用int4量化數(shù)據(jù)替換fp16，可以節(jié)省更大通信量。（回顧：這是發(fā)生在ZeRO 2/3階段，針對(duì)分片梯度進(jìn)行的梯度聚集和平均）

總體而言，與ZeRO-3 相比，ZeRO++通過(guò)這三項(xiàng)改進(jìn)將通信量減少了4倍。

拓展閱讀：

1. ZeRO：訓(xùn)練萬(wàn)億參數(shù)模型的內(nèi)存優(yōu)化:https://arxiv.org/abs/1910.02054

2. ZeRO教程:https://www.deepspeed.ai/tutorials/zero/

3. ZeRO++：針對(duì)巨型模型訓(xùn)練的超高效集體通信:https://arxiv.org/abs/2306.10209

4. ZeRO++教程:https://www.deepspeed.ai/tutorials/zeropp/

3 全分片數(shù)據(jù)并行

全分片數(shù)據(jù)并行（FSDP）是另一種旨在提高內(nèi)存效率、減少通信計(jì)算開銷以提高吞吐量的數(shù)據(jù)并行技術(shù)。FSDP的分片策略基于Xu等人和ZeRO的想法。FSDP有兩種分片策略：全分片（Full Sharding）和混合分片（Hybrid Sharding）。

全分片

這基本上與ZeRO-3相同，其中參數(shù)、優(yōu)化器狀態(tài)和梯度被分片到各個(gè)worker或設(shè)備上。下圖來(lái)自FSDP博客，展示了兩個(gè)設(shè)備之間不同操作的低層次可視化。

正如你所看到的，每個(gè)worker/設(shè)備僅持有權(quán)重的子集，并且可以按需進(jìn)行參數(shù)通信來(lái)計(jì)算中間激活和梯度。以下是來(lái)自PyTorch的FSDP教程：

“在前向傳播中運(yùn)行all_gather從所有rank中收集全部分片，以恢復(fù)該FSDP單元中的完整參數(shù)運(yùn)行前向計(jì)算丟棄剛剛收集的參數(shù)分片

在后向傳播中運(yùn)行all_gather從所有rank中收集全部分片，以恢復(fù)該FSDP單元中的完整參數(shù)運(yùn)行后向計(jì)算運(yùn)行reduce_scatter以同步梯度丟棄參數(shù)”

此處，“rank”指的是一個(gè)GPU worker。

下圖是另一篇論文中有幫助的可視化圖，特別是對(duì)比了全分片和混合分片：

混合分片

混合分片包括分片和復(fù)制。這意味著，對(duì)于給定數(shù)量的worker/ GPUW，分片僅在大小為F 的子集之間發(fā)生，并在不同子集之間復(fù)制。

具體來(lái)說(shuō)，假設(shè)我們想要跨2個(gè)節(jié)點(diǎn)進(jìn)行多節(jié)點(diǎn)訓(xùn)練，每個(gè)節(jié)點(diǎn)在GCP中都是一個(gè)a2-ultragpu-8g節(jié)點(diǎn)。每個(gè)節(jié)點(diǎn)有8個(gè)A100 GPU，共有16個(gè)worker。那么可以使用混合分片在每個(gè)節(jié)點(diǎn)內(nèi)部對(duì)模型參數(shù)進(jìn)行分片，然后在節(jié)點(diǎn)之間復(fù)制。這意味著，在每個(gè)前向/反向傳播中，每個(gè)節(jié)點(diǎn)內(nèi)部有類似的all-gather和reduce-scatter操作，即從其他GPU獲取模型參數(shù)（節(jié)點(diǎn)內(nèi)部），并計(jì)算中間激活和梯度。然后，我們需要進(jìn)一步在節(jié)點(diǎn)之間進(jìn)行另一個(gè)all-gather操作，以獲得在該訓(xùn)練步驟中處理的總小批量數(shù)據(jù)的平均梯度值。當(dāng)我們被迫處理分片參數(shù)（即不能選擇ZeRO 2/1時(shí)），并處于多節(jié)點(diǎn)設(shè)置中時(shí)，這種方法尤其有吸引力。這類似于ZeRO++中的“層次劃分（hierarchical partitioning）”。

4 實(shí)現(xiàn)

如何使用DeepSpeed和FSDP？

DeepSpeed ZeRO/FSDP的主要優(yōu)勢(shì)之一在于，當(dāng)實(shí)際上只處于數(shù)據(jù)并行設(shè)置時(shí)，我們可以在數(shù)據(jù)+張量并行的情況下節(jié)省內(nèi)存，提升吞吐量。這意味著，我們無(wú)需進(jìn)行任何臨時(shí)的架構(gòu)更改，也無(wú)需通過(guò)混亂的`.to()`設(shè)備轉(zhuǎn)換或任何定制操作來(lái)改變前向傳播過(guò)程。因此，ZeRO/FSDP可以在不同的架構(gòu)中運(yùn)行，實(shí)現(xiàn)良好的集成效果。ZeRO是在微軟的DeepSpeed庫(kù)中實(shí)現(xiàn)的，并已整合到Accelerate庫(kù)中。而FSDP則是PyTorch本身的一部分，同樣也整合到了Accelerate庫(kù)。因此，我們可以使用Trainer API（其后端使用了Accelerate）中的任一策略，或者直接使用Accelerate。

流水線并行和張量并行？

目前，流水線并行（PP）和張量（TP）并行需要進(jìn)行架構(gòu)更改和/或調(diào)整模型的前向傳播過(guò)程。如果想要在DeepSpeed中使用流水線并行，那么則需要更改模型的架構(gòu)定義。

此外，由于Transformers庫(kù)為每個(gè)模型實(shí)現(xiàn)了數(shù)十種特性，這使得對(duì)其進(jìn)行集成變得非常復(fù)雜。當(dāng)然，我們已經(jīng)可以通過(guò)device_map="auto"實(shí)現(xiàn)樸素模型并行，但在使用多個(gè)GPU時(shí)，這是一種非常糟糕的策略，它只在模型無(wú)法適配單個(gè)GPU的情況下才有意義（否則最好還是只使用單個(gè)GPU）。如果確實(shí)需要流水線并行和張量并行，那么目前最好的選擇是使用Megatron-LM，并專注于他們支持的模型（如BERT、GPT-2、T5、Llama）。

此外，還可以在Megatron-DeepSpeed中使用ZeRO支持的DP、DeepSpeed PP和Megatron TP，但這些方法僅適用于基于BERT、GPT-2和T5的模型訓(xùn)練。Stas Bekman曾試圖在HuggingFace中實(shí)現(xiàn)流水線并行和張量并行，但未能成功。以下摘自Stas Bekman的博文：

Transformer現(xiàn)狀：目前的情況是，沒有一個(gè)模型支持完全的流水線并行（full-PP）。GPT2和T5模型具有樸素模型并行（naive MP）支持。目前的主要障礙是無(wú)法將模型轉(zhuǎn)換為nn.Sequential，并將所有輸入轉(zhuǎn)化為張量。因?yàn)槟壳澳Ｐ椭邪S多特性，這些特性使轉(zhuǎn)換變得非常復(fù)雜，因此我們需要?jiǎng)h除特性才能實(shí)現(xiàn)轉(zhuǎn)換。

我們已經(jīng)嘗試了開源，但這個(gè)問題太過(guò)棘手，盡管已向Stas Bekman求助，但仍未能解決。

DeepSpeed ZeRO和FSDP會(huì)繼續(xù)存在嗎？

我認(rèn)為，由于其易用性，DeepSpeed的ZeRO和PyTorch的FSDP很可能會(huì)被保留下來(lái)，或者說(shuō)只會(huì)被非常相似的策略取代。在語(yǔ)言大模型領(lǐng)域，變化是唯一的不變。

隨著新模型、架構(gòu)、注意力實(shí)現(xiàn)、位置嵌入改進(jìn)等的出現(xiàn)，能夠在幾小時(shí)內(nèi)用一個(gè)架構(gòu)替換另一個(gè)架構(gòu)，并在10M個(gè)樣本上訓(xùn)練40億+參數(shù)的模型變得非常重要。即使新的DP + PP + TP策略最終能夠?qū)崿F(xiàn)比ZeRO驅(qū)動(dòng)的DP更高的吞吐量，可能也不會(huì)有太多的采納量。出于相似原因，我認(rèn)為，DeepSpeed ZeRO和FSDP會(huì)獲得更多關(guān)注度，它們的吞吐量將進(jìn)一步優(yōu)化，甚至可能出現(xiàn)針對(duì)純數(shù)據(jù)并行類別的定制計(jì)算集群配置。因此，如果我們堅(jiān)持僅采用DP，結(jié)果不會(huì)太差。

（當(dāng)然，這與OpenAI、Anthropic等所進(jìn)行的超大型語(yǔ)言模型訓(xùn)練無(wú)關(guān)）

5 高效微調(diào)

這又是一個(gè)熱門話題！接下來(lái)我將簡(jiǎn)要列舉一些最受歡迎的優(yōu)化方法：

混合精度

對(duì)于大模型訓(xùn)練而言，這是一個(gè)不言而喻的選擇。簡(jiǎn)而言之，權(quán)重、激活和梯度以半精度格式存儲(chǔ)，同時(shí)你有一個(gè)FP32/單精度格式的“主副本”權(quán)重。兩種常用的半精度格式是BF16（由Google Brain開發(fā)的“Brain Float 16”）和FP16。FP16需要額外的損失擴(kuò)展以防止梯度下溢，而BF16似乎不存在這些問題。

拓展閱讀：

1. 混合精度訓(xùn)練：https://arxiv.org/abs/1710.03740
2. 性能和可擴(kuò)展性：如何適配更大的模型并進(jìn)行更快的訓(xùn)練：https://huggingface.co/docs/transformers/v4.18.0/en/performance

參數(shù)高效微調(diào)（PEFT）

PEFT旨在通過(guò)凍結(jié)大部分模型權(quán)重，將其中的子集/少量額外參數(shù)設(shè)置為可訓(xùn)練，從而在微調(diào)過(guò)程中降低內(nèi)存需求。LoRA是最流行的PEFT，其中對(duì)模型參數(shù)進(jìn)行了低秩版本的權(quán)重更新微調(diào)。

另一種有效的PEFT是IA(3)，它在基于Transformer體系結(jié)構(gòu)中的鍵、值和前饋層中注入可訓(xùn)練向量。對(duì)于LoRA和IA(3)，添加的權(quán)重/向量可以與基礎(chǔ)權(quán)重合并，這意味著，在推理時(shí)無(wú)需進(jìn)行額外的計(jì)算（加法/乘法）。不足之處在于，其性能可能不及完全微調(diào)時(shí)的性能。然而，這一情況已迅速改變，實(shí)際上，只要你向所有線性層添加可訓(xùn)練權(quán)重，基于LoRA的方法可以達(dá)到完全微調(diào)的性能。

就IA(3)而言，在我對(duì)GPT-2（770M）這樣的小型模型進(jìn)行的實(shí)驗(yàn)中，發(fā)現(xiàn)IA(3)可以在可訓(xùn)練參數(shù)的數(shù)量不到1/10的情況下與LoRA的性能相匹配。然而，這仍需要更多社區(qū)實(shí)驗(yàn)，尤其是在LLama-2-7B或Falcon-40B的規(guī)模上。

拓展閱讀：

• 語(yǔ)言大模型的低秩調(diào)整：https://arxiv.org/abs/2106.09685
• PEFT中的LoRA概念指南：https://huggingface.co/docs/peft/conceptual_guides/lora

IA3概念指南：

• PEFT中IA3的概念指南：https://huggingface.co/docs/peft/conceptual_guides/ia3
• 使用LoRA高效微調(diào)T5-XXL：https://www.philschmid.de/fine-tune-flan-t5-peft

Flash Attention

Flash Attention是一種快速（速度得到提升?。?、內(nèi)存高效（節(jié)省內(nèi)存！）、精確（無(wú)需近似?。?、IO感知（通過(guò)考慮GPU內(nèi)存的不同級(jí)別進(jìn)行讀/寫操作?。┑淖⒁饬λ惴?。 使用 Flash Attention 2（自2023年7月份推出），你可以在A100 80GB上獲得220 TFLOPS以上的性能（其最高性能為315 TFLOPS）。 換句話說(shuō)，當(dāng)Flash Attention 1在2022年中問世時(shí)，它擁有當(dāng)時(shí)可達(dá)到的最佳吞吐量，最高可達(dá)124 TFLOPS。 此外，根據(jù)我們之前研究的部分DeepSpeed ZeRO論文，預(yù)計(jì)在2021年，Tesla V100的可實(shí)現(xiàn)峰值吞吐量約為70 TFLOPS！

目前，F(xiàn)lash Attention支持Ampere、Ada或Hopper NVIDIA GPU（如A100、RTX 3090、H100等），且僅支持半精度數(shù)據(jù)類型bf16/fp16。要在Transformers中使用Flash Attention，只需對(duì)LLama和Falcon進(jìn)行一處更改（將use_flash_attention=True?傳遞給AutoModel）。對(duì)于其他模型，則需手動(dòng)將forward中使用的注意力函數(shù)改為Flash-attention的高吞吐量版本，不過(guò)這一點(diǎn)正在快速取得進(jìn)展。Flash Attention很快將集成到PyTorch的?scaled_dot_product_attention?中，這樣就不必再依賴猴補(bǔ)?。╩onkey patches）。（這本應(yīng)在v2.1版本中實(shí)現(xiàn)，但有望在不久的將來(lái)實(shí)現(xiàn)）

Flash Attention v1.0集成在Optimum中也已經(jīng)有一段時(shí)間了，但不能使用填充詞元（padding tokens），這使得它的作用非常受限。

拓展閱讀：

• ELI5：Flash Attention：https://gordicaleksa.medium.com/eli5-flash-attention-5c44017022ad
• FlashAttention-2：更快的關(guān)注力與更好的并行性和工作分區(qū)?https://tridao.me/publications/flash2/flash2.pdf

梯度/激活檢查點(diǎn)

通常情況下，每次在前向傳播中，所有中間激活都會(huì)被保留在內(nèi)存中，因?yàn)樗鼈冊(cè)谟?jì)算反向傳播時(shí)是必需的。梯度/激活檢查點(diǎn)（checkpointing）是一種通過(guò)僅保留部分中間激活，并根據(jù)需要重新計(jì)算其余部分來(lái)減少內(nèi)存消耗的技術(shù)。其中牽涉到的權(quán)衡是額外的重新計(jì)算時(shí)間。據(jù)HuggingFace提供的經(jīng)驗(yàn)法則，梯度檢查點(diǎn)會(huì)使訓(xùn)練減慢約20%。當(dāng)模型層數(shù)為N 時(shí)，激活的內(nèi)存需求從O(N)下降到??(?)?。當(dāng)然，因?yàn)榭們?nèi)存中存儲(chǔ)的不僅僅是激活，所以內(nèi)存消耗可能沒有那么大。

拓展閱讀：

• 《將大型網(wǎng)絡(luò)適配到內(nèi)存中》：https://medium.com/tensorflow/fitting-larger-networks-into-memory-583e3c758ff9
• 《性能和可擴(kuò)展性：如何適配更大的模型并加速訓(xùn)練》：https://huggingface.co/docs/transformers/v4.18.0/en/performance

量化

量化是黑客的最愛，主要有兩種量化方法，下面我將簡(jiǎn)要介紹：

• 后訓(xùn)練量化（PTQ）：旨在實(shí)現(xiàn)高效推斷。LLM.int8()，GPTQ就屬于這一類。

• 量化感知訓(xùn)練：從術(shù)語(yǔ)的原始意義出發(fā)，這意味著從一開始就使用量化的權(quán)重和激活來(lái)訓(xùn)練模型，以備后續(xù)推理使用。這正是我們想要的 ——一種使用量化參數(shù)進(jìn)行訓(xùn)練的策略。QLoRA屬于這一類（在某種程度上，因?yàn)樗鼘⒘炕系接?xùn)練過(guò)程中）。QLoRA主要將基礎(chǔ)預(yù)訓(xùn)練模型權(quán)重量化為8/4bits，然后以浮點(diǎn)半精度/全精度訓(xùn)練額外的LoRA參數(shù)。這是一種十分強(qiáng)大的策略，能在單個(gè)擁有48GB vRAM的GPU上微調(diào)超過(guò)600億參數(shù)的模型。當(dāng)然，需要注意的是，其吞吐量會(huì)大大降低。原因很簡(jiǎn)單：每當(dāng)計(jì)算給定層的激活時(shí)，都會(huì)發(fā)生額外的去量化步驟（具體數(shù)字取決于硬件設(shè)置。例如，你可能能夠在8個(gè)A100上運(yùn)行DeepSpeed ZeRO 3的Falcon-40B訓(xùn)練時(shí)不進(jìn)行量化，在這里單獨(dú)使用量化毫無(wú)意義，即使你可以獲得更好的批次大小。不過(guò)，在只有2個(gè)A100的情況下就不一樣了）。

補(bǔ)充信息：完整的QLoRA論文值得一讀。他們的方法除了能在消費(fèi)級(jí)GPU上訓(xùn)練65B參數(shù)的模型（這是當(dāng)時(shí)最大的開源語(yǔ)言模型）之外，該論文還表示，基于LoRA的訓(xùn)練可以與完全微調(diào)相媲美（他們?cè)黾恿烁嗫捎?xùn)練的層，但通過(guò)對(duì)基礎(chǔ)權(quán)重進(jìn)行量化使其比原始LoRA配置更高效），并且數(shù)據(jù)集的質(zhì)量在監(jiān)督微調(diào)中至關(guān)重要（450K FLAN樣本比9K高質(zhì)量人工標(biāo)注樣本更糟糕）。

拓展閱讀：

• QLoRA：https://arxiv.org/abs/2305.14314
• 對(duì)Transformer模型進(jìn)行量化：https://huggingface.co/docs/transformers/v4.34.0/en/main_classes/quantization

梯度累積

梯度累積是一種提高有效批次大小的方法，盡管吞吐量會(huì)有所下降（有時(shí)甚至?xí)陆抵亮悖?/p>

比如，假設(shè)你的批次大小為2，梯度累積步長(zhǎng)為4。使用梯度累積，你會(huì)在每個(gè)訓(xùn)練步驟中進(jìn)行常規(guī)的前向和后向傳播，但是，你會(huì)在每4個(gè)訓(xùn)練步驟中調(diào)用一次優(yōu)化器步驟（在PyTorch中為optimizer.step()）。這意味著，在4個(gè)步驟中會(huì)進(jìn)行梯度累積，然后使用這8個(gè)樣本的平均梯度更新模型權(quán)重。這樣一來(lái)，你的批次大小增加了，這使得權(quán)重更新變得更加平穩(wěn)，但內(nèi)存消耗保持不變。實(shí)際上，使用多GPU/多節(jié)點(diǎn)訓(xùn)練的梯度累積可以增加批次大小，同時(shí)加快訓(xùn)練速度。因?yàn)樵谡Ｓ?xùn)練中，梯度在本地（由GPU處理的批次）被局部平均，并且每個(gè)訓(xùn)練步驟都需執(zhí)行一次全局歸約。而使用梯度累積，你可以在更大的間隔（即gradient_accumulation_steps）上執(zhí)行這一全局歸約操作。減少此類全局歸約可降低worker間的通信負(fù)擔(dān)（可能還包括節(jié)點(diǎn)間的通信），進(jìn)而提高訓(xùn)練速度。

拓展閱讀：

• 使用DeepSpeed API進(jìn)行訓(xùn)練：https://www.deepspeed.ai/training/

是否應(yīng)該一直嘗試增加批次大小？

閱讀完我們討論的所有內(nèi)容之后，這是一個(gè)吸引人的問題。答案是：否！要記得，我們的目標(biāo)是盡可能快地使用現(xiàn)有硬件，訓(xùn)練出可能的最優(yōu)神經(jīng)網(wǎng)絡(luò)。因此，即使你只使用了可用GPU內(nèi)存的75%（例如在A100中為60GB/80GB），也可能已經(jīng)達(dá)到了系統(tǒng)中的最大吞吐量。這意味著，進(jìn)一步增加批次大小將導(dǎo)致相應(yīng)時(shí)延的增加，結(jié)果是吞吐量不會(huì)提高，甚至可能降低。此外，即使你有一套強(qiáng)大的硬件配置來(lái)訓(xùn)練40B+參數(shù)的模型，僅通過(guò)上述內(nèi)存優(yōu)化來(lái)增加批次大小還不夠，因?yàn)檫@可能會(huì)影響收斂性能。因此，很難針對(duì)大模型展開一項(xiàng)良好的研究。

在Transformer時(shí)代之前有一篇論文（https://openreview.net/forum?id=H1oyRlYgg）顯示，大型批次大小可能會(huì)影響泛化性能。DeepSpeed提到，你可以找到更好的超參數(shù)/優(yōu)化器選擇，使更大型批次大小發(fā)揮作用，正如1-cycle 學(xué)習(xí)速率調(diào)度指南（https://www.deepspeed.ai/tutorials/one-cycle/）所證明的那樣。DeepSpeed的一位作者還提出，為使大規(guī)模訓(xùn)練運(yùn)行達(dá)到目標(biāo)收斂速率，全局批次大小通常是固定的?？傊?，更大型的批次大小并不一定更好！

到底多大的批次大小算大？

便于參考，BLOOM-176B預(yù)訓(xùn)練采用了366B個(gè)詞元，全局批次大小為2048。尤其是在微調(diào)階段，目前還不清楚批次大小大到何種程度才會(huì)影響模型的收斂。

6 實(shí)踐指南

總結(jié)

綜上所述，以下是嘗試在1000萬(wàn)以上規(guī)模的數(shù)據(jù)集上對(duì)10-100B+模型參數(shù)進(jìn)行實(shí)驗(yàn)和微調(diào)的實(shí)用指南（我有DeepSpeed的經(jīng)驗(yàn)，尚未使用FSDP，所以我將重點(diǎn)放在這上面）：

• 默認(rèn)情況下使用BF16/ FP16。BF16基本不需要其他配置參數(shù)，通常不會(huì)出現(xiàn)任何溢出問題（相反，F(xiàn)P16可能會(huì)因不同的損失擴(kuò)展因子導(dǎo)致不同的結(jié)果，并且由于動(dòng)態(tài)范圍較小，可能出現(xiàn)更多的溢出問題），因此非常方便。
• 使用LoRA，并將可訓(xùn)練參數(shù)添加到所有線性層。如果你想緊密遵循QLoRA，可在這里（https://github.com/artid?oro/qlora/blob/7f4e95a68dc076bea9b3a413d2b512eca6d004e5/qlora.py#L248）使用他們的實(shí)用函數(shù)。
• 如果你GPU支持Flash Attention，那么可以使用它。目前，F(xiàn)lash Attention 2可在HuggingFace的Llama 2和Falcon上使用，其他模型可能需要進(jìn)行調(diào)整。
• 使用梯度/激活檢查點(diǎn)。這將略微降低吞吐量。如果你使用了Flash Attention，可能就不需要梯度檢查點(diǎn)。參考Flash Attention論文（另見Tri Dao的建議，https://github.com/EleutherAI/gpt-neox/pull/725#issuecomment-1374134498）。

• 在你的數(shù)據(jù)加載器中使用高效的采樣器，如多重包采樣器（multi-pack sampler）。
• 如果你有多個(gè)GPU，首先嘗試使用BF16+LoRA+梯度檢查點(diǎn)+DeepSpeed ZeRO 3。
• GPU內(nèi)存有限時(shí)，可使用量化。QLoRA風(fēng)格的訓(xùn)練目前僅適用于DeepSpeed ZeRO 1/2。因此，即使它在模型參數(shù)方面非常高效，但在ZeRO 1/2中仍有參數(shù)冗余，并且吞吐量也會(huì)減少。
• 隨著GPU的增加（如8個(gè)V100或A100），DS ZeRO-3成為了最佳選擇。DS ZeRO-2也不錯(cuò)，但你可能會(huì)受到CPU內(nèi)存限制的影響（在模型初始化期間），因?yàn)槟Ｐ蛥?shù)會(huì)在所有工作節(jié)點(diǎn)上被復(fù)制。例如，如果要在具有8個(gè)GPU的節(jié)點(diǎn)上使用Falcon-40B模型，那么CPU內(nèi)存需要超過(guò)1.5TB，而云容器實(shí)例（AWS、GCP等）很少有這么大的內(nèi)存，所以會(huì)受到限制。當(dāng)然，在家用服務(wù)器上可能不適用，因?yàn)榧矣梅?wù)器沒有用于GPU間通信的NVLink。DS ZeRO-3具有更多的GP間通信，因此NVLink很重要。
• 在小規(guī)模多節(jié)點(diǎn)設(shè)置中，啟用了層次分區(qū)（或FSDP的混合分片）的DeepSpeed ZeRO-3似乎是最佳選擇。如果你有Infiniband互聯(lián)，大多可以使用普通的DeepSpeed ZeRO-3，并用于更大模型。
• 如果在上述所有優(yōu)化之后仍然批次大小不足，應(yīng)該使用梯度累積。對(duì)于大模型和多GPU/多節(jié)點(diǎn)設(shè)置，梯度累積可加快訓(xùn)練時(shí)間。
• 如果你的GPU內(nèi)存非常有限，可以激活CPU/磁盤卸載（通過(guò)硬盤，這必須是ZeRO-Infinity 的NVMe）。隨著Flash Attention 2的出現(xiàn)，我們需要對(duì)純GPU訓(xùn)練和GPU + NVMe/CPU卸載之間的吞吐量差距進(jìn)行另一項(xiàng)研究。我懷疑現(xiàn)在這個(gè)差距比以前大得多，因此只有在真正受限的情況下才進(jìn)行卸載（這就是為什么這是要嘗試的最后一個(gè)優(yōu)化）。ZeRO-Infinity 優(yōu)于ZeRO-Offload，對(duì)于這一點(diǎn)你必須使用ZeRO Stage 3。
• 計(jì)算有效的批次大小，并相應(yīng)調(diào)整超參數(shù)。一個(gè)通用的指導(dǎo)原則是，隨著有效批次大小的增加，要增加學(xué)習(xí)率。正如OpenAI的微調(diào)文檔中提到的，即使對(duì)于100B+的模型，這似乎仍然成立。
• 最后，在開始訓(xùn)練時(shí)，使用?htop?監(jiān)控RAM使用情況（有時(shí)會(huì)出現(xiàn)RAM OOM問題），同時(shí)使用?nvidia-smi?確保GPU沒有被數(shù)據(jù)預(yù)處理拖慢（你應(yīng)該爭(zhēng)取接近100%的GPU利用率，即使GPU內(nèi)存使用較少）。

關(guān)于超參數(shù)的更多內(nèi)容：以下為學(xué)習(xí)率縮放的相關(guān)評(píng)論。事實(shí)證明，在預(yù)訓(xùn)練過(guò)程中，隨著模型的增大（比如100B+），即使使用了更大的批次大小，通常來(lái)說(shuō)，我們還是應(yīng)該降低而非增加學(xué)習(xí)率，這一點(diǎn)可以在OpenAI的GPT-3論文和BLOOM論文中得到證實(shí)。要建立對(duì)這種模式的直覺十分困難，因此我們只能采取非常實(shí)證的方法，如果實(shí)驗(yàn)結(jié)果與之相悖，就迅速更新所有先驗(yàn)知識(shí)。

補(bǔ)充指南

在Transformers的文檔中提供了許多關(guān)于性能和可擴(kuò)展性的有用技巧。如果你關(guān)心在家搭建和管理服務(wù)器時(shí)可能遇到的問題，或者NVLink的具體作用，又或者想深入了解內(nèi)存管理，這些內(nèi)容都非常有價(jià)值。Stas Bekman系列（https://github.com/stas00/ml-engineering）收錄了許多與調(diào)試和性能相關(guān)的實(shí)用提示，可能會(huì)有所幫助。

7 關(guān)于DeepSpeed和FSDP的更多內(nèi)容

使用DeepSpeed ZeRO-3進(jìn)行多節(jié)點(diǎn)訓(xùn)練

對(duì)于普通ZeRO-3，你需要確?？绻?jié)點(diǎn)的參數(shù)通信不會(huì)成為一個(gè)巨大的瓶頸，否則，即便使用兩個(gè)節(jié)點(diǎn)，吞吐量的提升可能也非常有限。通過(guò)Stas Bekman的調(diào)查，我們對(duì)需要具備怎樣的跨節(jié)點(diǎn)網(wǎng)絡(luò)有了清晰的認(rèn)識(shí)：

對(duì)于大小為M?的模型，有N?個(gè)節(jié)點(diǎn)，每個(gè)節(jié)點(diǎn)有G?個(gè)GPU，每個(gè)節(jié)點(diǎn)在每個(gè)訓(xùn)練步驟中發(fā)送約48*M/Nbits的數(shù)據(jù)。即便是對(duì)于兩個(gè)8xA100節(jié)點(diǎn)，配備一個(gè)40B參數(shù)模型，每個(gè)節(jié)點(diǎn)在每個(gè)訓(xùn)練步驟中通信的數(shù)據(jù)量就達(dá)到了120GB！對(duì)于大規(guī)模訓(xùn)練（64GPU+），確實(shí)需要具備1000Gbps的InfiniBand互連。對(duì)于較小規(guī)模的多節(jié)點(diǎn)訓(xùn)練，可選擇100-400Gbps的帶寬。在亞馬遜EC2 P4d實(shí)例（配備了8個(gè)A100，用于機(jī)器學(xué)習(xí)訓(xùn)練的節(jié)點(diǎn)）上，通常使用彈性網(wǎng)絡(luò)適配器（Elastic Fabric Adapter，EFA）作為跨節(jié)點(diǎn)的網(wǎng)絡(luò)接口，根據(jù)規(guī)格，你可以獲得高達(dá)400Gbps的網(wǎng)絡(luò)帶寬，這相當(dāng)不錯(cuò)！實(shí)際上，你獲得的確切帶寬約為340Mbps，因此應(yīng)該規(guī)劃使用規(guī)格中列出的最大帶寬的80-85%。用于AWS P5實(shí)例的EFA v2配備了H100，速度快了整整8倍。

要測(cè)算你所獲得的跨節(jié)點(diǎn)網(wǎng)絡(luò)帶寬，可使用Stas Bekman的實(shí)用函(https://github.com/stas00/ml-engineering/blob/9a51114f8377350bfbf1764f23feac441e865401/multi-node/all_reduce_bench.py）。這個(gè)函數(shù)專門用于對(duì)所有歸約操作進(jìn)行基準(zhǔn)測(cè)試，因此能夠準(zhǔn)確測(cè)試訓(xùn)練過(guò)程中所能看到的情況。

建議：在沒有Infiniband的情況下，最好使用ZeRO++和層次分區(qū)（hpZ）。要實(shí)現(xiàn)這一點(diǎn)，你需要將zero_hpz_partition_size?配置參數(shù)設(shè)置為每個(gè)節(jié)點(diǎn)的GPU數(shù)量/rank。例如，如果你使用兩個(gè)節(jié)點(diǎn)進(jìn)行訓(xùn)練，每個(gè)節(jié)點(diǎn)有8個(gè)A100 GPU，那么zero_hpz_partition_size?將設(shè)置為8。

Accelerate/Transformer 支持：目前尚不清楚加速DeepSpeed集成是否支持層次分區(qū)（hpZ）。表面上看，應(yīng)該支持，因?yàn)榧铀賾?yīng)該整合DeepSpeed ZeRO的所有功能，而hpZ只是DeepSpeed配置文件中的一個(gè)參數(shù)更改。我已經(jīng)提了一個(gè)issue（2023年10月1日），如有需要，我會(huì)更新這篇帖子。

對(duì)于FSDP，你將使用混合分片策略（HYBRID_SHARD）。對(duì)于多節(jié)點(diǎn)訓(xùn)練（至少是小規(guī)模的，沒有Infiniband的情況下），這似乎是一個(gè)不錯(cuò)的選擇，但有用戶通過(guò)艱難的方式發(fā)現(xiàn)了這一點(diǎn)。Accelerate已經(jīng)支持帶有混合分片的FSDP。

DeepSpeed 內(nèi)存需求

當(dāng)你擁有新的基礎(chǔ)架構(gòu)設(shè)置，并希望嘗試DeepSpeed時(shí)，絕對(duì)應(yīng)該使用DeepSpeed的內(nèi)存估算器。

DeepSpeed ZeRO 1/2

deepspeed.runtime.zero.stage_1_and_2.estimate_zero2_model_states_mem_needs_all_live(model,
  num_gpus_per_node=1, num_nodes=1, additional_buffer_factor=1.5)

使用默認(rèn)緩沖因子（這是一個(gè)簡(jiǎn)單地?cái)U(kuò)展所有CPU和GPU內(nèi)存估算的估算因子），對(duì)于一個(gè)擁有3B參數(shù)的模型，在擁有1個(gè)節(jié)點(diǎn)和8個(gè)GPU的環(huán)境中，你將得到以下結(jié)果，來(lái)自于DeepSpeed文檔：

python -c 'from transformers import AutoModel; \
from deepspeed.runtime.zero.stage_1_and_2 import estimate_zero2_model_states_mem_needs_all_live; \
model = AutoModel.from_pretrained("t5-3b"); \
estimate_zero2_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)'
?
Estimated memory needed for params, optim states and gradients for a:
HW: Setup with 1 node, 8 GPUs per node.
SW: Model with 2851M total params.
  per CPU  |  per GPU |   Options
  127.48GB |   5.31GB | offload_optimizer=cpu
  127.48GB |  15.93GB | offload_optimizer=none

以Falcon-40b為例，使用ZeRO 1/2所需的CPU RAM大于1.5TB。你無(wú)需實(shí)際運(yùn)行訓(xùn)練/微調(diào)代碼來(lái)測(cè)試這一點(diǎn)，只需使用上述模型和硬件設(shè)置運(yùn)行上述命令，就會(huì)得到一個(gè)估算。對(duì)于FSDP，我沒有找到類似的內(nèi)存估算器，因此可直接使用DeepSpeed ZeRO-3對(duì)FSDP的全分片估算。

注意：再補(bǔ)充一點(diǎn)關(guān)于CPU RAM不足的情況。這可能很難調(diào)試，因?yàn)槟愕倪M(jìn)程將在日志中沒有任何信息的情況下失敗。添加PyTorch分布式調(diào)試標(biāo)志（NCCL_DEBUG=INFO, TORCH_DISTRIBUTED_DEBUG=INFO?）也無(wú)濟(jì)于事，因?yàn)檫@是一個(gè)RAM問題。我只是在初始訓(xùn)練階段通過(guò)監(jiān)視htop（如前文所述）發(fā)現(xiàn)了這一困難。

DeepSpeed ZeRO 3

python deepspeed.runtime.zero.stage3.estimate_zero3_model_states_mem_needs_all_live(model, \
  num_gpus_per_node=1, num_nodes=1, additional_buffer_factor=1.5)

在一個(gè)節(jié)點(diǎn)上，每個(gè)節(jié)點(diǎn)有8個(gè)GPU的硬件設(shè)置下，對(duì)于一個(gè)包含3B參數(shù)的模型，使用以下命令可以估算參數(shù)、優(yōu)化狀態(tài)和梯度所需的內(nèi)存：

python -c 'from transformers import AutoModel; \
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \
model = AutoModel.from_pretrained("t5-3b"); \
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)'
?
Estimated memory needed for params, optim states and gradients for a:
HW: Setup with 1 node, 8 GPUs per node.
SW: Model with 2851M total params, 32M largest layer params.
  per CPU  |  per GPU |   Options
   71.71GB |   0.12GB | offload_param=cpu , offload_optimizer=cpu , zero_init=1
  127.48GB |   0.12GB | offload_param=cpu , offload_optimizer=cpu , zero_init=0
   63.74GB |   0.79GB | offload_param=none, offload_optimizer=cpu , zero_init=1
  127.48GB |   0.79GB | offload_param=none, offload_optimizer=cpu , zero_init=0
    1.47GB |   6.10GB | offload_param=none, offload_optimizer=none, zero_init=1
  127.48GB |   6.10GB | offload_param=none, offload_optimizer=none, zero_init=0

這也讓你對(duì)不同的卸載策略和初始化有了一個(gè)概念！使用zero_init=1?時(shí)，模型權(quán)重將以一種內(nèi)存可擴(kuò)展的方式初始化，一旦分配后，權(quán)重將立即在你的worker之間分區(qū)。當(dāng)zero_init=0?時(shí)，CPU RAM需求可能會(huì)飆升（對(duì)于10B+模型和多個(gè)GPU來(lái)說(shuō)可能會(huì)達(dá)到TB 級(jí)別），因此你絕對(duì)應(yīng)該使用zero_init=1。在多GPU設(shè)置中使用CPU卸載也會(huì)增加大量RAM的需求，并且可能導(dǎo)致RAM OOM。使用估算器，之后如果確實(shí)需要，如CPU RAM不足的話，可以切換到NVMe卸載。

使用Accelerate

Accelerate 旨在提供一個(gè)統(tǒng)一的界面，用于啟動(dòng)各種分布式訓(xùn)練運(yùn)行，同時(shí)為你提供在純 PyTorch中編寫代碼的靈活性。目前看來(lái)，如果你打算在FSDP和DeepSpeed之間切換同一段代碼，可能會(huì)有一些注意事項(xiàng)。例如，對(duì)于FSDP，你必須在實(shí)例化優(yōu)化器之前調(diào)用accelerator.prepare(model)。我不確定相同的方法是否適用于 DeepSpeed（對(duì)于 DeepSpeed，你可以對(duì)一切進(jìn)行.prepare()調(diào)用）。還有一些其他注意事項(xiàng)，這里就不展開了，但可以查看以下Accelerate的文檔。

拓展閱讀：1.使用Accelerate進(jìn)行FSDP：https://huggingface.co/docs/accelerate/usage_guides/fsdp

2.使用PyTorch FSDP對(duì)Llama 2 70B進(jìn)行微調(diào)：https://huggingface.co/blog/ram-efficient-pytorch-fsdp

8開源代碼庫(kù)

短短幾個(gè)月內(nèi)，開源代碼庫(kù)已取得了長(zhǎng)足進(jìn)步。在這里，我主要想回答“如果我想立即開始使用開源代碼庫(kù)進(jìn)行微調(diào)，可以使用什么？應(yīng)該牢記什么？”等問題。如往常一樣，細(xì)節(jié)決定成敗。接下來(lái)我將總結(jié)FastChat和Axolotl的平臺(tái)功能，這是兩個(gè)最實(shí)用、最受歡迎的平臺(tái)。

FastChat

FastChat是一個(gè)用于微調(diào)、服務(wù)和評(píng)估LLM-based chatbots from LMSys的平臺(tái)。其功能包括：

Serving

你可以使用FastChat為L(zhǎng)lama、Falcon、WizardLM等模型提供服務(wù)。（支持的模型列表）添加新模型進(jìn)行推理/服務(wù)似乎非常簡(jiǎn)單，并支持因果模型（如 Llama）和序列到序列模型（如 T5）。他們還支持使用CPU卸載和不同的量化方案等服務(wù)。在幕后，F(xiàn)astChat使用了出色的vLLM庫(kù)進(jìn)行高效推理，這也是另一個(gè)LMSys開源項(xiàng)目。

微調(diào)

• 支持微調(diào)的模型有Llama、T5和Baichuan。如有錯(cuò)誤，煩請(qǐng)指正。主要的微調(diào)腳本僅適用于Llama模型，并有針對(duì)Llama特定的魔法數(shù)字。此外，還有針對(duì)T5和Baichuan模型的額外微調(diào)腳本。Falcon模型的訓(xùn)練支持仍有待解決。
• 專門針對(duì)（單輪/多輪）對(duì)話數(shù)據(jù)監(jiān)督微調(diào)，以訓(xùn)練聊天機(jī)器人。在類似FLAN的數(shù)據(jù)集上，用（instruction, response）對(duì)進(jìn)行指令微調(diào)是有可能的（單輪對(duì)話），但不可能混合因果語(yǔ)言建模數(shù)據(jù)集等其他數(shù)據(jù)集格式，或者在多個(gè)數(shù)據(jù)集上進(jìn)行訓(xùn)練。
• 支持LoRA和基于QLoRA的訓(xùn)練。提供的DeepSpeed配置是參考配置，適用于低資源環(huán)境（例如，1個(gè)V100 32GB用于微調(diào)Llama-13B），因此會(huì)默認(rèn)啟動(dòng)CPU卸載。請(qǐng)確保根據(jù)你的硬件設(shè)置進(jìn)行修改。更多參考信息請(qǐng)參閱訓(xùn)練文檔。
• FastChat使用Trainer API，并且與幾乎所有可用的開源訓(xùn)練包一樣，在訓(xùn)練和評(píng)估中僅支持單一的同質(zhì)數(shù)據(jù)集（在這種情況下，還必須是對(duì)話數(shù)據(jù)）。這意味著，你可以看到組合數(shù)據(jù)集的訓(xùn)練損失和評(píng)估損失，但在監(jiān)控運(yùn)行時(shí)不會(huì)看到其他信息。

評(píng)估

FastChat的評(píng)估包以MT-bench為基礎(chǔ)，是一個(gè)基于多輪對(duì)話的評(píng)估數(shù)據(jù)集。在評(píng)分方面，他們使用LLM作為評(píng)委，可以利用GPT-3.5或GPT-4等更強(qiáng)大的語(yǔ)言模型對(duì)模型輸出評(píng)分。

Axolotl

Axolotl是一個(gè)用于微調(diào)語(yǔ)言模型的大規(guī)模開源項(xiàng)目，具有以下顯著特點(diǎn)：

• 支持Llama、MPT、Falcon等各種因果語(yǔ)言模型。目前尚不支持序列到序列模型。
• 具有多種格式，支持在多個(gè)數(shù)據(jù)集上訓(xùn)練。在我看來(lái)，對(duì)Axolotl來(lái)說(shuō)，這是一個(gè)極其重要的功能。數(shù)據(jù)集格式的完整列表太過(guò)龐大，但你可以在類似FLAN的指令微調(diào)數(shù)據(jù)集、類似ShareGPT的基于對(duì)話的數(shù)據(jù)集以及簡(jiǎn)單的基于文本補(bǔ)全的數(shù)據(jù)集上進(jìn)行訓(xùn)練（用于因果語(yǔ)言建模的純文本/代碼）。Axolotl還支持預(yù)標(biāo)記數(shù)據(jù)集。
• 支持LoRA、QLoRA，以及多包采樣器，該采樣器將相似長(zhǎng)度的序列打包在一起，以避免在填充詞元上浪費(fèi)計(jì)算。
• Axolotl還使用Trainer API，并具有許多用于自定義評(píng)估和記錄的功能。你可以在MMLU上進(jìn)行評(píng)估，或者在本地基準(zhǔn)數(shù)據(jù)集上記錄訓(xùn)練期間的損失/準(zhǔn)確性。
• 此外，Axolotl還支持FSDP和DeepSpeed，主要是因?yàn)樗鼈冎恍枳孴rainer處理這一部分。Flash-Attention也適用于Llama、BTLM和新的Mistral等其他模型。
• 因?yàn)樵赥rainer上定制起來(lái)比較困難，所以Axolotl不支持對(duì)不同任務(wù)（如MMLU和MATH）進(jìn)行評(píng)估，不能分開可視化損失曲線。
• 對(duì)此，我有一個(gè)小小的疑問，除?prediction_step?和?training_step?等方法之外，Axolotl最終幾乎定制了Trainer的每個(gè)部分。如果他們已經(jīng)可以使用 torch.distributed.gather ，我想知道為什么不直接從一開始就用純PyTorch + Accelerate編寫所有代碼，這樣本就可以避免冗余。Trainer還提供了十幾種優(yōu)化器以供選擇，但監(jiān)督微調(diào)只需要其中的一小部分。

實(shí)用的微調(diào)指南

HuggingFace近期發(fā)布的幾篇值得閱讀的微調(diào)指南：

• 在2個(gè)8xA100-ultra-80GB節(jié)點(diǎn)上使用PEFT和FSDP對(duì)Llama 2 70B進(jìn)行微調(diào)：https://huggingface.co/blog/ram-efficient-pytorch-fsdp
• 在2個(gè)8X100-ultra-80GB節(jié)點(diǎn)上使PEFT和DeepSpeed對(duì)Falcon 180B進(jìn)行微調(diào)：https://medium.com/@sourabmangrulkar/falcon-180b-finetuning-using-peft-and-deepspeed-b92643091d99
• 在1個(gè)8xA100-ultra-80GB節(jié)點(diǎn)上使用PEFT和DeepSpeed對(duì)Falcon 180B進(jìn)行微調(diào)：https://www.philschmid.de/deepspeed-lora-flash-attention

9超大型模型訓(xùn)練

首先，在搭載8枚A100顯卡的單一節(jié)點(diǎn)上，我們可以訓(xùn)練多大的模型。通過(guò)上述優(yōu)化措施，目前我們可以在搭載8枚A100-80GB顯卡的DGX節(jié)點(diǎn)上微調(diào)Falcon 180B。根據(jù)HuggingFace的指南（https://www.philschmid.de/deepspeed-lora-flash-attention），你可以在153分鐘內(nèi)，使用15000個(gè)樣本，在dolly數(shù)據(jù)集上對(duì)一個(gè)包含180億參數(shù)的模型進(jìn)行3輪微調(diào)，其中有效批次大小為64，序列長(zhǎng)度為2048個(gè)詞元。由于文本塊是作為預(yù)處理步驟進(jìn)行的，因此不知道準(zhǔn)確的吞吐量數(shù)值，當(dāng)然我們也需要估算每個(gè)訓(xùn)練步驟所需的時(shí)間（因?yàn)榭傆?xùn)練時(shí)間包括了評(píng)估和檢查點(diǎn)所需的時(shí)間），但每秒吞吐量可能大于5個(gè)樣本，對(duì)于180B參數(shù)的模型來(lái)說(shuō)，這是相當(dāng)不錯(cuò)的表現(xiàn)。

如果想了解更多關(guān)于管理大規(guī)模訓(xùn)練運(yùn)行的信息，可以向Stas Bekman尋求幫助。前文已經(jīng)提到過(guò)，但我還是再次重申：開源社區(qū)要感謝Stas Bekman的地方很多，其中包括他對(duì)大規(guī)模訓(xùn)練的出色論述。大規(guī)模訓(xùn)練的調(diào)試、監(jiān)控、深入研究大規(guī)模多節(jié)點(diǎn)網(wǎng)絡(luò)設(shè)置等內(nèi)容都可以在他的合集中找到。不過(guò)其中一些關(guān)于分布式訓(xùn)練的資料和建議可能有點(diǎn)過(guò)時(shí)。文章來(lái)源地址http://www.zghlxwxcb.cn/news/detail-801372.html

其他人都在看

• GPU架構(gòu)與計(jì)算入門指南
• 為什么開源大模型終將勝出
• LoRA和QLoRA微調(diào)語(yǔ)言大模型
• 可復(fù)現(xiàn)的語(yǔ)言大模型推理性能指標(biāo)
• ChatGPT規(guī)?；?wù)的經(jīng)驗(yàn)與教訓(xùn)
• 微調(diào)語(yǔ)言大模型選LoRA還是全參數(shù)
• 開源語(yǔ)言大模型演進(jìn)史：向LLaMA2看齊

到了這里，關(guān)于語(yǔ)言大模型的分布式訓(xùn)練與高效微調(diào)指南的文章就介紹完了。如果您還想了解更多內(nèi)容，請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！