hfai.nn.CPUOffload | 模型训练的显存节省利器

在深度学习模型的训练中，研究者与开发者们尝尝会碰到显存不足的问题 (OOM, out of memory) ，比如模型参数规模大，或者训练过程中产生的额外开销大，又或者是程序代码的问题，没有足够的显存资源，对我们科研与开发产生了诸多限制。

以往，我们通过代码优化、梯度累计、半精度等等一系列方法，以降低深度学习模型训练的显存需求，然而随着模型参数规模的发展，业务数据越来越复杂，显存需求快速增长，那么除了这些 trick 之外，还有什么既简单、又高效的显存节省方法吗？本期文章介绍的主角， hfai.nn.CPUOffload，给您提供一条不一样的显存节省之路。

那么 hfai.nn.CPUOffload 为什么可以节省显存？如何使用？它和 PyTorch 原始版本有何不同？本期文章将为大家分享 CPUOffload 设计背后的故事，讲述幻方 AI 设计 CPUOffload 的理念，展现 CPUOffload 的性能和便捷。

概述

训练模型的过程中，显存占用主要包括如下三部分

1. 权重矩阵
2. 前向传播的中间向量
3. 反向传播的梯度矩阵

我们以 ResNet50 网络在输入尺寸为（256，3，224，224）时的计算过程作为示例，分析上述三个部分的显存占用。通过 hfai.utils.profile_memory 工具，我们可以很容易的获得模型各部分的细节，如下图所示（由于网络层数较多，只展示部分网络）：

每一行为一个 nn.Module 类型网络的详情信息：

• module name：nn.Module 类型网络的名称
• type：网络的类型
• parameter size: 网络参数量
• activation size: 中间变量的 tesnor 大小（不包括参数）
• #calls：被调用的次数
• input shape：输入的 tensor 形状
• output shape：输出的 tensor 形状
• peak mem：峰值显存，forward 过程中的峰值显存减去 forward 之前已经占用的显存
• forward time：网络 forward 的时间，如果多次调用，则时间累加

最后三行统计了模型的参数总量、中间变量总量、单次迭代前向传播的时间。从上图可以看到 ResNet50 在批大小为 256 时，参数总量为 97 MiB，中间变量为 20974 MiB，中间变量占比较大。而反向传播由于需要保存权重矩阵的梯度和中间变量的梯度，往往需要前向传播的 2 倍或更多的显存空间。

如果我们要减少整体显存占用，不考虑模型的分开存放，有如下两种思路：

第一种思路是使用更小的模型。当数据量小的时候，我们可以使用更小的模型进行拟合；但当数据量大、分布广时，我们需要使用更大的模型，以达到更好的拟合效果。

第二种思路是减小中间变量的显存占用。网络前向传播产生中间变量，反向传播使用中间变量，因此我们可以在中间变量没有被使用的时间，将其移到 CPU 上，以减少显存占用，这个过程我们称之为 CPUOffload。

CPUOffload 原理

中间变量在 PyTorch 中被称为 save tensor, 我们通过 torch.autograd.graph.saved_tensors_hooks 钩子可以在中间变量产生时（前向传播过程中）对其进行打包（pack），在中间变量被使用时（反向传播过程中）对其进行解包（unpack）。一个简单的使用例子如下：

def pack_hook(x):
    return (x.device, x.cpu())

def unpack_hook(packed):
    device, tensor = packed
    return tensor.to(device)

x = torch.randn(5, requires_grad=True)
with torch.autograd.graph.saved_tensors_hooks(pack, unpack):
    y = x * x

y.sum().backward()
torch.allclose(x.grad, (2 * x))

上面的代码是在中间变量产生时，将中间变量移动到 CPU 上，在需要使用中间变量时移动到 GPU 上，从 GPU 移动到 CPU 的过程为 Host2Device（H2D），从 CPU 移动到 GPU 的过程为 Device2Host（D2H）。

PyTorch 提供了 torch.autograd.graph.save_on_cpu 同样实现了此功能。我们可以很方便的对模型中的某些中间变量进行 CPUOffload，一个简单的使用例子如下：

