天天看點
傳回

[源碼分析] Facebook如何訓練超大模型 --- (3)

FSDP是Facebook 深度借鑒微軟ZeRO之後提出的PyTorch DDP更新版本，可以認為是對标微軟 ZeRO，目标是訓練超大規模模型。前文我們介紹了 FSDP 如何使用，本文從源碼角度來介紹 FSDP 如何實作 offload。

[源碼分析] Facebook如何訓練超大模型 --- (3)

目錄

    • 0x00 摘要
    • 0x01 ZeRO-Offload
      • 1.1 設計原則
      • 1.2 ZeRO
    • 0x02 解除安裝政策
      • 2.1 資料流圖
      • 2.2 限制CPU計算
      • 2.3 最小化計算量
      • 2.4 最大化記憶體節約
      • 2.5 唯一最優化政策
      • 2.6 ZeRO-Offload Schedule
        • 2.6.1 單機計劃
        • 2.6.2 多節點計劃
    • 0x03 FairScale Offload 使用
      • 3.1 思路
      • 3.2 使用
      • 3.3 配置
    • 0x04 源碼
      • 4.1 建構
        • 4.1.1 初始化
        • 4.1.2 切片
      • 4.2 ModelShard
        • 4.2.1 定義
        • 4.2.2 功能函數
      • 4.3 前向傳播
        • 4.3.1 前向傳播
        • 4.3.2 Hook
    • 0xFF

我們在前文介紹過，微軟 ZeRO 可以對一個萬億參數模型可以使用 8 路模型并行、64 路管道并行和 8 路資料并行在 4,096 個 NVIDIA A100 GPU 上進行擴充。

而FSDP（Fully Sharded Data Parallel）是Facebook 深度借鑒微軟ZeRO之後提出的PyTorch DDP更新版本，可以認為是對标微軟 ZeRO，其本質是 parameter sharding。Parameter sharding 就是把模型參數等切分到各個GPU之上。我們會以 Google，微軟和 Facebook 的論文，部落格以及代碼來進行學習分析。

前文我們介紹了 FSDP 如何實作參數分區，FSDP 也會和Offload一起使用，這兩項加起來就是ZeRO-offload的實作。本文基于原始論文 https://arxiv.org/pdf/2101.06840.pdf，官博https://www.deepspeed.ai/tutorials/zero-offload/ 和源碼來一起分析學習。

本系列其他文章如下：

[源碼解析] PyTorch 分布式之 ZeroRedundancyOptimizer

[論文翻譯] 分布式訓練 Parameter sharding 之 ZeRO

[論文翻譯] 分布式訓練 Parameter Sharding 之 Google Weight Sharding

[源碼分析] Facebook如何訓練超大模型---(1)

[源碼分析] Facebook如何訓練超大模型 --- (2)

基于 Zero Redundancy Optimizer 基礎之上，加利福尼亞大學默塞德分校和微軟的一組研究人員開發了 ZeRO-Offload。ZeRO-Offload 通過同時利用GPU和主控端 CPU 的計算和存儲資源，提升了在較少 GPU 資源下可以高效訓練的模型規模。

ZeRO-Offload 核心技術是在 ZeRO-2基礎之上将優化器狀态和梯度卸至 CPU 記憶體。優化器狀态在整個訓練過程中将消耗大部分 GPU 顯存，反向傳播過程中計算出來的梯度也占據了相當的顯存，把他們移到CPU，這樣盡管存在拷貝至 CPU 的開銷，但是節省的 GPU 顯存可用于訓練更大的模型，GPU 計算效率仍然可以提高。

ZeRO-offload 屬于CPU解除安裝技術，就是當GPU記憶體已滿時，可以将暫時未使用的資料解除安裝到CPU，并在以後需要時将其讀回（Rhu等人，2016）。ZeRO-offload 基于三個原則來設計：效率、可伸縮性和可用性。其背後的關鍵技術是：在 ZeRO-2 基礎上将優化器計算，優化器狀态和梯度解除安裝到 CPU 記憶體。這種方法讓 ZeRO-Offload 能最大程度降低拷貝至 CPU 導緻的計算效率損失，同時還實作了與原始ZeRO-2相同的效率，有時甚至更好。研究人員已經可以确定 CPU 和 GPU 之間資料分區和最佳計算政策。該方法涉及到的流程包括如何将梯度、優化器狀态和優化器計算分散到 GPU，以及如何在 GPU 上進行向前和向後計算。

下圖展示了 Zero-OffLoad 的架構：

[源碼分析] Facebook如何訓練超大模型 --- (3)

ZeRO-Offload 概述，圖來自 https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/

ZeRO-Offload與ZeRO一起工作，可将DL訓練擴充到多個GPU。ZeRO有三個階段，分别對應于三種不同的劃分：模型狀态、優化器狀态、梯度和參數的劃分，分别為ZeRO-1、ZeRO-2和ZeRO-3。

  • ZeRO-1隻對優化器狀态進行分區。
  • ZeRO-2除了對優化器狀态進行分區外，還對梯度進行分區，
  • ZeRO-3對所有模型狀态進行分區。

