大模型訓練（5）：Zero Redundancy Optimizer（ZeRO零冗余優化器）

0 英文縮寫

Large Language Model（LLM）大型語言模型
Data Parallelism（DP）數據并行
Distributed Data Parallelism（DDP）分布式數據并行
Zero Redundancy Optimizer（ZeRO）零冗余優化器
$P_{os}$ ：Partition Optimizer States 對優化器狀態進行切片
$P_G$ ：Partition Gradients 對梯度進行切片
$P_P$ ：Partition Parameters 對參數進行切片

1 背景

1.1 大模型訓練現狀

LLM在訓練時往往需要大量內存來存儲中間激活、權重等參數，百億模型甚至無法在單個GPU上進行訓練，使得模型訓練在某些情況下非常低效和不可能。這就需要進行多卡，或者多節點分布式訓練。在大規模深度學習模型訓練中有如下幾種主要的范式：

流水線并行
張量并行
數據并行
模型并行

目前訓練超大規模語言模型技術路線：

GPU（算力核心設備）
PyTorch（框架，代碼到驅動的鏈接者映射者）
Megatron-LM（實現模型并行、張量并行、流水線并行，樸素數據并行）
DeepSpeed（實現數據并行）

1.2 幾種并行原理簡述

PP并行的原理：當模型太大，一塊GPU放不下時，流水線并行將模型的不同階段（一般為不同層）放到不同的GPU上，通過將mini-batch切割成更細粒度的micro-batch，實現對訓練數據的流水線處理，提升GPU計算通訊比。同時通過re-materialization機制降低顯存消耗。
DP并行的優勢（更強易用性應用更加廣泛）：在實際應用中，流水線并行并不特別流行，主要原因是模型能否均勻切割，影響了整體計算效率，這就需要算法工程師做手調。因此，來介紹一種應用最廣泛，最易于理解的并行范式：數據并行。
數據并行的核心思想是：在各個GPU上都拷貝一份完整模型，各自吃一份數據，算一份梯度，最后對梯度進行累加來更新整體模型。理念不復雜，但到了大模型場景，巨大的存儲和GPU間的通訊量，就是系統設計要考慮的重點了。在本文以及后續文章中，我們將遞進介紹三種主流數據并行的實現方式：
- DP：最早的數據并行模式，一般采用參數服務器這一編程框架。實際中多用于單機多卡
- DDP：分布式數據并行，采用Ring AllReduce的通訊方式，實際中多用于多機場景
- ZeRO：（本文主要的核心內容）零冗余優化器。

2 DeepSpeed概述

2.1 基本理念與達成路徑

DeepSpeed是由Microsoft提供的分布式訓練工具，與其他框架相比，優勢在支持更大規模的模型和提供更多的優化策略和工具（例如 ZeRO 和 Offload 等）。看一下官網對于這個理念的描述：

Why would you want to use DeepSpeed with just one GPU?

It has a ZeRO-offload feature which can delegate some computations and memory to the host’s CPU and RAM, and thus leave more GPU resources for model’s needs - e.g. larger batch size, or enabling a fitting of a very big model which normally won’t fit.
It provides a smart GPU memory management system, that minimizes memory fragmentation, which again allows you to fit bigger models and data batches.

具體點說，DeepSpeed將當前時刻，訓練模型用不到的參數（包括模型參數、optimizer、梯度等），不計算或者緩存到CPU中，等到要用到了，再從其他GPU上拿或者從CPU挪到GPU。越多的參數被卸載掉，GPU的負擔就越小；但隨之的代價就是，更為頻繁的GPU-GPU或者CPU-GPU交互，極大增加了訓練推理的時間開銷。因此，DeepSpeed使用的一個核心要義是，時間開銷和顯存占用的權衡。

2.2 基本概念

在分布式計算環境中，需要理解幾個非常基礎的概念：

節點編號（node_rank）：分配給系統中每個節點的唯一標識符，用于區分不同計算機之間的通信。
全局進程編號（rank）：分配給整個系統中的每個進程的唯一標識符，用于區分不同進程之間的通信。
局部進程編號（local_rank）：分配給單個節點內的每個進程的唯一標識符，用于區分同一節點內的不同進程之間的通信。
全局總進程數（word_size）：在整個系統中運行的所有進程的總數，用于確定可以并行完成多少工作以及需要完成任務所需的資源數量。
主節點（master_ip+master_port）：在分布式計算環境中，主節點負責協調所有其他節點和進程的工作，為了確定主節點，我們需要知道它的IP地址和端口號。主節點還負責監控系統狀態、處理任務分配和結果匯總等任務，因此是整個系統的關鍵部分。

