[源码分析] 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 包括两个层（下面的前/后指的是从前向传播角度看的层之间关系）。