ZeRO-Offload 與ZeRO-2協同工作，是以我們将對其進行進一步讨論。

在ZeRO-2中，每個GPU都存儲了所有參數的副本，但在每個訓練步驟結束時的參數更新中，隻更新其中自己GPU負責的部分。由于每個GPU隻更新一部分參數，它們隻存儲進行更新所需的優化器狀态和梯度。在更新之後，每個GPU使用一個all-gather通信将其更新參數的部分發送給所有其他GPU。ZeRO-2的計算和通信具體描述如下。

  • 在前向傳播過程中，每個GPU計算不同mini-batch的損失。
  • 在後向傳播過程中，當計算出每個梯度之後，在擁有該梯度或部分梯度的GPU/GPU上會使用reduce算子對該梯度進行平均化。
  • 在後向傳播完成之後，每個GPU使用平均梯度來更新其部分參數和優化器狀态。
  • 更新之後，會進行一次all-gather以接收在其他GPU上計算的其餘參數更新。

下面就讓我們研讀一下論文内容。

ZeRO-Offload旨在通過在訓練期間将一些模型狀态從GPU解除安裝到CPU記憶體，進而在單個或多個GPU上實作高效的大型模型訓練。

如前所述，模型狀态：參數、梯度和優化器狀态，是大型模型訓練中記憶體瓶頸的主要來源。通過将這些模型狀态的一部分解除安裝到CPU，ZeRO-Offload可以訓練更大的模型。然而，确定最佳的解除安裝政策并非易事。有許多方法可以将模型狀态解除安裝到CPU記憶體中，每一種方法在CPU計算和GPU-CPU通信方面有不同的權衡。

為了确定最佳的解除安裝政策，ZeRO-Offload将DL訓練模拟成資料流圖，并使用第一原理來在CPU和GPU裝置之間對這個圖進行有效地劃分。ZeRO-Offload在三個關鍵方面對圖進行了優化：

  • i）隻在CPU上進行少量計算，以防止CPU成為性能瓶頸。和GPU相比，CPU的計算量是數量級減少。
  • ii）確定CPU和GPU記憶體之間的通信量最小；
  • iii）在實作最小通信量的同時，它可以最大限度地節省記憶體。

事實上，ZeRO-Offload可以在訓練過程中實作與非解除安裝訓練相媲美的高效率，而且它是獨特的最佳（unique optimal），這意味着沒有其他解決方案可以在不增加通信量或增加CPU計算的情況下提供更好的記憶體節省。

接下來将讨論獨特最優解除安裝政策的推導，該政策是專門為混合精度訓練與Adam優化器設計的。

[源碼分析] Facebook如何訓練超大模型 --- (3)

DL訓練的工作量可以表示為資料和計算的權重有向圖，如圖所示，其中圓形節點代表模型狀态（參數16，梯度16，參數32，動量32，方差32），矩形節點代表計算（向前、向後、參數更新）。圖中的邊代表節點之間的資料流，邊的權重是在任何給定的訓練疊代期間流經它的總資料量（以位元組為機關）。對于一個有M個參數的模型，在源節點産生fp16模型狀态的情況下，該圖中的邊的權重為2M，或者在源節點産生fp32模型狀态的情況下為4M。

GPU和CPU之間的解除安裝政策可以用這個圖的雙向分區來表示，比如分區中的計算節點将在擁有該分區的裝置上執行，而該分區中的資料節點将存儲在擁有該分區的裝置上。GPU和CPU之間必須通信的總資料量由兩個分區上運作的邊的權重給出。有許多方法可以對該圖進行分區。比如可以使用第一原理簡化資料流圖，以減少基于三個不同效率名額的可能選擇的數量：i）CPU計算量開銷，ii）通信開銷，以及iii）記憶體節省。

CPU計算吞吐量比GPU計算吞吐量慢多個數量級。是以，将大型計算圖解除安裝到CPU将嚴重限制訓練效率。是以，我們必須避免将計算密集型元件解除安裝到CPU上。

DL訓練每個疊代的計算複雜度通常由O(MB)給出，其中M是模型大小，B是有效batch size。為了避免CPU計算成為瓶頸，隻有那些計算複雜度低于O(MB)的計算才應該解除安裝到CPU上。這意味着計算複雜度為O(MB)的前向傳播和後向傳播必須在GPU上完成，而複雜度為O(MB)的剩餘計算（如範數計算、權重更新等）可能會解除安裝到CPU上。

基于這個簡單的觀察，我們将資料流圖中的前向和後向節點融合為一個超級節點（FWD-BWD），并将其配置設定到GPU上。

我們接下來分析最小化計算量（Minimizing Communication Volume）。

CPU記憶體帶寬至少比CPU和GPU之間的PCI-E帶寬快一個數量級，而GPU記憶體比CPU記憶體快一個數量級。是以，我們必須最小化CPU和GPU記憶體之間的通信量，以防止PCI-E帶寬成為訓練性能瓶頸。為此，我們必須首先确定模型狀态解除安裝政策的理論最小通信量。

