背景

在最近幾年，雖然大規模預訓練模型已經越來越普遍，但是關于如何訓練這些模型的內容卻很少有人關注，一般都是一些財大氣粗的企業或實驗室來訓練大模型并發布，然后中小型企業以及高校來使用。即便如此也有一些門檻，受限于機器配置，可能效果更好的大模型并不能直接加載到顯卡中，或者是單機多卡希望可以通過分布式的方法進行微調。
目前訓練超大規模語言模型主要有兩條技術路線：TPU + XLA + TensorFlow 和 GPU + PyTorch + Megatron-LM + DeepSpeed。前者由Google主導，由于TPU和自家云平臺GCP深度綁定，對于非Google開發者來說， 只可遠觀而不可把玩，后者背后則有NVIDIA、Meta、微軟等大廠加持，社區氛圍活躍，也更受到群眾歡迎。
這篇文章介紹的ZeRO方法，就是可以讓大模型變得更加親民，不管是預訓練、微調還是部署，即使沒有A100、V100、P100之類的GPU服務器，也可以讓加載大模型不是夢。

硬件48 * Intel? Xeon? CPU E5-2650 v4 @ 2.20GHz
8 * 1080Ti GPUs（11GB顯存）

軟件	DeepSpeed
模型	GPT2-3.5B-chinese

本次實驗的模型是由IDEA發布的 Wenzhong2.0-GPT2-3.5B-chinese，這是一個基于解碼器結構的單向語言模型，使用100G中文常用數據，32個A100訓練了28個小時，是目前最大的開源GPT2中文大模型。接下來我們將在8塊1080Ti顯卡上微調這個大模型。

分布式模型訓練

數據并行（Data Parallel, DP）

將模型復制到多張卡上，然后將數據切分到多個GPU上，并行訓練，每一輪結束之后進行參數同步。

數據并行是最常見的并行形式，因為它很簡單。在數據并行訓練中，數據集被分割成幾個碎片，每個碎片被分配到一個設備上。這相當于沿批次維度對訓練過程進行并行化。每個設備將持有一個完整的模型副本，并在分配的數據集碎片上進行訓練。在反向傳播之后，模型的梯度將被全部減少，以便在不同設備上的模型參數能夠保持同步。


數據并行由于簡單易實現，應用最為廣泛，當然這不表示它沒有”缺點“，每張卡都存儲一個模型，此時顯存就成了模型規模的天花板。如果我們能減少模型訓練過程中的顯存占用，那不就可以訓練更大的模型了？一個簡單的觀察是，如果有2張卡，那么系統中就存在2份模型參數，如果有4張卡，那么系統中就存在4份模型參數，如果有N張卡，系統中就存在N份模型參數，但其實其中N-1份都是冗余的，我們有必要讓每張卡都存一個完整的模型嗎？系統中能否只有一個完整模型，每張卡都存 1/N 參數，卡數越多，每張卡的顯存占用越少，這樣越能訓練更大規模的模型。

優點：加快訓練速度
缺點：模型大小上限為單個GPU顯存

模型并行（Model Parallel, MP）

在數據并行訓練中，一個明顯的特點是每個 GPU 持有整個模型權重的副本。這就帶來了冗余問題。另一種并行模式是模型并行，即模型被分割并分布在一個設備陣列上。通常有兩種類型的并行：張量并行和流水線并行。張量并行是在一個操作中進行并行計算，如矩陣-矩陣乘法。流水線并行是在各層之間進行并行計算。因此，從另一個角度來看，張量并行可以被看作是層內并行，流水線并行可以被看作是層間并行。

張量并行（Tensor Parallel, TP）

將模型參數切分成分塊矩陣，分配到不同的GPU上，參數更新時再進行同步。

張量并行將一個大的張量進行切分，然后分配到不同的GPU上，每個GPU只處理張量的一部分，并且只有在一些需要所有張量計算的時候才進行聚合。


張量并行需要在不同設備之間傳遞結果，因此不建議在跨節點之間進行張量并行，除非網絡速度非常快。

優點：可以不受單張GPU顯存限制，訓練更大的模型。
缺點：計算/通信效率低。

流水線并行（Pipeline Parallel, PP）

將模型按層進行分組，然后分配到不同的GPU，在GPU之間進行前向傳播和反向傳播。

原生的流水線并行是指將模型層組劃分到到多個GPU上，然后只需要簡單地將數據從一個GPU移動到另一個GPU。流水線并行的實現相對比較簡單，只需要將不同的層通過.to()方法切換即可。


