本文，我們將瞭解如何基於 PyTorch 最新的 完全分片資料並行 (Fully Sharded Data Parallel，FSDP) 功能用 Accelerate 庫來訓練大模型。
動機 🤗
隨著機器學習 (ML) 模型的規模、大小和引數量的不斷增加，ML 從業者發現在自己的硬體上訓練甚至載入如此大的模型變得越來越難。 一方面，人們發現大模型與較小的模型相比，學習速度更快 (資料和計算效率更高) 且會有顯著的提升 [1]; 另一方面，在大多數硬體上訓練此類模型變得令人望而卻步。
分散式訓練是訓練這些機器學習大模型的關鍵。 大規模分散式訓練 領域最近取得了不少重大進展，我們將其中一些最突出的進展總結如下:
-
使用 ZeRO 資料並行 - 零冗餘最佳化器 [2]
-
階段 1: 跨資料並行程式 / GPU 對
最佳化器狀態進行分片 -
階段 2: 跨資料並行程式/ GPU 對
最佳化器狀態 + 梯度進行分片 -
階段 3: 跨資料並行程式 / GPU 對
最佳化器狀態 + 梯度 + 模型引數進行分片 -
CPU 解除安裝: 進一步將 ZeRO 階段 2 的
最佳化器狀態 + 梯度解除安裝到 CPU 上 [3] -
張量並行 [4]: 模型並行的一種形式，透過對各層引數進行精巧的跨加速器 / GPU 分片，在實現平行計算的同時避免了昂貴的通訊同步開銷。
-
流水線並行 [5]: 模型並行的另一種形式，其將模型的不同層放在不同的加速器 / GPU 上，並利用流水線來保持所有加速器同時執行。舉個例子，在第 2 個加速器 / GPU 對第 1 個 micro batch 進行計算的同時，第 1 個加速器 / GPU 對第 2 個 micro batch 進行計算。
-
3D 並行 [3]: 採用
ZeRO 資料並行 + 張量並行 + 流水線並行的方式來訓練數百億引數的大模型。例如，BigScience 176B 語言模型就採用了該並行方式 [6]。
本文我們主要關注 ZeRO 資料並行，更具體地講是 PyTorch 最新的 完全分片資料並行 (Fully Sharded Data Parallel，FSDP) 功能。 DeepSpeed 和 FairScale 實現了 ZeRO 論文的核心思想。我們已經將其整合到了 transformers 的 Trainer 中，詳見博文 透過 DeepSpeed 和 FairScale 使用 ZeRO 進行更大更快的訓練[10]。最近，PyTorch 已正式將 Fairscale FSDP 整合進其 Distributed 模組中，並增加了更多的最佳化。
Accelerate 🚀: 無需更改任何程式碼即可使用 PyTorch FSDP
我們以基於 GPT-2 的 Large (762M) 和 XL (1.5B) 模型的因果語言建模任務為例。
以下是預訓練 GPT-2 模型的程式碼。其與 此處 的官方因果語言建模示例相似，僅增加了 2 個引數 n_train (2000) 和 n_val (500) 以防止對整個資料集進行預處理/訓練，從而支援更快地進行概念驗證。
執行 accelerate config 命令後得到的 FSDP 配置示例如下:
compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: FSDP
fsdp_config:
min_num_params: 2000
offload_params: false
sharding_strategy: 1
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 2
use_cpu: false
多 GPU FSDP
本文我們使用單節點多 GPU 上作為實驗平臺。我們比較了分散式資料並行 (DDP) 和 FSDP 在各種不同配置下的效能。我們可以看到，對 GPT-2 Large(762M) 模型而言，DDP 尚能夠支援其中某些 batch size 而不會引起記憶體不足 (OOM) 錯誤。但當使用 GPT-2 XL (1.5B) 時，即使 batch size 為 1，DDP 也會失敗並出現 OOM 錯誤。同時，我們看到，FSDP 可以支援以更大的 batch size 訓練 GPT-2 Large 模型，同時它還可以使用較大的 batch size 訓練 DDP 訓練不了的 GPT-2 XL 模型。
硬體配置: 2 張 24GB 英偉達 Titan RTX GPU。
GPT-2 Large 模型 (762M 引數) 的訓練命令如下:
export BS=#`try with different batch sizes till you don't get OOM error,
#i.e., start with larger batch size and go on decreasing till it fits on GPU`
time accelerate launch run_clm_no_trainer.py \
--model_name_or_path gpt2-large \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--per_device_train_batch_size $BS
--per_device_eval_batch_size $BS
--num_train_epochs 1
--block_size 12
FSDP 執行截圖:
| 並行方法 | 最大 Batch Size ($BS) | 大致訓練時間 (分鐘) | 備註 |
|---|---|---|---|
| DDP | 7 | 15 | |
| DDP + FP16 | 7 | 8 | |
| FSDP (配置: SHARD_GRAD_OP) | 11 | 11 | |
| FSDP (配置: min_num_params = 1M + FULL_SHARD) | 15 | 12 | |
| FSDP (配置: min_num_params = 2K + FULL_SHARD) | 15 | 13 | |
| FSDP (配置: min_num_params = 1M + FULL_SHARD + CPU 解除安裝) | 20 | 23 | |
| FSDP (配置: min_num_params = 2K + FULL_SHARD + CPU 解除安裝) | 22 | 24 |
表 1: GPT-2 Large (762M) 模型 FSDP 訓練效能基準測試
從表 1 中我們可以看到，相對於 DDP 而言，FSDP 支援更大的 batch size，在不使用和使用 CPU 解除安裝設定的情況下 FSDP 支援的最大 batch size 分別可達 DDP 的 2 倍及 3 倍。從訓練時間來看，混合精度的 DDP 最快，其後是分別使用 ZeRO 階段 2 和階段 3 的 FSDP。由於因果語言建模的任務的上下文序列長度 ( --block_size ) 是固定的，因此 FSDP 在訓練時間上加速還不是太高。對於動態 batch size 的應用而言，支援更大 batch size 的 FSDP 可能會在訓練時間方面有更大的加速。目前，FSDP 的混合精度支援在 transformers 上還存在一些 問題。一旦問題解決，訓練時間將會進一步顯著縮短。
使用 CPU 解除安裝來支援放不進 GPU 視訊記憶體的大模型訓練
訓練 GPT-2 XL (1.5B) 模型的命令如下:
export BS=#`try with different batch sizes till you don't get OOM error,
#i.e., start with larger batch size and go on decreasing till it fits on GPU`
time accelerate launch run_clm_no_trainer.py \
--model_name_or_path gpt2-xl \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--per_device_train_batch_size $BS
--per_device_eval_batch_size $BS
--num_train_epochs 1
--block_size 12
| 並行方法 | 最大 Batch Size ($BS) | GPU 數 | 大致訓練時間 (小時) | 備註 |
|---|---|---|---|---|
| DDP | 1 | 1 | NA | OOM Error RuntimeError: CUDA out of memory. Tried to allocate 40.00 MiB (GPU 0; 23.65 GiB total capacity; 22.27 GiB already allocated; 20.31 MiB free; 22.76 GiB reserved in total by PyTorch) |
| DDP | 1 | 2 | NA | OOM Error RuntimeError: CUDA out of memory. Tried to allocate 40.00 MiB (GPU 0; 23.65 GiB total capacity; 22.27 GiB already allocated; 20.31 MiB free; 22.76 GiB reserved in total by PyTorch) |
| DDP + FP16 | 1 | 1 | NA | OOM Error RuntimeError: CUDA out of memory. Tried to allocate 40.00 MiB (GPU 0; 23.65 GiB total capacity; 22.27 GiB already allocated; 20.31 MiB free; 22.76 GiB reserved in total by PyTorch) |
| FSDP (配置: min_num_params = 2K) | 5 | 2 | 0.6 | |
| FSDP (配置: min_num_params = 2K + CPU 解除安裝) | 10 | 1 | 3 | |
| FSDP (配置: min_num_params = 2K + CPU 解除安裝) | 14 | 2 | 1.16 |
表 2: GPT-2 XL (1.5B) 模型上的 FSDP 基準測試
從表 2 中，我們可以觀察到 DDP (帶和不帶 fp16) 甚至在 batch size 為 1 的情況下就會出現 CUDA OOM 錯誤，從而無法執行。而開啟了 ZeRO- 階段 3 的 FSDP 能夠以 batch size 為 5 (總 batch size = 10 (5 \(\times\) 2) ) 在 2 個 GPU 上執行。當使用 2 個 GPU 時，開啟了 CPU 解除安裝的 FSDP 還能將最大 batch size 進一步增加到每 GPU 14。 開啟了 CPU 解除安裝的 FSDP 可以在單個 GPU 上訓練 GPT-2 1.5B 模型，batch size 為 10。這使得機器學習從業者能夠用最少的計算資源來訓練大模型，從而助力大模型訓練民主化。
Accelerate 的 FSDP 整合的功能和限制
下面，我們深入瞭解以下 Accelerate 對 FSDP 的整合中，支援了那些功能，有什麼已知的限制。
支援 FSDP 所需的 PyTorch 版本: PyTorch Nightly 或 1.12.0 之後的版本。
命令列支援的配置:
- 分片策略: [1] FULL_SHARD, [2] SHARD_GRAD_OP
- Min Num Params: FSDP 預設自動包裝的最小引數量。
- Offload Params: 是否將引數和梯度解除安裝到 CPU。
如果想要對更多的控制引數進行配置，使用者可以利用 FullyShardedDataParallelPlugin ，其可以指定 auto_wrap_policy 、 backward_prefetch 以及 ignored_modules 。
建立該類的例項後，使用者可以在建立 Accelerator 物件時把該例項傳進去。
有關這些選項的更多資訊，請參閱 PyTorch FullyShardedDataParallel 程式碼。
接下來，我們體會下 min_num_params 配置的重要性。以下內容摘自 [8]，它詳細說明了 FSDP 自動包裝策略的重要性。
(圖源: 連結)
當使用 default_auto_wrap_policy 時，如果該層的引數量超過 min_num_params ，則該層將被包裝在一個 FSDP 模組中。官方有一個在 GLUE MRPC 任務上微調 BERT-Large (330M) 模型的示例程式碼，其完整地展示瞭如何正確使用 FSDP 功能，其中還包含了用於跟蹤峰值記憶體使用情況的程式碼。
fsdp_with_peak_mem_tracking.py
我們利用 Accelerate 的跟蹤功能來記錄訓練和評估期間的峰值記憶體使用情況以及模型準確率指標。下圖展示了 wandb 實驗臺 頁面的截圖。
我們可以看到，DDP 佔用的記憶體是使用了自動模型包裝功能的 FSDP 的兩倍。不帶自動模型包裝的 FSDP 比帶自動模型包裝的 FSDP 的記憶體佔用更多，但比 DDP 少得多。與 min_num_params=1M 時相比， min_num_params=2k 時帶自動模型包裝的 FSDP 佔用的記憶體略少。這凸顯了 FSDP 自動模型包裝策略的重要性，使用者應該調整 min_num_params 以找到能顯著節省記憶體又不會導致大量通訊開銷的設定。如 [8] 中所述，PyTorch 團隊也在為此開發自動配置調優工具。
需要注意的一些事項
-
PyTorch FSDP 會自動對模型子模組進行包裝、將引數攤平並對其進行原位分片。因此，在模型包裝之前建立的任何最佳化器都會被破壞並導致更多的記憶體佔用。因此，強烈建議在對模型呼叫
prepare方法後再建立最佳化器，這樣效率會更高。對單模型而言，如果沒有按照順序呼叫的話，Accelerate會丟擲以下告警資訊，並自動幫你包裝模型並建立最佳化器。FSDP Warning: When using FSDP, it is efficient and recommended to call prepare for the model before creating the optimizer
即使如此，我們還是推薦使用者在使用 FSDP 時用以下方式顯式準備模型和最佳化器:
model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+ model = accelerator.prepare(model)
optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
- model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(model,
- optimizer, train_dataloader, eval_dataloader, lr_scheduler
- )
+ optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+ optimizer, train_dataloader, eval_dataloader, lr_scheduler
+ )
-
對單模型而言，如果你的模型有多組引數，而你想為它們設定不同最佳化器超參。此時，如果你對整個模型統一呼叫
prepare方法，這些引數的組別資訊會丟失，你會看到如下告警資訊:FSDP Warning: When using FSDP, several parameter groups will be conflated into a single one due to nested module wrapping and parameter flattening.
告警資訊表明，在使用 FSDP 對模型進行包裝後，之前建立的引數組資訊丟失了。因為 FSDP 會將巢狀式的模組引數攤平為一維陣列 (一個陣列可能包含多個子模組的引數)。舉個例子，下面是 GPU 0 上 FSDP 模型的有名稱的引數 (當使用 2 個 GPU 時，FSDP 會把第一個分片的引數給 GPU 0， 因此其一維陣列中大約會有 55M (110M / 2) 個引數)。此時，如果我們在 FSDP 包裝前將 BERT-Base 模型的 [bias, LayerNorm.weight] 引數的權重衰減設為 0，則在模型包裝後，該設定將無效。原因是，你可以看到下面這些字串中均已不含這倆引數的名字，這倆引數已經被併入了其他層。想要了解更多細節，可參閱本 問題 (其中寫道: 原模型引數沒有 .grads 屬性意味著它們無法單獨被最佳化器最佳化 (這就是我們為什麼不能支援對多組引數設定不同的最佳化器超參) )。
{
'_fsdp_wrapped_module.flat_param': torch.Size([494209]),
'_fsdp_wrapped_module._fpw_module.bert.embeddings.word_embeddings._fsdp_wrapped_module.flat_param': torch.Size([11720448]),
'_fsdp_wrapped_module._fpw_module.bert.encoder._fsdp_wrapped_module.flat_param': torch.Size([42527232])
}
- 如果是多模型情況，須在建立最佳化器之前呼叫模型
prepare方法，否則會丟擲錯誤。 FSDP 目前不支援混合精度，我們正在等待 PyTorch 修復對其的支援。
工作原理 📝
(圖源: 連結)
上述工作流概述了 FSDP 的幕後流程。我們先來了解一下 DDP 是如何工作的，然後再看 FSDP 是如何改進它的。在 DDP 中，每個工作程式 (加速器 / GPU) 都會保留一份模型的所有引數、梯度和最佳化器狀態的副本。每個工作程式會獲取不同的資料，這些資料會經過前向傳播，計算損失，然後再反向傳播以生成梯度。接著，執行 all-reduce 操作，此時每個工作程式從其餘工作程式獲取梯度並取平均。這樣一輪下來，每個工作程式上的梯度都是相同的，且都是全域性梯度，接著最佳化器再用這些梯度來更新模型引數。我們可以看到，每個 GPU 上都保留完整副本會消耗大量的視訊記憶體，這限制了該方法所能支援的 batch size 以及模型尺寸。
FSDP 透過讓各資料並行工作程式分片儲存最佳化器狀態、梯度和模型引數來解決這個問題。進一步地，還可以透過將這些張量解除安裝到 CPU 記憶體來支援那些 GPU 視訊記憶體容納不下的大模型。在具體執行時，與 DDP 類似，FSDP 的每個工作程式獲取不同的資料。在前向傳播過程中，如果啟用了 CPU 解除安裝，則首先將本地分片的引數搬到 GPU/加速器。然後，每個工作程式對給定的 FSDP 包裝模組/層執行 all-gather 操作以獲取所需的引數，執行計算，然後釋放/清空其他工作程式的引數分片。在對所有 FSDP 模組全部執行該操作後就是計算損失，然後是後向傳播。在後向傳播期間，再次執行 all-gather 操作以獲取給定 FSDP 模組所需的所有引數，執行計算以獲得區域性梯度，然後再次釋放其他工作程式的分片。最後，使用 reduce-scatter 操作對區域性梯度進行平均並將相應分片給對應的工作程式，該操作使得每個工作程式都可以更新其本地分片的引數。如果啟用了 CPU 解除安裝的話，梯度會傳給 CPU，以便直接在 CPU 上更新引數。
如欲深入瞭解 PyTorch FSDP 工作原理以及相關實驗及其結果，請參閱 [7,8,9]。
問題
如果在 accelerate 中使用 PyTorch FSDP 時遇到任何問題，請提交至 accelerate。
但如果你的問題是跟 PyTorch FSDP 配置和部署有關的 - 你需要提交相應的問題至 PyTorch。
參考文獻
[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
[10] Fit More and Train Faster With ZeRO via DeepSpeed and FairScale
英文原文: https://hf.co/blog/pytorch-fsdp
原文作者: Sourab Mangrulkar，Sylvain Gugger
譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。