模型狀态解除安裝政策的最小通信量為4M（M是模型大小）。請注意，在将前向和後向融合為單個超級節點後，資料流圖中的每個節點都是一個循環的一部分。是以，此圖的任何分區都需要在至少兩條邊上做切割。每條邊的權重至少為2M，導緻總通信量至少為4M。

如果我們選擇将通信量限制在這個最小值，我們可以大大簡化資料流圖，并将分區政策的數量減少到較少數量。

建立fp32超級節點：請注意，任何不将fp32模型放在同一位置的分區政策都表明其生産者和消費者節點無法實作4M的最小通信量。這樣的分區必須在至少在如下兩條邊上切分：一條權重為4M的邊和另一條至少2M的邊，進而産生至少6M的通信量。是以，為了實作最小通信量，所有解除安裝政策必須将fp32模型狀态與其生産者和消費者算子放在一起，即fp32模型狀态（動量32、方差32和p32）必須與Param Update和 float2half 計算放在同一位置。

此限制允許我們将資料流圖中的所有上述fp32資料和計算節點視為一個超級節點，我們稱之為Update super。我們在圖2中展示了這個簡化的資料流圖，它僅由四個節點組成：FWD-BWD超級節點、p16資料節點、g16資料節點和更新超級節點。

p16配置設定：為了實作最小通信量，p16必須與FWD-BWD Super位于同一位置，因為這兩個節點之間的邊緣權重為4M。如果這兩個節點分開，通信量将會增加到6M（4M+2M）。由于我們已經将節點FWD-BWD Super配置設定給GPU以限制CPU上的計算，p16也必須配置設定給GPU。

我們接下來看看如何最大化記憶體節約（Maximizing Memory Savings）。

在簡化資料流圖以最小化通信量之後，隻剩下g16和Update Super需要被配置設定。請注意，在這一點上，所有的分區結果都會導緻最小的通信量，是以我們可以進一步調整選擇，以最大限度地節省GPU的記憶體。表1顯示了所有有效的分區政策所帶來的記憶體節省，這些政策使通信量最小。通過将g16和Update Super解除安裝到CPU，可以實作8倍的最大記憶體節省。

[源碼分析] Facebook如何訓練超大模型 --- (3)

ZeRO-Offload在CPU記憶體中配置設定所有的fp32模型狀态以及fp16梯度，它也在CPU中計算參數更新。fp16的參數保留在GPU上，前向和後向的計算也在GPU上完成。

我們通過簡化我們的資料流圖來得出這個解除安裝政策，并排除了所有其他的分區政策，因為其他政策或者不能限制CPU的計算，或者無法最小化通信量，或無法最大限度地節省記憶體。是以，ZeRO-Offload不僅在上述名額上是最優的，而且是唯一的；不可能有其他政策能比ZeRO-Offload節省更多的記憶體，而不增加CPU的計算複雜性或産生額外的GPU-CPU通信量。

在這一節中，我們将讨論基于我們的解除安裝政策，如何在單GPU系統上實作ZeRO-Offload的具體計算和通信schedule。然後，我們将展示如何通過将我們的解除安裝政策與ZeRO資料并行和模型并行結合起來，把這個schedule擴充到多GPU系統上有效工作。

ZeRO-2 在每個 GPU 上儲存一部分優化器狀态量和梯度，ZeRO-Offload 繼承了 ZeRO-2 的劃分優化器狀态量和梯度的方法。和 ZeRO-2 不同之處在于，ZeRO-Offload 把優化器狀态量和梯度移到了本機記憶體上。即，ZeRO-Offload 對資料進行分區，使：

  • fp16參數存儲在GPU中。
  • fp32參數儲存在CPU記憶體中。
  • fp16梯度儲存在CPU記憶體中。
  • 所有優化器狀态（如fp32動量、方差）在整體訓練過程中都儲存在CPU記憶體中。

在計算時：

  • 我們首先通過前向傳播計算損失。由于fp16參數已在GPU上，是以這部分計算不需要CPU通信。
  • 在損失的反向傳播過程中，在反向排程的不同點計算不同參數的梯度。
    • 可以在計算每個參數後立即将這些梯度單獨或分組傳輸到CPU記憶體。是以，在将梯度傳輸到CPU記憶體之前，隻需少量記憶體即可臨時保留GPU記憶體上的梯度。
    • 每個梯度傳輸可以與反向圖的剩餘部分上的反向傳播重疊，進而允許ZeRO-Offload隐藏通信成本的重要部分。
  • 反向傳播後，ZeRO-Offload 直接在CPU上更新fp32參數和剩餘優化器狀态（如動量和方差），并将更新後的fp32參數從CPU記憶體複制為GPU記憶體上的fp16參數。下圖以圖解的方式顯示了ZeRO-Offload的每個步驟中的計算和通信，
    • 當梯度到了 CPU 之後，劃分後的優化狀态變量就會并行在 CPU 上進行更新（圖中的 p update）。
    • 當更新完成之後，劃分後的參數就被移回GPU，接下來會用 all gather 操作進行更新（(圖中的 g swap）。
    • 通過使用不同 CUDA stream 來讓通信（如 g offload 和 g swap）和計算（如反向傳播和 p update) 重疊起來，通信隐藏在計算之中，這樣可以提高訓練效率。