流水線并行的核心思想是，模型按層分割成若干塊，每塊都交給一個設備。在前向傳遞過程中，每個設備將中間的激活傳遞給下一個階段。在后向傳遞過程中，每個設備將輸入張量的梯度傳回給前一個流水線階段。這允許設備同時進行計算，并增加了訓練的吞吐量。流水線并行訓練的一個缺點是，會有一些設備參與計算的冒泡時間，導致計算資源的浪費。


優點：層內計算/通信效率增加
缺點：存在空閑等待時間

數據并行+流水線并行

有了這幾種并行方式之后，組合起來才能發揮更好的效果，數據并行和流水線并行相對來說比較簡單，先將數據分成兩份，放到兩個節點上，然后在每個節點的多個GPU上再執行流水線并行。

數據并行+流水線并行+模型并行

為了再進一步的提升效率，可以將三種并行方式都結合起來，稱之為3D并行，如下圖所示。


站在三維直角坐標系的角度看，x表示的是同一層內的張量并行，y軸表示的是同一節點內的流水線并行，z軸表示的是不同節點間的數據并行。

DeepSpeed

ZeRO（Zero Redundancy Optimizer）

類似于張量并行進行切分，支持多種offload技術。
目標：優化存儲效率的同時還能保持較高的計算和通信效率。

為了能夠在比較普通的機器上也能微調大模型，我們首先需要分析一下模型訓練過程中都有哪些部分需要消耗存儲空間。在進行深度學習訓練的時候，有4大部分的顯存開銷，分別是模型參數（Parameters），模型參數的梯度（Gradients），優化器狀態（Optimizer States）以及中間激活值（Intermediate Activations）。

優化模型占用空間

在訓練過程中當然模型占用的空間是最大的，但是現有的方法中，不管是數據并行DP還是模型并行MP都不能很好的解決。數據并行有很好的計算/通信效率，但是由于模型復制了多份，導致空間利用率很差，而模型并行雖然內存利用率高，但是由于對模型的進行了很精細的拆分，導致計算/通信效率很低。除此之外，所有這些方法都靜態保存了整個訓練過程中所需的所有模型參數，但實際上并不是整個訓練期間都需要這些內容。

ZeRO-DP

基于上述問題，提出了ZeRO-DP技術，即ZeRO驅動的數據并行，兼顧數據并行的計算/通信效率和模型并行的空間效率。首先ZeRO-DP會對模型狀態進行分區，避免了復制模型導致的冗余，然后在訓練期間使用動態通信調度保留數據并行的計算粒度和通信量，也能維持一個類似的計算/通信效率。
ZeRO-DP有三個優化階段：① 優化器狀態分區、② 梯度分區、③ 參數分區。

優化器狀態分區（Optimizer State Partitioning, Pos）：在與數據并行保持相同通信的情況下可以降低4倍空間占用；

① + 梯度分區（Gradient Partitioning, Pos+g）：在與數據并行保持相同通信量的情況下可以降低8倍空間占用；

① + ② + 參數分區（Parameter Partitioning, Pos+g+p）：空間占用減少量與GPU的個數呈線性關系 ，通信量增加50%。

模型參數（fp16）、模型梯度（fp16）和Adam狀態（fp32的模型參數備份，fp32的momentum和fp32的variance）。假設模型參數量 Φ ，則共需要 2Φ+2Φ+(4Φ+4Φ+4Φ)=4Φ+12Φ=16Φ 字節存儲。

優化剩余空間

優化了模型的空間利用率之后，接下來就要優化激活值、臨時緩沖區和不可用空間碎片。

ZeRO-R

對于激活值，需要在前向傳播的過程中存儲，在反向傳播的過程中利用。ZeRO-R激活分區識別和刪除模型并行中冗余的激活函數值進行優化，并且在適當的時候會向CPU進行換入換出。

模型訓練過程中經常會創建一些大小不等的臨時緩沖區，比如對梯度進行AllReduce啥的，ZeRO-R預先創建一個固定的緩沖區，訓練過程中不再動態創建，如果要傳輸的數據較小，則多組數據bucket后再一次性傳輸，以達到內存和計算效率的平衡。

在訓練過程中，不同張量的生命周期不同，會導致產生空間碎片，進而可能會導致連續空間分配失敗。ZeRO-R根據張量的不同生命周期主動管理空間，防止產生空間碎片。