2.3 通信策略

DeepSpeed 還提供了 mpi、gloo 和 nccl 等通信策略，可以根據具體情況進行選擇和配置。

mpi 是一種跨節點通信庫，常用于 CPU 集群上的分布式訓練
gloo 是一種高性能的分布式訓練框架，支持 CPU 和 GPU 上的分布式訓練
nccl 是 NVIDIA 提供的 GPU 專用通信庫，被廣泛應用于 GPU 上的分布式訓練

在使用 DeepSpeed 進行分布式訓練時，可以根據具體情況選擇合適的通信庫，例如在 CPU 集群上進行分布式訓練，可以選擇 mpi 和 gloo；如果是在 GPU 上進行分布式訓練，可以選擇 nccl。

2.4 為什么需要DeepSpeed

ZeRO減少內存占用，用 3D 并行化優化并實現萬億參數模型訓練：
- pytorch官方提供的分布式訓練工具Accelerate只支持nvlink，而T4，3090這類顯卡是PIX ，檢測方式：nvidia-smi topo -m；
- DeepSpeed 實現了三種并行方法的靈活組合：ZeRO 支持的數據并行，流水線并行和張量切片模型并行。3D 并行性適應了不同工作負載的需求，以支持具有萬億參數的超大型模型，同時實現了近乎完美的顯存擴展性和吞吐量擴展效率。此外，其提高的通信效率使用戶可以在網絡帶寬有限的常規群集上以 2-7 倍的速度訓練有數十億參數的模型。
Offload 使 GPU 單卡能夠訓練 10 倍大的模型為了同時利用 CPU 和 GPU 內存來訓練大型模型，擴展了 ZeRO-2。我們的用戶在使用帶有單張英偉達 V100 GPU 的機器時，可以在不耗盡顯存的情況下運行多達 130 億個參數的模型，模型規模擴展至現有方法的10倍，并保持有競爭力的吞吐量。此功能使數十億參數的模型訓練更加大眾化，并為許多深度學習從業人員打開了一扇探索更大更好的模型的窗戶。
混合精度訓練
- 通過 DeepSpeed Sparse Attention 用6倍速度執行10倍長的序列： DeepSpeed提供了稀疏 attention kernel ——一種工具性技術，可支持長序列的模型輸入，包括文本輸入，圖像輸入和語音輸入。與經典的稠密 Transformer 相比，它支持的輸入序列長一個數量級，并在保持相當的精度下獲得最高 6 倍的執行速度提升。它還比最新的稀疏實現快 1.5–3 倍。此外，我們的稀疏 kernel 靈活支持稀疏格式，使用戶能夠通過自定義稀疏結構進行創新。
- 1 比特 Adam 減少 5 倍通信量： Adam 是一個在大規模深度學習模型訓練場景下的有效的（也許是最廣為應用的）優化器。然而，它與通信效率優化算法往往不兼容。因此，在跨設備進行分布式擴展時，通信開銷可能成為瓶頸。我們推出了一種 1 比特 Adam 新算法，以及其高效實現。該算法最多可減少 5 倍通信量，同時實現了與Adam相似的收斂率。在通信受限的場景下，我們觀察到分布式訓練速度提升了 3.5 倍，這使得該算法可以擴展到不同類型的 GPU 群集和網絡環境。

2.5 訓練介紹

對應deepspeed參數

ZeRO：這對應DeepSpeed工具中的ZeRO方式，分別是zero_optimization.stage=0/1/2/3
Offload：ZeRO-Offload 通過利用主機CPU上的計算和內存資源來執行優化器，從而減少此類模型的GPU計算和內存需求。卸載通過zero_optimization.offload_optimizer.device設置
gradient_checkpointing : 降低深度學習模型訓練過程中內存消耗的技術
混合精度：在 DeepSpeed 中，可以通過在配置文件中設置 bf16.enabled: true 來啟用 BF16 混合精度訓練，減少占用內存。混合精度訓練是指在訓練過程中同時使用FP16（半精度浮點數）和FP32（單精度浮點數）兩種精度的技術。
DeepSpeed的推理優化技術：
- Deep fusion：如下圖，紅色虛線框是以該單位為優化Kernel，對應的數字是優化的效率倍數
- Inference-customized GeMM

