本文，我們將了解如何基于 PyTorch 最新的?完全分片數(shù)據(jù)并行 (Fully Sharded Data Parallel，F(xiàn)SDP)?功能用?Accelerate?庫(kù)來(lái)訓(xùn)練大模型。

動(dòng)機(jī) ??

隨著機(jī)器學(xué)習(xí) (ML) 模型的規(guī)模、大小和參數(shù)量的不斷增加，ML 從業(yè)者發(fā)現(xiàn)在自己的硬件上訓(xùn)練甚至加載如此大的模型變得越來(lái)越難。?一方面，人們發(fā)現(xiàn)大模型與較小的模型相比，學(xué)習(xí)速度更快 (數(shù)據(jù)和計(jì)算效率更高) 且會(huì)有顯著的提升 [1]; 另一方面，在大多數(shù)硬件上訓(xùn)練此類模型變得令人望而卻步。

分布式訓(xùn)練是訓(xùn)練這些機(jī)器學(xué)習(xí)大模型的關(guān)鍵。大規(guī)模分布式訓(xùn)練?領(lǐng)域最近取得了不少重大進(jìn)展，我們將其中一些最突出的進(jìn)展總結(jié)如下:

1. 使用 ZeRO 數(shù)據(jù)并行 - 零冗余優(yōu)化器 [2]

2. 階段 1: 跨數(shù)據(jù)并行進(jìn)程 / GPU 對(duì)優(yōu)化器狀態(tài)?進(jìn)行分片

3. 階段 2: 跨數(shù)據(jù)并行進(jìn)程/ GPU 對(duì)優(yōu)化器狀態(tài) + 梯度?進(jìn)行分片

4. 階段 3: 跨數(shù)據(jù)并行進(jìn)程 / GPU 對(duì)優(yōu)化器狀態(tài) + 梯度 + 模型參數(shù)?進(jìn)行分片

5. CPU 卸載: 進(jìn)一步將 ZeRO 階段 2 的優(yōu)化器狀態(tài) + 梯度?卸載到 CPU 上 [3]

6. 張量并行 [4]: 模型并行的一種形式，通過(guò)對(duì)各層參數(shù)進(jìn)行精巧的跨加速器 / GPU 分片，在實(shí)現(xiàn)并行計(jì)算的同時(shí)避免了昂貴的通信同步開銷。

7. 流水線并行 [5]: 模型并行的另一種形式，其將模型的不同層放在不同的加速器 / GPU 上，并利用流水線來(lái)保持所有加速器同時(shí)運(yùn)行。舉個(gè)例子，在第 2 個(gè)加速器 / GPU 對(duì)第 1 個(gè) micro batch 進(jìn)行計(jì)算的同時(shí)，第 1 個(gè)加速器 / GPU 對(duì)第 2 個(gè) micro batch 進(jìn)行計(jì)算。

8. 3D 并行 [3]: 采用?ZeRO 數(shù)據(jù)并行 + 張量并行 + 流水線并行?的方式來(lái)訓(xùn)練數(shù)百億參數(shù)的大模型。例如，BigScience 176B 語(yǔ)言模型就采用了該并行方式 [6]。

本文我們主要關(guān)注 ZeRO 數(shù)據(jù)并行，更具體地講是 PyTorch 最新的?完全分片數(shù)據(jù)并行 (Fully Sharded Data Parallel，F(xiàn)SDP)?功能。DeepSpeed?和?FairScale?實(shí)現(xiàn)了 ZeRO 論文的核心思想。我們已經(jīng)將其集成到了?transformers?的?Trainer?中，詳見博文 通過(guò) DeepSpeed 和 FairScale 使用 ZeRO 進(jìn)行更大更快的訓(xùn)練[10]。最近，PyTorch 已正式將 Fairscale FSDP 整合進(jìn)其 Distributed 模塊中，并增加了更多的優(yōu)化。