ZeRO與模型并行

既然ZeRO解決了數據并行中空間利用率低的問題，那么再進一步，怎么與模型并行進行配合使用呢？用了ZeRO之后，模型并行就僅僅被用在訓練超大模型，ZeRO-DP在減少每個設備的空間占用方面跟模型并行一樣有效，或者當模型并行不能均勻的劃分模型時有效。
在分布式模型訓練過程中，數據并行非常易于使用，但是模型并行的門檻相對比較高，需要AI工程師手動調整模型，甚至還需要系統開發人員自定義分布式算子，當然也有一些已有的工作，比如Megatron-LM。

混合精度訓練（Mixed-Precision Training）

FP16

正常模型所使用的參數是Float浮點數類型，長度為4個字節，也就是32位。在一些不太需要高精度計算的應用中，比如圖像處理和神經網絡中，32位的空間有一些浪費，因此就又出現了一種新的數據類型，半精度浮點數，使用16位（2字節）來存儲浮點值，簡稱FP16。
FP16的格式包括三部分：① 1 bit的符號位，② 5 bit的指數位寬，③ 11 bit的尾數精度（10位顯式存儲，1位隱式存儲）。

0 01111 0000000000 = 1 0 01111 0000000001 = 1 + 2?10 = 1.0009765625 （1之后的最接近的數）0 11110 1111111111 = 65504 （max half precision） 0 11111 0000000000 = infinity 1 11111 0000000000 = ?infinity

但是在大型語言模型中直接使用FP16會有一些問題，如下圖所示，這是huggingface訓練一個104B的模型時的loss，可以發現非常的不穩定，除此之外Facebook最近發布的175B的OPT模型也報告了相同的問題。


問題在于FP16的指數位寬只有5 bit，因此可能會出現溢出的問題，它能表示的最大整數就是65504，也就是說一旦權重超過這個值就會溢出。

BF16

為了解決FP16的問題，Google的人工智能研究小組就開發了一種新的浮點數格式——BF16（Brain Floating Point），用于降低存儲需求，提高機器學習算法的計算速度。
BF16的格式包括三部分：① 1 bit的符號位，② 8 bit的指數位寬，③ 8 bit的尾數精度（7位顯式存儲，1位隱式存儲）。

3f80 = 0 01111111 0000000 = 1 c000 = 1 10000000 0000000 = -27f7f = 0 11111110 1111111 = (2^8 ? 1) × 2^{?7} × 2^{127} ≈ 3.38953139 × 10^{38}（bfloat16 精度的最大有限正值）

---

模型參數不管是使用BF16還是FP16，在優化器中更新的權重參數必須是用FP32的形式保存的，因此16位的浮點數僅用于計算。

ZeRO-Offload

ZeRO說到底是一種數據并行方案，可是很多人只有幾張甚至一張卡，顯存加起來都不夠，那怎么辦呢？在操作系統中，當內存不足時，可以選擇一些頁面進行換入換出，為新的數據騰出空間。類比一下，既然是因為顯存不足導致一張卡訓練不了大模型，那么ZeRO-Offload的想法就是：顯存不足，內存來補。
在一個典型的服務器上，CPU 可以輕松擁有幾百GB的內存，而每個 GPU 通常只有16或32GB的內存。相比于昂貴的顯存，內存比較廉價，之前的很多工作都是聚焦在內存顯存的換入換出，并沒有用到CPU的計算能力，也沒有考慮到多卡的場景。ZeRO-Offload則是將訓練階段的某些模型狀態從GPU和顯存卸載到CPU和內存。
當然ZeRO-Offload并不希望為了最小化顯存占用而犧牲計算效率， 否則的話還不如直接使用CPU和內存，因為即使將部分GPU的計算和顯存卸載到CPU和內存，肯定要涉及到GPU和CPU、顯存和內存的通信，而通信成本一般是非常高的，此外GPU的計算效率比CPU的計算效率高了好幾個數量積，因此也不能讓CPU參與過多的計算。
在ZeRO-Offload的策略中，將模型訓練過程看作數據流圖，用圓形節點表示模型狀態，比如參數、梯度和優化器狀態，用矩形節點表示計算操作，比如前向傳播、反向傳播和參數更新，邊表示數據流向。下圖是某一層的一次迭代過程（iteration/step），使用了混合精讀訓練，前向計算（FWD）需要用到上一次的激活值（activation）和本層的參數（parameter），反向傳播（BWD）也需要用到激活值和參數計算梯度，