2.6 ZeRO思路整理

前文描述提要：存儲主要分為兩大塊：Model States和Residual States

Model State：指和模型本身息息相關的，必須存儲的內容，具體包括：
- Optimizer States：是 Optimizer 在進行梯度更新時所需要用到的數據，例如 SGD 中的 Momentum、Adam優化算法中的Momentum（動量）和Variance（方差）。
- Gradients：模型梯度 $G$ ，是在反向傳播后所產生的梯度信息，其決定了參數的更新方向。
- Parameters：模型參數 $W$ ，也就是我們在整個過程中通過數據“學習”的信息。
Residual States：指并非模型必須的，但在訓練過程中會額外產生的內容，具體包括：
- Activation：激活值。在流水線并行中曾詳細介紹過。在backward過程中使用鏈式法則計算梯度時會用到。有了它算梯度會更快，但它不是必須存儲的，因為可以通過重新做Forward來算它。
- Temporary Buffers：臨時存儲。例如把梯度發送到某塊GPU上做加總聚合時產生的存儲。
- Unusable Fragment Memory：碎片化的存儲空間。雖然總存儲空間是夠的，但是如果取不到連續的存儲空間，相關的請求也會被fail掉。對這類空間浪費可以通過內存整理來解決。

通過前文知道了什么東西會占存儲，以及它們占了多大的存儲之后，微軟開發ZeRO是為了克服數據并行性和模型并行性的限制，同時實現兩者的優點。注意到，在整個訓練中，有很多states并不會每時每刻都用到，舉例來說；

Adam優化下的optimizer states只在最終做update時才用到
數據并行中，gradients只在最后做AllReduce和updates時才用到
參數 $W$ 只在做forward和backward的那一刻才用到

所以，ZeRO想了一個簡單粗暴的辦法：如果數據算完即廢，等需要的時候，我再想辦法從個什么地方拿回來，那不就省了一筆存儲空間嗎？沿著這個思路，我們逐一來看ZeRO是如何遞進做存儲優化的。

第1種優化：針對模型狀態內存的優化（ZeRO-DP優化）
- ZeRO 將模型參數分成了三個部分：Optimizer States、Gradient 和 Model Parameter。
- ZeRO-0：禁用所有類型的分片，僅使用 DeepSpeed 作為 DDP
- ZeRO-1 Stage 1：即為 $P_{os}$
  - Optimizer State Partitioning 只對optimizer進行切片后分布式保存（每一個節點僅存部分）
  - 分割Optimizer states。優化器參數被劃分到多個memory上，每個momoey上的進程只負責更新它自己那部分參數。減少了4倍的內存，通信容量與數據并行性相同
- ZeRO-2 Stage 2：即為 $P_G$
  - Gradient Partitioning 對optimizer和grad進行切片后分布式保存（每一個節點僅存部分）
  - 分割Optimizer States與Gradients。每個memory，只保留它分配到的optimizer state所對應的梯度。這很合理，因為梯度和Optimizer是緊密聯系在一起的。只知道梯度，不知道Optimizer state，是沒有辦法優化模型參數的。
- ZeRO-3 Stage 3: 即為 $P_P$
  - Parameter Partitioning (ZeRO stage 3) 對optimizer、grad和模型參數進行切片后分布式保存（每一個節點僅存部分）
  - 分割Optimizer States、Gradients與Parameters，或者說，不同的layer. ZeRO-3會在forward和backward的時候，自動將模型參數分配到多個memory。ZeRO-Stage3將模型參數分片到不同的GPU上，通過交換節點間通信來降低顯存占用，但需要進行額外的通信操作，因此可能會導致訓練速度的下降。
第2種優化：針對殘差狀態內存的優化：ZeRO-R優化：
- 對residual states的優化（activation）
- 靈活設置保存部分activation或者每個GPU維護部分activation
- 固定大小的內存buffer
- 碎片化的存儲空間進行重新整合
第3種優化：Custom mixed precision training handling：混合精度
第4種優化：Offload優化
- ZeRO-Offload to CPU and NVMe：
  - 將模型參數分布在CPU和GPU上，通過CPU去計算一部分梯度，從而減少顯存占用，但也會帶來一定的計算開銷。
  - ZeRO-Infinity是ZeRO-3的拓展，將forward中間結果保存到內存、硬盤（NVMe）等緩存中，然后在需要時進行加載或重計算，進一步降低顯存占用
