[原始碼分析] 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/