[源碼分析] Facebook如何訓練超大模型 --- (3)

下圖以僞代碼的形式顯示了具體的計劃。

[源碼分析] Facebook如何訓練超大模型 --- (3)

ZeRO-Offload 可以有效地擴充到數百個GPU。ZeRO-Offload 保留ZeRO Stage-2（優化器狀态和梯度分區）的模型狀态分區政策，同時将分區的梯度、優化器狀态和相應的參數更新解除安裝到CPU。

在解除安裝之前進行分區的主要好處是，對于具有1個以上GPU的系統，每個資料并行程序隻負責更新參數的子集。從所有資料并行GPU到CPU的聚合通信量保持不變，而且并行使用CPU資源共同計算單個權重更新。是以，總的CPU更新時間随着資料并行度的增加而減少，

因為CPU計算資源随着計算節點數量的增加而線性增加。這允許ZeRO-Offload 實作非常好的可伸縮性，因為CPU優化器步驟的減少抵消了跨GPU的通信開銷。ZeRO-Offload 在不同的GPU之間劃分梯度和優化器狀态，每個GPU将其擁有的分區解除安裝到CPU記憶體中，并在整個教育訓練過程中保持該分區。

在反向傳播過程中，ZeRO-Offload 使用GPU上的reduce scatter計算并且平均梯度，每個資料并行程序（GPU）僅将屬于其分區的平均梯度解除安裝到CPU記憶體上（下圖中的 g offload）并且把自己不負責的部分丢棄掉。

一旦梯度在CPU上可用，優化器狀态分區将由CPU上的每個資料并行程序并行更新。更新後，參數分區移回GPU，然後在GPU上執行類似于ZeRO-2的all gather操作來收集所有參數。下圖顯示了ZeRO-Offload 的data placement模型參數、梯度和優化器狀态。

[源碼分析] Facebook如何訓練超大模型 --- (3)

ZeRO-Offload資料并行排程的詳細資訊如代碼圖所示。上述all gather操作在代碼圖中顯示為一系列廣播操作。

以下思路結合了FairScale的文檔和自己的思考。

一般來說，大型模型往往會導緻OOM錯誤，而FairScale

OffloadModel

API使使用者能夠在有限的GPU資源上訓練大型模型，進而實作了大規模分布式訓練。

OffloadModel

支援混合精度訓練、可以使用激活檢查點減少記憶體占用，以及使用微批來處理降低通信量。

FairScale Offload 受到

Layer-to-Layer <https://arxiv.org/abs/2002.05645>

和

Zero-Offload <https://arxiv.org/abs/2101.06840>

的深度啟發，OffloadModel使用CPU存儲整個模型、優化器狀态和梯度。OffloadModel然後将一層（或多個層）加載到GPU上，以便在向前和向後傳播過程中進行訓練。層與層邊界的中間激活也存儲在CPU上，并根據向後傳播的需要複制到GPU。完成後向傳播後，模型的所有參數将使用位于CPU上的梯度進行更新，具體可以參見下面的示例圖。

[源碼分析] Facebook如何訓練超大模型 --- (3)

Offload 的執行有一個假定條件：模型假定為nn.Sequential模型，并根據參數數量（幾乎）平均分片到nn.Modules 清單之中。每個 nn.Module 現在包含整個模型的一部分，我們稱之為模型分片（model shards）。

在這個假定條件基礎之上，Offload 具體采用了以下方法來進行具體實作：

  • 在每次疊代中，從CPU複制每個模型分片到GPU，然後使用小批量（minibatch）資料計算前向傳播，并把模型分片從GPU複制回CPU。在後向傳播過程中，重複相同的過程。本文對應了此項具體實作。
  • 優化器保留在CPU上，在運作optimizer.step之前，梯度和參數都會移動到CPU上。這確定了CPU可以更新參數并保持優化器狀态。優化器部分文章對應了此項具體實作。具體可以參見 self.move_grads_to_cpu 選項。
  • 如果啟用了激活檢查點，我們将使用torch.autograd.Function來禁用FW過程中的計算圖構造，并在給定分片的FW過程完成後把中間激活從GPU複制到CPU。BW過程中執行相反複制操作。後續 Activation 文章會講述此項實作。
  • 可以使用微批次（Micro-batches）實作更大的吞吐量，并抵消從CPU<->GPU移動模型參數和激活的成本。微批次技術允許您指定大的小批次，這些小批次被分解為微批次（micro-batches），并在每次疊代時饋送到模型分片。簡言之，這是一種允許在給定時間在模型分片之上進行更多計算的方法，以抵消從CPU<->GPU複制的成本。