第5種優化：A range of fast CUDA-extension-based optimizers

看上去比較高大上，可能讓你很難專心去理解，但實際上，這個概念非常簡單。這只是通常的 DDP，只是沒有每個 GPU 都復制完整的模型參數、梯度和優化器狀態，而是每個 GPU 只存儲其中的一部分。在隨后的運行過程中，當需要給定層的完整層參數時，所有 GPU 同步以相互提供它們缺失的部分 —— 僅此而已。

3 混合精度模型

后文都是用如下混合精度實現方式來計算模型在訓練時需要的存儲大小，假設模型的參數 $W$ 大小是 $\Phi$ （此處可以理解為參數數量），以byte為單位，存儲如下：

必存（共計 $12\Phi$ ）：
- Parameters（FP32占4個字節，共 $\Phi$ 個）： $W_{必存}=4\Phi$
- momentum（FP32占2個字節，共 $\Phi$ 個）： $M=4\Phi$
- variance（FP32占2個字節，共 $\Phi$ 個）： $4\Phi$
中間值（共計 $4\Phi$ ）：
- Parameters（FP16）： $W_{中間}=2\Phi$
- Gradients（FP16）： $G=2\Phi$

因為采用了Adam優化，所以才會出現momentum和variance，當然你也可以選擇別的優化辦法。因此這里為了更通用些，記模型必存的數據大小為 $K\Phi$ 。因此最終內存開銷為： $2\Phi+2\Phi+K\Phi$

4 ZeRO-DP

ZeRO-DP（Zero Redundancy Optimizer-Data Parallelism）是來自于論文《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》中的一種顯存優化方法ZeRO的核心部分。通過該方法可以大幅度的優化顯存占用，從而在有限的資源下訓練更大的模型。ZeRO通過在數據并行進程中劃分模型狀態（參數，梯度和優化器狀態），而不是全復制它們，從而消除了數據并行進程中的內存冗余。它在訓練期間使用動態通信計劃，以在分布式設備之間共享必要的狀態，以保持計算粒度和數據并行性的通信量。

ZeRO驅動的數據并行性，它允許每個設備的內存使用量隨數據并行性的程度線性擴展，并產生與數據并行性相似的通信量。 ZeRO支持的數據并行性可以適合任意大小的模型，只要聚合的設備內存足夠大以共享模型狀態即可。

針對模型狀態的存儲優化（去除冗余），ZeRO使用的方法是分片（partition），即每張卡只存 $\frac{1}{N}$ 的模型狀態量，這樣系統內只維護一份模型狀態。

4.0 存儲模型與前置知識

作出如下假設：模型參數 $W$ 的個數為 $\Phi$ ，梯度個數也為 $\Phi$ ，GPU個數為 $N$

對單個GPU使用DDP來說：

Reduce-Scatter階段，通訊量為 $\frac{(N?1)\Phi}{N}$ （每一個數據塊大小 $\frac{\Phi}{N}$ ，ring上走 $N ? 1$ 次就可以在某一個GPU上完成歸并）（如果考慮發送與收到則需要兩份如此通訊量）
All-Gather階段，通訊量為 $\frac{(N?1)\Phi}{N}$ （每一個數據塊大小 $\frac{\Phi}{N}$ ，ring上走 $N ? 1$ 塊數據次就可以把某一塊GPU上完成歸并的數據塊下發至每一個GPU）（如果考慮發送與收到則需要兩份如此通訊量）
單卡單向總通訊量為 $\frac{2(N?1)\Phi}{N}$ ，隨著 $N$ 的增大，可以近似為 $2\Phi$ ，雙向則為 $4\Phi$
全卡單向總通訊量為 $2N\Phi$ ，雙向則為 $4N\Phi$

一般互聯帶寬為雙向帶寬，即同時實現相同帶寬的收與發。假設收與發的帶寬均為 $B$ 。

DP	收	發	收發通訊時間
Server（單卡）	$N\Phi$	$N\Phi$	$\frac{N\Phi}{B}$
Worker（單卡）	$\Phi$	$\Phi$	$\frac{\Phi}{B}$
集群視角（以server為瓶頸）	$N\Phi$	$N\Phi$	$\frac{N\Phi}{B}$

