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

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

0x00 摘要

我們在前文介紹過，微軟 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)

0x01 ZeRO-Offload

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

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

1.1 設計原則

ZeRO-offload 屬於CPU卸載技術，就是當GPU內存已滿時，可以將暫時未使用的數據卸載到CPU，並在以後需要時將其讀回（Rhu等人，2016）。ZeRO-offload 基於三個原則來設計：效率、可伸縮性和可用性。其背後的關鍵技術是：在 ZeRO-2 基礎上將優化器計算，優化器狀態和梯度卸載到 CPU 內存。這種方法讓 ZeRO-Offload 能最大程度降低拷貝至 CPU 導致的計算效率損失，同時還實現了與原始ZeRO-2相同的效率，有時甚至更好。研究人員已經可以確定 CPU 和 GPU 之間數據分區和最佳計算策略。該方法涉及到的流程包括如何將梯度、優化器狀態和優化器計算分散到 GPU，以及如何在 GPU 上進行向前和向後計算。

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

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

1.2 ZeRO

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上計算的其餘參數更新。

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

0x02 卸載策略

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優化器設計的。

2.1 數據流圖

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

GPU和CPU之間的卸載策略可以用這個圖的雙向分區來表示，比如分區中的計算節點將在擁有該分區的設備上執行，而該分區中的數據節點將存儲在擁有該分區的設備上。GPU和CPU之間必須通信的總數據量由兩個分區上運行的邊的權重給出。有許多方法可以對該圖進行分區。比如可以使用第一原理簡化數據流圖，以減少基於三個不同效率指標的可能選擇的數量：i）CPU計算量開銷，ii）通信開銷，以及iii）內存節省。

2.2 限制CPU計算

CPU計算吞吐量比GPU計算吞吐量慢多個數量級。因此，將大型計算圖卸載到CPU將嚴重限制訓練效率。因此，我們必須避免將計算密集型組件卸載到CPU上。

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

基於這個簡單的觀察，我們將數據流圖中的前向和後向節點融合爲一個超級節點（FWD-BWD），並將其分配到GPU上。

2.3 最小化計算量

我們接下來分析最小化計算量（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。

2.4 最大化內存節約

我們接下來看看如何最大化內存節約（Maximizing Memory Savings）。

在簡化數據流圖以最小化通信量之後，只剩下g16和Update Super需要被分配。請注意，在這一點上，所有的分區結果都會導致最小的通信量，所以我們可以進一步調整選擇，以最大限度地節省GPU的內存。表1顯示了所有有效的分區策略所帶來的內存節省，這些策略使通信量最小。通過將g16和Update Super卸載到CPU，可以實現8倍的最大內存節省。

2.5 唯一最優化策略

ZeRO-Offload在CPU內存中分配所有的fp32模型狀態以及fp16梯度，它也在CPU中計算參數更新。fp16的參數保留在GPU上，前向和後向的計算也在GPU上完成。

我們通過簡化我們的數據流圖來得出這個卸載策略，並排除了所有其他的分區策略，因爲其他策略或者不能限制CPU的計算，或者無法最小化通信量，或無法最大限度地節省內存。因此，ZeRO-Offload不僅在上述指標上是最優的，而且是唯一的；不可能有其他策略能比ZeRO-Offload節省更多的內存，而不增加CPU的計算複雜性或產生額外的GPU-CPU通信量。

2.6 ZeRO-Offload Schedule

在這一節中，我們將討論基於我們的卸載策略，如何在單GPU系統上實現ZeRO-Offload的具體計算和通信schedule。然後，我們將展示如何通過將我們的卸載策略與ZeRO數據並行和模型並行結合起來，把這個schedule擴展到多GPU系統上有效工作。

2.6.1 單機計劃

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) 重疊起來，通信隱藏在計算之中，這樣可以提高訓練效率。

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

2.6.2 多節點計劃

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模型參數、梯度和優化器狀態。

ZeRO-Offload數據並行調度的詳細信息如代碼圖所示。上述all gather操作在代碼圖中顯示爲一系列廣播操作。

0x03 FairScale Offload 使用

3.1 思路

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

一般來說，大型模型往往會導致OOM錯誤，而FairScale OffloadModelAPI使用戶能夠在有限的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上的梯度進行更新，具體可以參見下面的示例圖。

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複製的成本。

3.2 使用

具體使用樣例如下，首先會進行常規配置，並且定義了一個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()

3.3 配置

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``*.    

0x04 源碼

4.1 構建

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

4.1.1 初始化

因爲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

4.1.2 切片

初始化代碼之中使用了_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

4.2 ModelShard

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

4.2.1 定義

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

4.2.2 功能函數

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

• 轉發函數，就是直接調用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)

4.3 前向傳播

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

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

4.3.1 前向傳播

從設計思路可知，在每次迭代中，前向傳播從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

4.3.2 Hook

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 的動作其實相同，但是邏輯相反。

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

0xFF

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/