Accelerate ??: 無(wú)需更改任何代碼即可使用 PyTorch FSDP

我們以基于 GPT-2 的 Large (762M) 和 XL (1.5B) 模型的因果語(yǔ)言建模任務(wù)為例。

以下是預(yù)訓(xùn)練 GPT-2 模型的代碼。其與 此處 的官方因果語(yǔ)言建模示例相似，僅增加了 2 個(gè)參數(shù)?n_train?(2000) 和?n_val?(500) 以防止對(duì)整個(gè)數(shù)據(jù)集進(jìn)行預(yù)處理/訓(xùn)練，從而支持更快地進(jìn)行概念驗(yàn)證。

run_clm_no_trainer.py

運(yùn)行?accelerate config?命令后得到的 FSDP 配置示例如下:?

多 GPU FSDP

本文我們使用單節(jié)點(diǎn)多 GPU 上作為實(shí)驗(yàn)平臺(tái)。我們比較了分布式數(shù)據(jù)并行 (DDP) 和 FSDP 在各種不同配置下的性能。我們可以看到，對(duì) GPT-2 Large(762M) 模型而言，DDP 尚能夠支持其中某些 batch size 而不會(huì)引起內(nèi)存不足 (OOM) 錯(cuò)誤。但當(dāng)使用 GPT-2 XL (1.5B) 時(shí)，即使 batch size 為 1，DDP 也會(huì)失敗并出現(xiàn) OOM 錯(cuò)誤。同時(shí)，我們看到，F(xiàn)SDP 可以支持以更大的 batch size 訓(xùn)練 GPT-2 Large 模型，同時(shí)它還可以使用較大的 batch size 訓(xùn)練 DDP 訓(xùn)練不了的 GPT-2 XL 模型。

硬件配置: 2 張 24GB 英偉達(dá) Titan RTX GPU。

GPT-2 Large 模型 (762M 參數(shù)) 的訓(xùn)練命令如下:

FSDP 運(yùn)行截屏:

表 1: GPT-2 Large (762M) 模型 FSDP 訓(xùn)練性能基準(zhǔn)測(cè)試

從表 1 中我們可以看到，相對(duì)于 DDP 而言，F(xiàn)SDP?支持更大的 batch size，在不使用和使用 CPU 卸載設(shè)置的情況下 FSDP 支持的最大 batch size 分別可達(dá) DDP 的?2 倍及 3 倍。從訓(xùn)練時(shí)間來(lái)看，混合精度的 DDP 最快，其后是分別使用 ZeRO 階段 2 和階段 3 的 FSDP。由于因果語(yǔ)言建模的任務(wù)的上下文序列長(zhǎng)度 (?--block_size?) 是固定的，因此 FSDP 在訓(xùn)練時(shí)間上加速還不是太高。對(duì)于動(dòng)態(tài) batch size 的應(yīng)用而言，支持更大 batch size 的 FSDP 可能會(huì)在訓(xùn)練時(shí)間方面有更大的加速。目前，F(xiàn)SDP 的混合精度支持在?transformers?上還存在一些 問(wèn)題。一旦問(wèn)題解決，訓(xùn)練時(shí)間將會(huì)進(jìn)一步顯著縮短。

使用 CPU 卸載來(lái)支持放不進(jìn) GPU 顯存的大模型訓(xùn)練

訓(xùn)練 GPT-2 XL (1.5B) 模型的命令如下:

表 2: GPT-2 XL (1.5B) 模型上的 FSDP 基準(zhǔn)測(cè)試