DDP	收	發	收發通訊時間
單卡Reduce-Scatter階段	$\Phi$	$\Phi$	$\frac{\Phi}{B}$
單卡All-Gather階段	$\Phi$	$\Phi$	$\frac{\Phi}{B}$
單卡All-Reduce	$2\Phi$	$2\Phi$	$\frac{2\Phi}{B}$
集群視角All-Reduce階段	$2N\Phi$	$2N\Phi$	$\frac{2N\Phi}{NB}=\frac{2\Phi}{B}$

DP中Server為瓶頸，搬運數據量均阻塞在Server的通訊能力上。DDP把通訊量均衡負載到了每一時刻的每個Worker上，當越來越多的GPU分布在距離較遠的機器上時，DP的通訊時間是會增加的，但是DDP可以基本不變。

4.1 $P_{os}$ ：分割優化狀態

首先，從 optimizer state 開始優化。將optimizer state分成若干份，每塊GPU上各自維護一份（不再在單塊GPU上完成所有的optimizer state的數據）。這樣就減少了相當一部分的顯存開銷。如下圖：

此時，整體數據并行的流程如下（數據量模型參考章節4.0、混合精度參考章節3）：

step1：每塊GPU上存一份完整的參數 $W_{中間}$ （使用4.0中數據個數為 $\Phi$ ，數據格式為FP16，實際存儲空間 $2\Phi$ ）， $W_{必存}$ 參考step4中的optimizer states
step2：將一個batch的數據分成 $N$ 份（X1、X2、X3）（上圖為X=3），每塊GPU各吃一份，做完一輪foward和backward后，各得一份梯度 $G$ ，（數據個數為 $\Phi$ ，數據格式為FP16）
step3：對分散在不同GPU上的梯度做一次AllReduce，得到完整的梯度 $G$ （數據個數為 $\Phi$ ，數據格式為FP16，實際存儲空間 $2\Phi$ ），產生單卡通訊量 $2\Phi$ （通訊量參考4.0 DDP表格中第4行，注意這里的通訊量沒有考慮精度，僅衡量個數）
step4：得到完整梯度 $G$ ，就可以對 $W_{必存}$ 做更新。我們知道 $W_{必存}$ 的更新由optimizer states、 $W_{必存}$ 原始值和梯度 $G$ 共同決定。這里就是 $P_{os}$ 的核心，每塊GPU上只存部分optimizer states與部分 $W_{必存}$ ，因此只能將部分對應的 $W_{中間}$ 進行更新（下圖藍色部分，數據個數為 $\frac{\Phi}{N}$ ，數據格式為FP16）
- 細節說明1： $W_{必存}$ （數據個數 $\Phi$ ，數據精度FP32，實際存儲空間 $4\Phi$ ），momentum（數據個數 $\Phi$ ，數據精度FP32，實際存儲空間 $4\Phi$ ），variance（數據個數 $\Phi$ ，數據精度FP32，實際存儲空間 $4\Phi$ ?）。因為采用了Adam優化，所以才會出現momentum和variance，當然你也可以選擇別的優化辦法。因此這里為了更通用些，記模型用于更新的數據大小為 $K\Phi$ 。因此分片后實際內存開銷為： $\frac{K}{N}\Phi$
- 細節說明2：為什么 $G$ 需要做Allreduce，同時存下完整的數據，因為對于每一片GPU都只看到了部分數據，需要所有的梯度歸并之后才能看到全局的、經過所有batchsize數據反應的真是梯度，即對于每一個GPU來說確實只需要對應數據的分片梯度，但是這里有一個地方容易搞混G1G2G3說的是不同數據回傳出來的所有參數的梯度，需要求和歸并后，才能得到對應數據的準確分片梯度
step5：每塊GPU上都有部分 $W_{中間}$ 沒有完成更新（圖中白色部分）。所以我們需要對 $W_{中間}$ 做一次All-Gather，從別的GPU上把更新好的部分 $W_{中間}$ 取回來。產生單卡通訊量 $\Phi$ （通訊量參考4.0 DDP表格中第3行，注意這里的通訊量沒有考慮精度，僅衡量個數），完成刷新后的 $W_{中間}$ （數據個數為 $\Phi$ ，數據格式為FP16，實際存儲空間 $2\Phi$ ）