具體使用樣例如下，首先會進行正常配置，并且定義了一個Sequential模型。

from torch.utils.data.dataloader import DataLoader
from torchvision.datasets import FakeData
from torchvision.transforms import ToTensor

# 引入Offload
from fairscale.experimental.nn.offload import OffloadModel 

# 定義訓練配置
num_inputs = 8
num_outputs = 8
num_hidden =  4
num_layers =  2
batch_size =  8

# 資料加載
transform = ToTensor()
dataloader = DataLoader(
    FakeData(
        image_size=(1, num_inputs, num_inputs),
        num_classes=num_outputs,
        transform=transform,
    ),
    batch_size=batch_size,
)

# 定義了Sequential模型，注意前面提到的：模型假定為nn.Sequential模型，并根據參數數量（幾乎）平均分片到nn.Modules 清單之中。
model = torch.nn.Sequential(
    torch.nn.Linear(num_inputs * num_inputs, num_hidden),
    *([torch.nn.Linear(num_hidden, num_hidden) for _ in range(num_layers)]),
    torch.nn.Linear(num_hidden, num_outputs),
)

然後，要使用OffloadModel API，我們應該使用 OffloadModel 來包裝模型，包裝時，使用者可以指定：

  • 用于計算向前和向後傳播的裝置。
  • 模型将存儲在其上的offload 裝置。
  • 模型應分片的片數。
  • 預設情況下，激活檢查點處于關閉狀态，微批次數為1。 
offload_model = OffloadModel( # 使用 OffloadModel 來包裝模型
    model=model, # 原生模型
    device=torch.device("cuda"), # 用于計算向前和向後傳播的裝置
    offload_device=torch.device("cpu"), # 模型将存儲在其上的offload 裝置
    num_slices=3, # 模型應分片的片數
    checkpoint_activation=True,
    num_microbatches=1,
)

torch.cuda.set_device(0)
device = torch.device("cuda")

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(offload_model.parameters(), lr=0.001) # 使用OffloadModel

# To train 1 epoch.
offload_model.train() # 使用 OffloadModel
for batch_inputs, batch_outputs in dataloader:
    batch_inputs, batch_outputs = batch_inputs.to("cuda"), batch_outputs.to("cuda")
    start = time.time_ns()
    optimizer.zero_grad()
    inputs = batch_inputs.reshape(-1, num_inputs * num_inputs)
    with torch.cuda.amp.autocast():
        output = model(inputs) # 前向傳播
        loss = criterion(output, target=batch_outputs)
        loss.backward() # 反向傳播
    optimizer.step()

Offload 有如下配置，在使用時候可以注意。

move_params_to_cpu (bool, Optional):
    if ``True``, offload FP32 params to CPU. This is only relevant when
    *``mixed_precision``* is ``True``.
    
cpu_offload (bool, Optional):
    if ``True``, offload FP32 params to CPU. This is only relevant when
    *``mixed_precision``* is ``True``. Note: This arg will be deprecated in favor of
    *``move_params_to_cpu``* in an upcoming release.  
    
move_grads_to_cpu (bool, Optional):
    move gradient shard to CPU after reduction. This is useful when
    combined with CPU-based optimizers. It defaults to the value of
    *``cpu_offload``*.

我們接着看看如何建構一個 OffloadModel。

因為Python語言的特點，在初始化函數中可以看到 OffloadModel 的内部成員變量。傳遞的參數基本都直接配置到内部成員變量之中，除了model需要特殊處理。關于模型處理，回憶一下前面提到的：模型假定為nn.Sequential模型，并根據參數數量（幾乎）平均分片到nn.Modules 清單之中。

具體操作是：看看模型是否是list類型，如果是，說明已經分片好了，則直接把每一層用ModelShard封裝到 model_slices，否則先調用_split 進行切片再封裝到 model_slices。

class OffloadModel(nn.Module):
    def __init__(
        self,
        model: Any,
        device: torch.device,
        offload_device: torch.device = torch.device("cpu"),
        num_slices: int = 3,
        checkpoint_activation: bool = False,
        num_microbatches: int = 1,
    ):
        super().__init__()

        self.device = device # 計算裝置
        self.offload_device = offload_device # 設定解除安裝裝置，一般來說就是cpu
        # List of model shards that will be placed on/off the device.
        self.model_slices: List[nn.Module] = [] # 存儲原生模型的分片

        if type(model) == list: # list代表已經分片好了
            # This is already sharded using the auto shard functinality.
            for i, m in enumerate(model):
                self.model_slices.append( # 直接把每一層用ModelShard封裝
                    ModelShard(cpu_model_shard=m, device=device, offload_device=offload_device, index=i,)
                )
        else:
            # Slice the model into roughly equivalent sequential shards.
            splits = _split(model, num_slices) # 否則先split

            for i, split in enumerate(splits): # 周遊split分區結果
                # Add one model handling this slice
                self.model_slices.append( # 然後把每一個分區用ModelShard封裝
                    ModelShard(
                        cpu_model_shard=nn.Sequential(*split), device=device, offload_device=offload_device, index=i,
                    )
                )

        # Expose a unified view of the slices
        self._model = torch.nn.Sequential(*self.model_slices) # 最後生成一個nn.Sequential

        # intermediate activations at the slice boundaries.
        self._activations: List[Tuple] = []

        # Currently we only support microbatches with activation checkpointing.
        if not checkpoint_activation and num_microbatches > 1:
            raise RuntimeError("We currently only support microbatches with activation checkpointing.")

        # Bool indicating if we want to checkpoint activation on the host.
        self._checkpoint_activation = checkpoint_activation

        # Number of microbatches to run per batch on the device
        self._num_microbatches = num_microbatches