從表 2 中，我們可以觀察到 DDP (帶和不帶 fp16) 甚至在 batch size 為 1 的情況下就會(huì)出現(xiàn) CUDA OOM 錯(cuò)誤，從而無(wú)法運(yùn)行。而開啟了 ZeRO- 階段 3 的 FSDP 能夠以 batch size 為 5 (總 batch size = 10 (5??2) ) 在 2 個(gè) GPU 上運(yùn)行。當(dāng)使用 2 個(gè) GPU 時(shí)，開啟了 CPU 卸載的 FSDP 還能將最大 batch size 進(jìn)一步增加到每 GPU 14。開啟了 CPU 卸載的 FSDP 可以在單個(gè) GPU 上訓(xùn)練 GPT-2 1.5B 模型，batch size 為 10。這使得機(jī)器學(xué)習(xí)從業(yè)者能夠用最少的計(jì)算資源來(lái)訓(xùn)練大模型，從而助力大模型訓(xùn)練民主化。

Accelerate 的 FSDP 集成的功能和限制

下面，我們深入了解以下 Accelerate 對(duì) FSDP 的集成中，支持了那些功能，有什么已知的限制。

支持 FSDP 所需的 PyTorch 版本: PyTorch Nightly 或 1.12.0 之后的版本。

命令行支持的配置:

1. 分片策略: [1] FULL_SHARD, [2] SHARD_GRAD_OP

2. Min Num Params: FSDP 默認(rèn)自動(dòng)包裝的最小參數(shù)量。

3. Offload Params: 是否將參數(shù)和梯度卸載到 CPU。

如果想要對(duì)更多的控制參數(shù)進(jìn)行配置，用戶可以利用?FullyShardedDataParallelPlugin?，其可以指定?auto_wrap_policy?、?backward_prefetch?以及?ignored_modules?。

創(chuàng)建該類的實(shí)例后，用戶可以在創(chuàng)建 Accelerator 對(duì)象時(shí)把該實(shí)例傳進(jìn)去。

有關(guān)這些選項(xiàng)的更多信息，請(qǐng)參閱 PyTorch FullyShardedDataParallel 代碼。

接下來(lái)，我們體會(huì)下?min_num_params?配置的重要性。以下內(nèi)容摘自 [8]，它詳細(xì)說(shuō)明了 FSDP 自動(dòng)包裝策略的重要性。

(圖源: 鏈接)

當(dāng)使用?default_auto_wrap_policy?時(shí)，如果該層的參數(shù)量超過(guò)?min_num_params?，則該層將被包裝在一個(gè) FSDP 模塊中。官方有一個(gè)在 GLUE MRPC 任務(wù)上微調(diào) BERT-Large (330M) 模型的示例代碼，其完整地展示了如何正確使用 FSDP 功能，其中還包含了用于跟蹤峰值內(nèi)存使用情況的代碼。

fsdp_with_peak_mem_tracking.py

我們利用 Accelerate 的跟蹤功能來(lái)記錄訓(xùn)練和評(píng)估期間的峰值內(nèi)存使用情況以及模型準(zhǔn)確率指標(biāo)。下圖展示了 wandb 實(shí)驗(yàn)臺(tái) 頁(yè)面的截圖。

我們可以看到，DDP 占用的內(nèi)存是使用了自動(dòng)模型包裝功能的 FSDP 的兩倍。不帶自動(dòng)模型包裝的 FSDP 比帶自動(dòng)模型包裝的 FSDP 的內(nèi)存占用更多，但比 DDP 少得多。與?min_num_params=1M?時(shí)相比，?min_num_params=2k?時(shí)帶自動(dòng)模型包裝的 FSDP 占用的內(nèi)存略少。這凸顯了 FSDP 自動(dòng)模型包裝策略的重要性，用戶應(yīng)該調(diào)整?min_num_params?以找到能顯著節(jié)省內(nèi)存又不會(huì)導(dǎo)致大量通信開銷的設(shè)置。如 [8] 中所述，PyTorch 團(tuán)隊(duì)也在為此開發(fā)自動(dòng)配置調(diào)優(yōu)工具。

需要注意的一些事項(xiàng)