step2和step3可以用下圖表示，作為下圖后藍色塊再allgather到所有GPU上變成全部藍色：

做完 $P_{os}$ 后，顯存和通訊量的情況如下：

	顯存占用	實跑顯存 $K = 12$ , $\Phi=7.5B$ , $N = 64$	單卡通訊量
DDP	$(2+2+K)\Phi$	120GB	$2\Phi$
$P_{os}$	$(2+2+\frac{K}{N})\Phi$	31.4GB	$3\Phi$

表格說明如下：

顯存占用：考慮實際的byte，考慮混合精度，考慮中間變量，具體細節參考3，分別是
- step1與step5中的： $W_{中間}$ （數據個數為 $\Phi$ ，數據格式為FP16，實際存儲空間 $2\Phi$ ）
- step3中的： $G$ （數據個數為 $\Phi$ ，數據格式為FP16，實際存儲空間 $2\Phi$ ）
- step4中的：OS與 $W_{必存}$ 分片：（數據個數為 $\frac{\Phi}{N}$ ，數據格式為FP32，由于優化手段不一樣，實際存儲空間 $\frac{K}{N}\Phi$ ）
單卡通訊量：沒有考慮byte，可以理解為僅包含個數信息，但是不影響比例關系比較

假設各變量大小如表格第二列所示，那么 $P_{os}$ 在增加1.5倍單卡通訊開銷的基礎上，將單卡存儲降低了接近 $N$ 倍。看起來是個還不錯的trade-off，那么還能做得更好嗎

4.2 $P_{os}$ + $P_{G}$ ：分割優化狀態與梯度

現在，更近一步，把梯度也拆開，每個GPU格子維護一塊梯度。

此時，整體數據并行的流程如下（數據量模型參考章節4.0、混合精度參考章節3）：

step1：每塊GPU上存一份完整的參數 $W_{中間}$ （使用4.0中數據個數為 $\Phi$ ，數據格式為FP16，實際存儲空間 $2\Phi$ ）， $W_{必存}$ 參考step4中的optimizer states
step2：將一個batch的數據分成 $N$ 份（X1、X2、X3）（上圖為X=3），每塊GPU各吃一份，做完一輪foward和backward后，各得一份梯度 $G$ ，下圖中綠色+白色
step3：對分散在不同GPU上的梯度做一次Reduce-Scatter，保證每個GPU上所維持的那塊梯度是聚合梯度。例如對GPU1，它負責維護G1，因此其他的GPU只需要把G1對應位置的梯度發給GPU1做加總就可。匯總完畢后，保留下圖中的綠色，白色塊對本GPU無用，可以從顯存中移除。所以不同GPU分別維護分片G1/G2/G3（數據個數為 $\frac{\Phi}{N}$ ，數據格式為FP16，實際存儲空間 $\frac{2\Phi}{N}$ ），單卡通訊量 $\Phi$ 。（通訊量參考4.0 DDP表格中第2行，注意這里的通訊量沒有考慮精度，僅衡量個數）
step4：每塊GPU用自己對應的 $O$ 和 $G$ 去更新相應的 $W_{必存}$ （參考章節4.1，數據個數為 $\frac{\Phi}{N}$ ，數據格式為FP32，由于優化手段不一樣，實際存儲空間 $\frac{K}{N}\Phi$ ）
step5（同4.1）：每塊GPU上都有部分 $W_{中間}$ 沒有完成更新。所以我們需要對 $W_{中間}$ 做一次All-Gather，從別的GPU上把更新好的部分 $W_{中間}$ 取回來。產生單卡通訊量 $\Phi$ （通訊量參考4.0 DDP表格中第2行，注意這里的通訊量沒有考慮精度，僅衡量個數），完成刷新后的 $W_{中間}$ （數據個數為 $\Phi$ ，數據格式為FP16，實際存儲空間 $2\Phi$ ）

Step2和Step3見下圖：

做完 $P_{os}$ + $P_{G}$ 后，顯存和通訊量的情況如下：

	顯存占用	實跑顯存 $K = 12$ , $\Phi=7.5B$ , $N = 64$	單卡通訊量
DDP	$(2+2+K)\Phi$	120GB	$2\Phi$
$P_{os}$	$(2+2+\frac{K}{N})\Phi$	31.4GB	$3\Phi$
$P_{os}$ + $P_G$	$(2+\frac{2+K}{N})\Phi$	16.6GB	$2\Phi$