如果用Adam優化器進行參數更新（Param update），邊添加權重，物理含義是數據量大小（單位是字節），假設模型參數量是 M ，在混合精度訓練的前提下，邊的權重要么是2M（fp16），要么是4M（fp32）。


**現在要做的就是沿著邊把數據流圖切分為兩部分，分布對應GPU和CPU，**計算節點（矩形節點）落在哪個設備，哪個設備就執行計算，數據節點（圓形）落在哪個設備，哪個設備就負責存儲，將被切分的邊權重加起來，就是CPU和GPU的通信數據量。
ZeRO-Offload的切分思路是：圖中有四個計算類節點：FWD、BWD、Param update和float2half，前兩個計算復雜度大致是 O(MB) ， B 是batch size，后兩個計算復雜度是 O(M) 。為了不降低計算效率，將前兩個節點放在GPU，后兩個節點不但計算量小還需要和Adam狀態打交道，所以放在CPU上，Adam狀態自然也放在內存中，為了簡化數據圖，將前兩個節點融合成一個節點FWD-BWD Super Node，將后兩個節點融合成一個節點Update Super Node。如下圖右邊所示，沿著gradient 16和parameter 16兩條邊切分。


現在的計算流程是，在GPU上面進行前向和后向計算，將梯度傳給CPU，進行參數更新，再將更新后的參數傳給GPU。為了提高效率，可以將計算和通信并行起來，GPU在反向傳播階段，可以待梯度值填滿bucket后，一邊計算新的梯度一邊將bucket傳輸給CPU，當反向傳播結束，CPU基本上已經有最新的梯度值了，同樣的，CPU在參數更新時也同步將已經計算好的參數傳給GPU，如下圖所示。


到目前為止還都是單卡的場景，在多卡場景中，ZeRO-Offload可以利用ZeRO-2，將優化器狀態和梯度進行切分，每張卡只保留$\frac{1}{N}$，結合上ZeRO-Offload同樣是將這$\frac{1}{N}$的優化器狀態和梯度卸載到內存，在CPU上進行參數更新。在多卡場景，利用CPU多核并行計算，每張卡至少對應一個CPU進程，由這個進程負責進行局部參數更新。
并且CPU和GPU的通信量和 N 無關，因為傳輸的是fp16 gradient和fp16 parameter，總的傳輸量是固定的，由于利用多核并行計算，每個CPU進程只負責 1N 的計算，反而隨著卡數增加節省了CPU計算時間。

ZeRO-Infinity

從GPT-1到GPT-3，兩年時間內模型參數0.1B增加到175B，而同期，NVIDIA的顯卡是從V100的32GB顯存增加A100的80GB，顯然，顯寸的提升速度遠遠趕不上模型模型增長的速度，這就是內存墻問題。


所謂內存墻，其實就是硬件之間的通信會遇到瓶頸，不管是GPU之間還是GPU和CPU之間，通信的效率遠低于計算的效率。在過去20年中，硬件的峰值計算能力增加了90,000倍，但是內存/硬件互連帶寬卻只是提高了30倍。


ZeRO-Infinity這篇論文就是用來解決內存墻的問題，但是它并不是打通了內存墻，而是通過NVMe。NVMe是一種接口規范，最開始是用在SSD固態硬盤上的，之前的SSD用的都是SATA接口規范，數據傳輸速率最大6Gbps，而NVMe就是一種新的技術，NVMe固態硬盤擁有高達20Gbps的理論傳輸速度。
ZeRO-Infinity就是依靠 CPU 甚至是 NVMe 磁盤來訓練大型模型。主要的想法是，在不使用張量時，將其卸載回 CPU 內存或 NVMe 磁盤。通過使用異構系統架構，有可能在一臺機器上容納一個巨大的模型。

實踐

目標：微調 Wenzhong2.0-GPT2-3.5B-chinese

估計使用ZeRO-3配置需要的參數、優化器狀態和梯度的使用需求