初始化代碼之中使用了_split 方法來切分，這就對應了前面思路之中提到的：模型假定為nn.Sequential模型，并根據參數數量（幾乎）平均分片到nn.Modules 清單之中。每個 nn.Module 現在包含整個模型的一部分，我們稱之為模型分片（model shards）。

我們具體看看代碼，就能知道是如何大緻進行均勻分區的。

def _split(modules: nn.Sequential, number_splits: int) -> List[List[nn.Module]]:
    # 設定最小切分數目
    number_splits = min(len(modules), number_splits) 
    # 生成切分之後的容器
    splits: List[List[nn.Module]] = [[] for _ in range(number_splits)]

    # Count the number of parameters per exposed layer, use that as a proxy for memory footprint
    # 計算modules的每層參數的元素數目之和
    # p.numel()作用是擷取tensor中一共包含多少個元素，比如 torch.randn(3,3) 是9個元素
    total_number_params = sum([sum(p.numel() for p in m.parameters()) for m in modules])
    # 每個分區應該得到的元素個數
    number_parameters_per_shard = total_number_params // number_splits

    current_shard = 0

    for m in modules: # 周遊module的層
        for p in m.parameters(): # 周遊每層的參數
            p.data = p.data.pin_memory() # 把參數放到鎖頁記憶體，這樣其轉到GPU會更快。
        # Number of parameters in the current shard
        # 看看目前分區的元素數目
        current_shard_params = sum(p.numel() for sm in splits[current_shard] for p in sm.parameters())

        # This shard is big enough, point to the next one
        # 如果目前分區夠大了，就跳到下一個分區
        if (
            current_shard_params > 0
            and current_shard_params + sum(p.numel() for p in m.parameters()) > number_parameters_per_shard
            and current_shard < number_splits - 1
        ):
            current_shard += 1

        # 把m這層放到splits目前分區
        splits[current_shard].append(m) 

    # 列印出來每個分區大小    
    for i, split in enumerate(splits):
        current_shard_params = sum(p.numel() for sm in split for p in sm.parameters())
        logging.info(f"Shard {i} holds {current_shard_params/1e6:.2f}M parameters")

    return splits

Sequential模型的每個module被封裝為ModelShard，是以我們繼續看看ModelShard。

ModelShard的作用是封裝模型的一個分片，這樣可以在給定裝置上的FW和BW過程之中動态加載所使用的參數。重要成員變量是：

  • model_shard ：Sequential模型的一個分片，每個分區包含一個或者多個層。
  • device ：計算裝置。
  • offload_device ：解除安裝目标裝置。
  • cpu_to_gpu_stream ：從cpu到gpu的CUDA流。
  • gpu_to_cpu_stream ：從gpu到cpu的CUDA流。

具體定義如下：

class ModelShard(nn.Module):
    """
    Wrap one shard of the model, make it possible to load parameters on the
    fly for the FW and BW pass on the given device.
    """

    def __init__(
        self, cpu_model_shard: nn.Module, device: torch.device, offload_device: torch.device, index: int,
    ):
        super().__init__()
        self.model_shard = cpu_model_shard # 模型分片
        self.index = index

        # Save all the parameter sizes to be able to restore them
        self.device = device # 計算裝置
        torch.cuda.device(self.device)

        self.offload_device = offload_device

        self.model_shard.to(offload_device) # 先把模型放到CPU上
        self._cpu_to_gpu_stream = torch.cuda.Stream(device=self.device) # 生成stream
        self._gpu_to_cpu_stream = torch.cuda.Stream(device=self.device) # 生成stream

其基礎函數可以分類如下：

  • 轉發函數，就是直接調用module對應的函數，比如forward，train。
  • 基礎拷貝函數，就是把module拷貝到參數對應的裝置之上，比如 to，to_device。
  • 功能函數，就是在特定的stream之上把module拷貝到特定的裝置上，比如forward_load方法就是專門在_cpu_to_gpu_stream之上把模型拷貝到device之上，即在前向傳播時候進行 CPU --> GPU 的拷貝。 
def forward(self, *inputs):  # type: ignore
    return self.model_shard(*inputs) if isinstance(inputs, tuple) else self.model_shard(inputs)

def to(self, device: torch.device) -> "ModelShard":  # type: ignore
    # Make sure that the lookahead and lookback shards are not captured by this call
    self.model_shard.to(device)
    return self