和DDP相比，上述例子中存儲降了8倍，單卡通訊量持平，（通信量的優化主要是因為不需要全部G進行allreduce通信）好像更牛皮了呢！那么，還可以優化嗎？

4.3 $P_{os}$ + $P_G$ + $P_P$ ：分割優化狀態、梯度與參數

看到這里，也許你有點感覺了，ZeRO的思想就是：萬物皆可切，萬物皆可拋。所以現在，我們把參數也切開。每塊GPU置維持對應的optimizer states，gradients和parameters（即 $W$ ）。

此時，整體數據并行的流程如下（數據量模型參考章節4.0、混合精度參考章節3）：

step1：每塊GPU上只保存參數 $W_{中間}$ 的部分切片（數據個數為 $\frac{\Phi}{N}$ ，數據格式為FP16，實際存儲空間 $\frac{2\Phi}{N}$ ）
step2：將一個batch的數據分成 $N$ 份，每塊GPU各吃一份，做forward時，對 $W_{中間}$ 做一次All-Gather，取回分布在別的GPU上的 $W_{中間}$ ，得到一份完整的 $W_{中間}$ ，單卡通訊量 $\Phi$ 。forward做完，立刻把不是自己維護的 $W_{中間}$ 拋棄。
step3：做backward時，對 $W_{中間}$ 做一次All-Gather，取回完整的 $W_{中間}$ ，單卡通訊量 $\Phi$ 。backward做完，立刻把不是自己維護的 $W_{中間}$ 拋棄。各得一份梯度 $G$
step4：對分散在不同GPU上的梯度做一次Reduce-Scatter，保證每個GPU上所維持的那塊梯度是聚合梯度，其余不相關梯度扔掉。（分片 $G$ ：數據個數為 $\frac{\Phi}{N}$ ，數據格式為FP16，實際存儲空間 $\frac{2\Phi}{N}$ ），單卡通訊量 $\Phi$
step5：每塊GPU用自己對應的 $O$ 和 $G$ 去更新相應的 $W_{必存}$ （參考章節4.1，數據個數為 $\frac{\Phi}{N}$ ，數據格式為FP32，由于優化手段不一樣，實際存儲空間 $\frac{K}{N}\Phi$ ）
step6（同4.1）：每塊GPU上都有部分 $W_{中間}$ 沒有完成更新。由于只維護部分 $W_{中間}$ ，因此無需再對 $W_{中間}$ 做任何AllReduce操作，此時就回到了step1的狀態（數據個數為 $\frac{\Phi}{N}$ ，數據格式為FP16，實際存儲空間 $\frac{2\Phi}{N}$ ）。

做完 $P_{os}$ + $P_{G}$ + $P_P$ 后，顯存和通訊量的情況如下：

	顯存占用	實跑顯存 $K = 12$ , $\Phi=7.5B$ , $N = 64$	單卡通訊量
樸素DP	$(2+2+K)\Phi$	120GB	$2\Phi$
$P_{os}$	$(2+2+\frac{K}{N})\Phi$	31.4GB	$3\Phi$
$P_{os}$ + $P_G$	$(2+\frac{2+K}{N})\Phi$	16.6GB	$2\Phi$
$P_{os}$ + $P_G$ + $P_p$	$(\frac{2+2+K}{N})\Phi$	1.9GB	$3\Phi$

到這一步，我們用1.5倍的通訊開銷，換回近120倍的顯存。只要梯度計算和異步更新做的好，通訊時間大部分可以被計算時間隱藏，因此這樣的額外通訊開銷，也是劃算的。

到這里，我們可以放出原始論文中的說明圖了。

4.4 ZeRO VS 模型并行

模型并行，是指在forward和backward的過程中，只需要用自己維護的那塊 $W$ 來計算就行。即同樣的輸入 $X$ ，每塊GPU上各算模型的一部分，最后通過某些方式聚合結果。

大家可能會想，既然ZeRO都把參數 $W$ 給切了，那它應該是個模型并行呀？為什么要歸到數據并行里呢？

其實ZeRO是模型并行的形式，數據并行的實質。