python -c 'from transformers import GPT2Model; \ from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \ model = GPT2Model.from_pretrained("IDEA-CCNL/Wenzhong2.0-GPT2-3.5B-chinese"); \ estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)' Some weights of the model checkpoint at IDEA-CCNL/Wenzhong2.0-GPT2-3.5B-chinese were not used when initializing GPT2Model: ['lm_head.weight'] - This IS expected if you are initializing GPT2Model from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model). - This IS NOT expected if you are initializing GPT2Model from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model). Estimated memory needed for params, optim states and gradients for a: HW: Setup with 1 node, 8 GPUs per node. SW: Model with 3556M total params, 154M largest layer params.per CPU | per GPU | Options89.42GB | 0.58GB | offload_param=cpu , offload_optimizer=cpu , zero_init=1158.98GB | 0.58GB | offload_param=cpu , offload_optimizer=cpu , zero_init=079.49GB | 1.40GB | offload_param=none, offload_optimizer=cpu , zero_init=1158.98GB | 1.40GB | offload_param=none, offload_optimizer=cpu , zero_init=06.91GB | 8.03GB | offload_param=none, offload_optimizer=none, zero_init=1158.98GB | 8.03GB | offload_param=none, offload_optimizer=none, zero_init=0

實際訓練中每塊顯卡的占用情況

實驗

配置文件

#!/bin/bashset -x -eecho "START TIME: $(date)" MICRO_BATCH_SIZE=1 ROOT_DIR=$(pwd)ZERO_STAGE=3config_json="$ROOT_DIR/training_config.json"cat <<EOT >$config_json {"train_micro_batch_size_per_gpu": $MICRO_BATCH_SIZE,"steps_per_print": 1000,"gradient_clipping": 1,"zero_optimization": {"stage": ${ZERO_STAGE},"allgather_partitions": false,"allgather_bucket_size": 2e8,"overlap_comm": true,"reduce_scatter": true,"reduce_bucket_size": 2e8,"contiguous_gradients": true,"offload_optimizer": {"device": "cpu","pin_memory": true},"offload_param": {"device": "cpu","pin_memory": true},"stage3_max_live_parameters" : 2e8,"stage3_max_reuse_distance" : 2e8,"stage3_prefetch_bucket_size": 2e8,"stage3_param_persistence_threshold": 2e8,"sub_group_size" : 2e8,"round_robin_gradients": true},"bf16": {"enabled": true},"optimizer": {"type": "Adam","params": {"lr": 1e-5,"betas": [0.9,0.95],"eps": 1e-8,"weight_decay": 1e-2}},"scheduler": {"type": "WarmupLR","params":{"warmup_min_lr": 5e-6,"warmup_max_lr": 1e-5}} } EOTexport PL_DEEPSPEED_CONFIG_PATH=$config_json TRAINER_ARGS="--max_epochs 1 \--num_nodes 2 \--gpus 8 \--strategy deepspeed_stage_${ZERO_STAGE}_offload \--default_root_dir $ROOT_DIR \--dirpath $ROOT_DIR/ckpt \--save_top_k 3 \--monitor train_loss \--mode min \--save_last \ "DATA_DIR=/home/liuzhaofeng/nlg_pipeline/gpt2/dialog/datasets DATA_ARGS="--data_dir $DATA_DIR \--max_seq_length 64 \--train_batchsize $MICRO_BATCH_SIZE \--valid_batchsize $MICRO_BATCH_SIZE \--train_data test_train.txt \--valid_data test.txt \--test_data test.txt "PRETRAINED_MODEL_PATH="IDEA-CCNL/Wenzhong2.0-GPT2-3.5B-chinese" MODEL_ARGS="--pretrained_model_path ${PRETRAINED_MODEL_PATH} \--output_save_path $ROOT_DIR/predict.json \--learning_rate 1e-4 \--weight_decay 0.1 \--warmup 0.01 \ "SCRIPTS_PATH=${ROOT_DIR}/finetune_gpt2.pyexport CMD=" \$SCRIPTS_PATH \$TRAINER_ARGS \$MODEL_ARGS \$DATA_ARGS \ "export NCCL_IB_DISABLE=1 export NCCL_DEBUG_SUBSYS=INIT,P2P,SHM,NET,GRAPH,ENV export NCCL_SOCKET_IFNAME=enp129s0f1torchrun ${CMD}

壓力測試

BATCH_SIZEmax_seq_length內存占用顯存占用

1	64	153G	5669M
	128	153G	6041M
	256	154G	7139M
	512	154G	9383M
2	64	153G	6367M
	128	153G	7439M
	256	153G	9651M
	512	OOM
3	64	167G	6431M
	128	166G	8273M
	256	163G	11103M

總結

以上是生活随笔為你收集整理的震惊！我竟然在1080Ti上加载了一个35亿参数的模型（ZeRO, Zero Redundancy Optimizer）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。