def train(self, mode: bool = True) -> "ModelShard":
    # Make sure that the lookahead and lookback shards are not captured by this call
    self.model_shard.train(mode)
    return self

def to_device(self) -> None:
    self.model_shard.to(device=self.device, non_blocking=True)

def forward_load(self, non_blocking: bool = True) -> None:
    with torch.cuda.stream(self._cpu_to_gpu_stream):
        # Restore all the parameter buffers
        self.model_shard.to(device=self.device, non_blocking=non_blocking)

# Ignore the following function for code coverage since the backward pass
# is triggered by C++ code and cannot be calculated when overriding
# autograd.Function
def backward_load(self, non_blocking: bool = True) -> None:  # pragma: no cover
    with torch.cuda.stream(self._cpu_to_gpu_stream):
        self.model_shard.to(self.device, non_blocking=non_blocking)

def forward_drop(self, non_blocking: bool = True) -> None:
    with torch.cuda.stream(self._gpu_to_cpu_stream):
        self.model_shard.to(self.offload_device, non_blocking=non_blocking)

# Ignore the following function for code coverage since the backward pass
# is triggered by C++ code and cannot be calculated when overriding
# autograd.Function
def backward_drop(self, non_blocking: bool = True) -> None:  # pragma: no cover
    with torch.cuda.stream(self._gpu_to_cpu_stream):
        self.model_shard.to(self.offload_device, non_blocking=non_blocking)

有了上面的基礎，我們來看看 OffloadModel 的 forward 方法。

Offload 在每一步訓練之中，會将一層（或一系列層）加載到GPU上，用于向前和向後傳遞，并根據需要将中間激活複制到GPU上。一旦給定分片的向前或向後傳播完成，它将再次移回CPU。是以我們看看在前向傳播之中如何加載GPU，并且何時移回CPU。

從設計思路可知，在每次疊代中，前向傳播從CPU複制每個模型分片到GPU，然後使用小批量（minibatch）資料計算前向傳播，并把模型分片從GPU複制回CPU。在後向傳播過程中，重複相同的過程。

前向傳播的具體邏輯是：

  • 如果設定了 _checkpoint_activation，則調用 OffloadFunction 把激活檢查點解除安裝到CPU之上，直接傳回（我們會在後續文章進行分析）。
  • 否則就執行 Offload，具體就是從前往後周遊模型，對于每一層，會做如下操作：
    • 前一層的激活放入計算裝置上。
    • 拿到本層的輸入，前一層的激活就是本層的輸入。
    • 用前一層的激活進行前向傳播計算。
    • 調用ShardSyncLayer 配置hook (discard/load slices FW and BW)。
    • 把本層計算結果插入到_activations，後續将成為下一層的輸入。
    • 把本層計算結果拷貝到CPU。
  • 傳回最後一個激活，就是整體計算結果，把結果放到GPU之上。

具體代碼如下：

def forward(self, *inputs: Any, **_: Any) -> Any:
    # `apply` calls the `forward` function of the `OffloadFunction` class
    # and the `forward` function calls `inputs` on the first model shard.
    # Please see https://pytorch.org/docs/stable/autograd.html#function for more details.

    # We need the second param to be a dummy input to enable the
    # backward pass to be triggered for integer inputs.
    
    # 注意，如果設定了_checkpoint_activation，就直接傳回了。
    if self._checkpoint_activation:
        return OffloadFunction.apply(*inputs, torch.tensor([], requires_grad=True), self)

    self._activations = []
    for index in range(-1, len(self.model_slices)): # 從前往後周遊模型
        if index >= 0:
            # 本層激活放入裝置上
            self._activations[index] = tuple([a.cuda() for a in list(self._activations[index])])
            inputs = self._activations[index] # 前一層的激活就是本層的輸入
            inputs = self.model_slices[index](*inputs) # 用前一層的激活進行前向傳播計算
            
        # Call the custom autograd hooks (discard/load slices FW and BW)
        # 調用ShardSyncLayer hook
        inputs = ShardSyncLayer.apply(inputs, index, self.model_slices, self)
        self._activations.append(inputs) # 把本層計算結果插入到_activations，後續将成為下一層的輸入
        if index >= 0:
            # 把本層計算結果拷貝到CPU
            self._activations[index] = tuple([a.cpu() for a in list(self._activations[index])])

    result = self._activations[-1] # 傳回最後一個激活，就是整體計算結果
    result = tuple([r.cuda() for r in result]) # 結果放到GPU之上
    return result[0] if len(result) == 1 else result

ShardSyncLayer 就是Hook，其是模型分片之間的同步點，這裡就是做加載/移除等工作，不涉及具體前向後向計算工作。

  • 在向前傳播中，它會移除前一個分片中的參數，并加載下一個分片的參數。
  • 在後向傳播時，它會做相反的動作。從設計思路可知，在後向傳播過程中，重複與前向傳播相同的過程。