• PyTorch FSDP 會(huì)自動(dòng)對(duì)模型子模塊進(jìn)行包裝、將參數(shù)攤平并對(duì)其進(jìn)行原位分片。因此，在模型包裝之前創(chuàng)建的任何優(yōu)化器都會(huì)被破壞并導(dǎo)致更多的內(nèi)存占用。因此，強(qiáng)烈建議在對(duì)模型調(diào)用?prepare?方法后再創(chuàng)建優(yōu)化器，這樣效率會(huì)更高。對(duì)單模型而言，如果沒有按照順序調(diào)用的話，?Accelerate?會(huì)拋出以下告警信息，并自動(dòng)幫你包裝模型并創(chuàng)建優(yōu)化器。

  FSDP Warning: When using FSDP, it is efficient and recommended to call prepare for the model before creating the optimizer

即使如此，我們還是推薦用戶在使用 FSDP 時(shí)用以下方式顯式準(zhǔn)備模型和優(yōu)化器:

• 對(duì)單模型而言，如果你的模型有多組參數(shù)，而你想為它們?cè)O(shè)置不同優(yōu)化器超參。此時(shí)，如果你對(duì)整個(gè)模型統(tǒng)一調(diào)用?prepare?方法，這些參數(shù)的組別信息會(huì)丟失，你會(huì)看到如下告警信息:

  FSDP Warning: When using FSDP, several parameter groups will be conflated into a single one due to nested module wrapping and parameter flattening.

告警信息表明，在使用 FSDP 對(duì)模型進(jìn)行包裝后，之前創(chuàng)建的參數(shù)組信息丟失了。因?yàn)?FSDP 會(huì)將嵌套式的模塊參數(shù)攤平為一維數(shù)組 (一個(gè)數(shù)組可能包含多個(gè)子模塊的參數(shù))。舉個(gè)例子，下面是 GPU 0 上 FSDP 模型的有名稱的參數(shù) (當(dāng)使用 2 個(gè) GPU 時(shí)，F(xiàn)SDP 會(huì)把第一個(gè)分片的參數(shù)給 GPU 0， 因此其一維數(shù)組中大約會(huì)有 55M (110M / 2) 個(gè)參數(shù))。此時(shí)，如果我們?cè)?FSDP 包裝前將 BERT-Base 模型的 [bias, LayerNorm.weight] 參數(shù)的權(quán)重衰減設(shè)為 0，則在模型包裝后，該設(shè)置將無(wú)效。原因是，你可以看到下面這些字符串中均已不含這倆參數(shù)的名字，這倆參數(shù)已經(jīng)被并入了其他層。想要了解更多細(xì)節(jié)，可參閱本 問(wèn)題 (其中寫道:?原模型參數(shù)沒有 .grads 屬性意味著它們無(wú)法單獨(dú)被優(yōu)化器優(yōu)化 (這就是我們?yōu)槭裁床荒苤С謱?duì)多組參數(shù)設(shè)置不同的優(yōu)化器超參)?)。

• 如果是多模型情況，須在創(chuàng)建優(yōu)化器之前調(diào)用模型?prepare?方法，否則會(huì)拋出錯(cuò)誤。

• FSDP 目前不支持混合精度，我們正在等待 PyTorch 修復(fù)對(duì)其的支持。

工作原理 ??

(圖源: 鏈接)

上述工作流概述了 FSDP 的幕后流程。我們先來(lái)了解一下 DDP 是如何工作的，然后再看 FSDP 是如何改進(jìn)它的。在 DDP 中，每個(gè)工作進(jìn)程 (加速器 / GPU) 都會(huì)保留一份模型的所有參數(shù)、梯度和優(yōu)化器狀態(tài)的副本。每個(gè)工作進(jìn)程會(huì)獲取不同的數(shù)據(jù)，這些數(shù)據(jù)會(huì)經(jīng)過(guò)前向傳播，計(jì)算損失，然后再反向傳播以生成梯度。接著，執(zhí)行 all-reduce 操作，此時(shí)每個(gè)工作進(jìn)程從其余工作進(jìn)程獲取梯度并取平均。這樣一輪下來(lái)，每個(gè)工作進(jìn)程上的梯度都是相同的，且都是全局梯度，接著優(yōu)化器再用這些梯度來(lái)更新模型參數(shù)。我們可以看到，每個(gè) GPU 上都保留完整副本會(huì)消耗大量的顯存，這限制了該方法所能支持的 batch size 以及模型尺寸。