class TorchCPUOffload(nn.Module):
    def __init__(self, module):
        super().__init__()
        self.module = module

    def forward(self, *args, **kwargs):
        with torch.autograd.graph.save_on_cpu(pin_memory=True):
            return self.module(*args, **kwargs)

model = nn.Sequential(
    nn.Linear(10, 100),
    TorchCPUOffload(nn.Linear(100, 100)),
    nn.Linear(100, 10),
)

x = torch.randn(10)
loss = model(x).sum()
loss.backward()

上面的代码定义了一个包括三层全连接层的网络，通过 torch.autograd.graph.save_on_cpu 的方法，第二层全连接层所产生的中间变量，会在不用时移动到 CPU 中，以减少显存占用。

PyTorch CPUOffload 分析

如上述所述，我们可以实现中间变量的 CPUOffload，但在实际优化测试过程中，我们发现如下问题：

第一个问题是 torch.autograd.graph.save_on_cpu 无法自定义 CPUOffload 比例，进入该环境的模型所产生的中间变量，只能全部 CPUOffload。如果想要细粒度的控制 CPUOffload 的比例，则需要开发者自己计算模型各层的中间变量占比，并在代码中增加对应的钩子，以达到按比例 CPUOffload，实现起来较为复杂。

第二个问题是 torch.autograd.graph.save_on_cpu 速度慢，CPUOffload 是以时间换取显存，但是使用 pytorch 默认的版本，显存节省的成本过高。我们以 ResNet50 为例进行速度和显存的测试，具体数据如下表格所示，可以看到在节省百分之十的显存时，模型训练速度慢了将近一倍。

（以上实验在 python3.6、 torch1.10.0+cu113 环境下完成）

hfai.nn.CPUOffload 使用方法

为了解决如上的问题，幻方研发设计了 hfai.nn.CPUOffload，与 pytorch torch.autograd.graph.save_on_cpu 的上述问题对应，hfai.nn.CPUOffload 进行了如下优化：

首先，hfai.nn.CPUOffload 提供了 offload_ratio 参数，自动将对应比例的中间变量进行 CPUOffload，使得开发者可以进行更加灵活的选择和更加简单的开发。

另外 hfai.nn.CPUOffload 通过使用独立的 torch.cuda.Stream 进行 H2D 和 D2H，使得整体速度有了较大提升。

以 ResNet50 为例子，这里为大家展示 hfai.nn.CPUOffload 的使用方法和性能。

from hfai.nn import CPUOffload
import torchvision

model = torchvision.models.resnet50().cuda()
x = torch.randn(256, 3, 244, 224).cuda()
with CPUOffload(offload_ratio=0.1, tag='resnet'):
    out = model(x)

hfai.nn.CPUOffload 的使用十分简单，在上面的代码中，通过简单的一行代码，即可实现将特定比例的中间变量进行 CPUOffload，其中 offload_ratio 指定了 CPUOffload 的比例，tag 标签标识了当前 CPUOffload 的模型名称。

另外 hfai.nn.CPUOffload 的使用十分灵活，我们可以通过不同的标签来进行不同模型的 CPUOffload；在不同的迭代中，我们也可以为相同的模型选择不同的 CPUOffload 比例。

hfai.nn.CPUOffload 性能对比

我们将 hfai.nn.CPUOffload 与 torch.autograd.graph.save_on_cpu 在相同模型和相同输入的情况下进行了比较，比较结果如下表所示，可以看到在不同的 offload_ratio 下，hfai.nn.CPUOffload 的训练时间都比 torch.autograd.graph.save_on_cpu 更快。在更低的 offload_ratio 下，hfai.nn.CPUOffload 的独立 Stream 能够保证计算过程和 tensor 移动拥有更好的并行性，因此加速更加明显，在 offload_ratio 为 0.1 时达到了 43.0% 的提速。

看到这里，您是不是也跃跃欲试了呢？我们欢迎广大研究者和开发者安装体验 hfai，来试用这一深度学习训练神器！