ShardSyncLayer 不會更改或建立任何輸出，而是将輸入轉發到輸出。在代碼中幾個TODO注釋比較有意思，可能是開發者之間沒有做好工作交接，是以有疑惑 _。

# TODO(anj-s): Are these redundant in the backward pass?
# TODO(anj-s): Why do we need to do this?

具體如下：

class ShardSyncLayer(torch.autograd.Function):
    """
     The shard sync layer is a synchronization point between model shards.
     - In the forward pass, it drops parameters in the previous shard and
     loads parameters for the next shard.
     - In the backward pass, it does the reverse.
     It does not change or create any outputs at all, instead it just
     forwards the input as the output.
     NOTE: see https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function
     """

    @staticmethod
    @_conditional_amp_fwd_decorator  # type: ignore
    def forward(ctx: Any, inputs: Any, index: int, model_slices: Any, model_instance: Any) -> Any:
        drop_index = index # 本層
        load_index = index + 1 # 下一層
        max_slices = len(model_slices)

        if drop_index >= 0:
            # Move shard from device to offload device.
            model_slices[drop_index].forward_drop() # 解除安裝本層

        if load_index < max_slices:
            # Load shard from offload device to device.
            model_slices[load_index].forward_load() # 需要把下一層加載到GPU

        ctx.index = index
        ctx.model_slices = model_slices
        ctx.model_instance = model_instance

        return inputs if isinstance(inputs, tuple) else (inputs,)

    # Ignore the following function for code coverage since the backward pass
    # is triggered by C++ code and cannot be calculated when overriding
    # autograd.Function
    @staticmethod
    @_conditional_amp_bwd_decorator
    def backward(ctx, *grad_outputs):  # type: ignore # pragma: no cover

        # 從前向計算圖角度看，反向傳播需要把前向計算的下一層釋放，本層加載
        load_index = ctx.index # 本層
        drop_index = load_index + 1 # 下一層 
        model_slices = ctx.model_slices
        model_instance = ctx.model_instance

        # TODO(anj-s): Are these redundant in the backward pass?
        if drop_index == len(model_slices): # 如果是分區的最後一層
            # Drop the last activation since it is still on the CPU
            # after the loss.backward() call.
            # 把激活放回到GPU，但是這一步驟好像重複了，在fw之中已經做了，這也是代碼維護者的疑問
            model_instance._activations[-1] = tuple([a.cuda() for a in list(model_instance._activations[-1])])

        if drop_index < len(model_slices):
            # Move shard from device to offload device.
            model_slices[drop_index].backward_drop() # 把分片從計算裝置移動到offload裝置
            model_instance._activations[drop_index] = tuple(
                [a.cpu() for a in list(model_instance._activations[drop_index])]
            )

        if load_index >= 0:
            # Load shard from offload device to device.
            model_slices[load_index].backward_load() # 把分片從offload 裝置加載到計算裝置
            model_instance._activations[load_index] = tuple( # 激活加載到計算裝置
                [a.cuda() for a in list(model_instance._activations[load_index])]
            )

        # The returned variables need to mirror the forward inputs
        # TODO(anj-s): Why do we need to do this?
        if isinstance(grad_outputs, tuple):
            return grad_outputs[0], None, None, None

        return grad_outputs, None, None, None

我們總結一下邏輯圖，假設有兩個 ModelShard，每個 ModelShard 包括兩個層（下面的前/後指的是從前向傳播角度看的層之間關系）。

  • 前向傳播時候，ShardSyncLayer 會把計算圖之中前一個ModelShard參數移動到CPU，加載後一個ModelShard參數到GPU。
  • 後向傳播時候，ShardSyncLayer 會把計算圖之中後一個ModelShard參數移動到CPU，加載前一個ModelShard參數到GPU。
  • 前向後向傳播之中，ShardSyncLayer 的動作其實相同，但是邏輯相反。
[源碼分析] Facebook如何訓練超大模型 --- (3)

至此，Offload 分析完畢，下一篇介紹混合精度相關，敬請期待。

https://arxiv.org/pdf/2101.06840.pdf

https://www.deepspeed.ai/tutorials/zero-offload/

DeepSpeed: Extreme-scale model training for everyone

https://www.microsoft.com/en-us/research/blog/zero-infinity-and-deepspeed-unlocking-unprecedented-model-scale-for-deep-learning-training/

https://www.microsoft.com/en-us/research/blog/zero-2-deepspeed-shattering-barriers-of-deep-learning-speed-scale/

https://www.marktechpost.com/2021/02/01/microsoft-and-the-university-of-california-merced-introduces-zero-offload-a-novel-heterogeneous-deeplearning-training-technology-to-train-multi-billion-parameter-models-on-a-single-gpu/

上一篇: netty系列之:請netty再愛UDT一次
下一篇: mysql常用運算符和常用函數1、常用算術運算符2、常用比較運算符3、邏輯運算符4、指派運算符5、常用控制流程函數6、常用字元串函數7、數學函數8、時間日期函數9、類轉換函數10、聚合函數