FSDP 通過(guò)讓各數(shù)據(jù)并行工作進(jìn)程分片存儲(chǔ)優(yōu)化器狀態(tài)、梯度和模型參數(shù)來(lái)解決這個(gè)問(wèn)題。進(jìn)一步地，還可以通過(guò)將這些張量卸載到 CPU 內(nèi)存來(lái)支持那些 GPU 顯存容納不下的大模型。在具體運(yùn)行時(shí)，與 DDP 類似，F(xiàn)SDP 的每個(gè)工作進(jìn)程獲取不同的數(shù)據(jù)。在前向傳播過(guò)程中，如果啟用了 CPU 卸載，則首先將本地分片的參數(shù)搬到 GPU/加速器。然后，每個(gè)工作進(jìn)程對(duì)給定的 FSDP 包裝模塊/層執(zhí)行 all-gather 操作以獲取所需的參數(shù)，執(zhí)行計(jì)算，然后釋放/清空其他工作進(jìn)程的參數(shù)分片。在對(duì)所有 FSDP 模塊全部執(zhí)行該操作后就是計(jì)算損失，然后是后向傳播。在后向傳播期間，再次執(zhí)行 all-gather 操作以獲取給定 FSDP 模塊所需的所有參數(shù)，執(zhí)行計(jì)算以獲得局部梯度，然后再次釋放其他工作進(jìn)程的分片。最后，使用 reduce-scatter 操作對(duì)局部梯度進(jìn)行平均并將相應(yīng)分片給對(duì)應(yīng)的工作進(jìn)程，該操作使得每個(gè)工作進(jìn)程都可以更新其本地分片的參數(shù)。如果啟用了 CPU 卸載的話，梯度會(huì)傳給 CPU，以便直接在 CPU 上更新參數(shù)。

如欲深入了解 PyTorch FSDP 工作原理以及相關(guān)實(shí)驗(yàn)及其結(jié)果，請(qǐng)參閱 [7,8,9]。

問(wèn)題

如果在 accelerate 中使用 PyTorch FSDP 時(shí)遇到任何問(wèn)題，請(qǐng)?zhí)峤恢?accelerate。

但如果你的問(wèn)題是跟 PyTorch FSDP 配置和部署有關(guān)的 - 你需要提交相應(yīng)的問(wèn)題至 PyTorch。

參考文獻(xiàn)

[1] Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers

[2] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

[3] DeepSpeed: Extreme-scale model training for everyone - Microsoft Research

[4] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

[5] Introducing GPipe, an Open Source Library for Efficiently Training Large-scale Neural Network Models

[6] Which hardware do you need to train a 176B parameters model?

[7] Introducing PyTorch Fully Sharded Data Parallel (FSDP) API | PyTorch

[8] Getting Started with Fully Sharded Data Parallel(FSDP) — PyTorch Tutorials 1.11.0+cu102 documentation

[9] Training a 1 Trillion Parameter Model With PyTorch Fully Sharded Data Parallel on AWS | by PyTorch | PyTorch | Mar, 2022 | Medium

[10] Fit More and Train Faster With ZeRO via DeepSpeed and FairScale

英文原文:?https://hf.co/blog/pytorch-fsdp

原文作者: Sourab Mangrulkar，Sylvain Gugger

譯者: Matrix Yao (姚偉峰)，英特爾深度學(xué)習(xí)工程師，工作方向?yàn)?transformer-family 模型在各模態(tài)數(shù)據(jù)上的應(yīng)用及大規(guī)模模型的訓(xùn)練推理。