對ZeRO來說，它做forward和backward的時候，是需要把各GPU上維護的 $W$ 聚合起來的，即本質上還是用完整的 $W$ 進行計算。它是不同的輸入 $X$ ，完整的參數 $W$ ，最終再做聚合。

5 ZeRO-R

說完了以上對model states的顯存優化，現在來看對residual states的優化。

3.1 $P a$ ：Partitioned Activation Checkpointing

前面說過，對activation的存儲是靈活的。不像optimizer states，gradients和parameters對模型更新是必須的，activation只是起到加速梯度計算的作用。因此，在哪幾層保存activation，保存哪些activation都是可以靈活設置的。同樣，我們也可以仿照以上切割方式，每塊GPU上只維護部分的activation，需要時再從別的地方聚合過來就行。需要注意的是，activation對顯存的占用一般會遠高于模型本身，通訊量也是巨大的，所以這塊要靈活、有效地實驗設計。

3.2 $C_B$ ：Constant Size Buffer

固定大小的內存buffer，它的目的在于：

提升帶寬利用率。當GPU數量上升，GPU間的通訊次數也上升，每次的通訊量可能下降（但總通訊量不會變）。數據切片小了，就不能很好利用帶寬了。所以這個buffer起到了積攢數據的作用：等數據積攢到一定大小，再進行通訊。
使得存儲大小可控。在每次通訊前，積攢的存儲大小是常量，是已知可控的。更方便使用者對訓練中的存儲消耗和通訊時間進行預估。

3.3 $M_D$ ：Memory Defragmentation

在前文提過，設置機制，對碎片化的存儲空間進行重新整合，整出連續的存儲空間。防止出現總存儲足夠，但連續存儲不夠而引起的存儲請求fail

6 ZeRO-Offload與ZeRO-Infinity

最后，簡單介紹一下ZeRO-Offload。它的核心思想是：顯存不夠，內存來湊。

如果我把要存儲的大頭卸載（offload）到CPU上，而把計算部分放到GPU上，這樣比起跨機，是不是能既降顯存，也能減少一些通訊壓力呢？

Offload優化
- ZeRO-Offload和ZeRO-Stage3是DeepSpeed中的不同的Zero-Redundancy Optimization技術，用于加速分布式訓練，主要區別在資源占用和通信開銷方面。
- ZeRO-Offload將模型參數分布在CPU和GPU上，通過CPU去計算一部分梯度，從而減少顯存占用，但也會帶來一定的計算開銷。
ZeRO-Infinity是ZeRO-3的拓展。允許通過使用 NVMe 固態硬盤擴展 GPU 和 CPU 內存來訓練大型模型。ZeRO-Infinity 需要啟用 ZeRO-3。

6.1 ZeRO-Offload

ZeRO-Offload的核心思路就是讓CPU和內存也參與到訓練中去，回顧一下前文用到的訓練流程的圖，ZeRO-Offload就是把這個流程用上圖的方式把fp32參數的更新和float2half操作拆分到了CPU和內存上計算，而前向和后向傳播依然由GPU負責

具體做法是：

forward和backward計算量高，因此和它們相關的部分，例如參數W（fp16），activation，就全放入GPU。
update的部分計算量低，因此和它相關的部分，全部放入CPU中。例如W(fp32)，optimizer states（fp32）和gradients(fp16)等。

具體切分如下圖：

6.2 ZeRO-infinity

也是同理，它們在解決的事情都是：找個除GPU之外的地方，存數據。

7 總結

在DP中，每個GPU上都拷貝一份完整的模型，每個GPU上處理batch的一部分數據，所有GPU算出來的梯度進行累加后，再傳回各GPU用于更新參數
DP多采用參數服務器這一編程框架，一般由若個計算Worker和1個梯度聚合Server組成。Server與每個Worker通訊，Worker間并不通訊。因此Server承擔了系統所有的通訊壓力。基于此DP常用于單機多卡場景。
異步梯度更新是提升計算通訊比的一種方法，延遲更新的步數大小決定了模型的收斂速度。
Ring-AllReduce通過定義網絡環拓撲的方式，將通訊壓力均衡地分到每個GPU上，使得跨機器的數據并行（DDP）得以高效實現。
DP和DDP的總通訊量相同，但因負載不均的原因，DP需要耗費更多的時間搬運數據
ZeRO：通過增加部分通信的代價，減少本GPU中不必要存儲或者計算的數據，進一步縮小GPU的內存開銷