介绍

在本文中，我们将探索 10 种易于使用的技术来减少 PyTorch 中的内存使用量。 这些技术是累积性的，这意味着我们可以将它们相互叠加应用。

我们将开始使用 PyTorch 的 Torchvision 库中的Torchvision library来提供简单的代码示例，可以在自己的计算机上执行这些示例，而无需下载和安装太多代码和数据集依赖项。 独立的基线训练脚本由约 100 行代码组成（忽略空格和代码注释）

以下是我们将要介绍的部分和技术的概述：

1.微调视觉转换器

2.自动混合精度训练

3.低精度训练

4.减少批量大小的训练

5.梯度累积和微批次

6.选择更精简的优化器

7.在目标设备上实例化模型

8.分布式训练和张量分片

9.激活检查点
10.参数卸载

11.将上述应用综合起来

虽然我们在这里使用Vision Transformers（the ViT-L-16 model from the paper An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale），但本文中使用的所有技术也可以转移到其他模型 ：卷积网络、大型语言模型 (LLM) 等。

此外，在使用上述Vision Transformers示例一次介绍一种技术后，我们将应用这些技术在文本分类任务上训练 BigBird-Roberta LLM。 如果没有这些技术，就不可能在消费类硬件上训练这样的模型。

1.微调 Vision Transformer

为了简化实验的 PyTorch 代码，我们将引入开源 Fabric 库，它允许我们通过少量的代码应用各种先进的 PyTorch 技术（自动混合精度训练、多 GPU 训练、张量分片等）。 （而不是几十行）代码。

简单的 PyTorch 代码和使用 Fabric 的修改后的代码之间的区别很微妙，只涉及很小的修改，如下面的代码中突出显示的：

 

如上所述，这些微小的变化现在提供了一个利用 PyTorch 中的高级功能的途径，正如我们稍后将看到的，而无需重组任何现有代码。

总结一下上图，将普通 PyTorch 代码转换为 PyTorch+Fabric 的主要 3 个步骤如下：

作为快速健全性检查，使用普通 PyTorch 和 PyTorch with Fabric 的预测性能和内存消耗保持完全相同（+/- 由于随机性而出现预期波动）：

普通 PyTorch (01_pytorch-vit.py)：

已用时间 17.94 分钟
使用内存：26.79 GB
测试准确度95.85%
PyTorch 与 Fabric (01-2_pytorch-fabric.py)

已用时间 17.88 分钟
使用内存：26.84 GB
测试准确率96.06%
作为可选练习，可以尝试代码并替换

 2.基于Fabric的自动混合精度

现在只需更改一行代码就可以尝试先进的技术，例如混合精度和分布式训练。

我们将从混合精度训练开始，这已成为最近训练深度神经网络的规范。

应用混合精度训练

我们可以应用混合精度训练，只需进行一个小的修改，改变

fabric = Fabric(accelerator="cuda", devices=1)

更改为：

fabric = Fabric(accelerator="cuda", devices=1, precision="16-mixed")

结果，我们的内存消耗从 26.84 GB 减少到 18.21 GB，而没有牺牲预测精度，如下所示。

作为回报，混合精度训练不仅可以减少内存使用，还可以将运行时间减少 6 倍（从 17.88 分钟减少到 3.45 分钟），这是一个很好的额外好处； 然而，这篇文章的重点是内存消耗，以免使其进一步复杂化。

什么是混合精度训练？

混合精度训练同时使用 16 位和 32 位精度，以确保精度不损失。 16 位表示形式的梯度计算比 32 位格式快得多，并且可以节省大量内存。 这种策略是有益的，特别是当我们内存或计算受限时。

它被称为“混合”而不是“低”精度训练，因为我们不会将所有参数和操作转移到 16 位浮点数。 相反，我们在训练期间在 32 位和 16 位操作之间切换，因此出现了术语“混合”精度。

如下图所示，混合精度训练包括将权重转换为较低精度（FP16）以加快计算速度，计算梯度，将梯度转换回较高精度（FP32）以实现数值稳定性，并用缩放后的权重更新原始权重。 梯度。

这种方法可以实现有效的训练，同时保持神经网络的准确性和稳定性。

 有关更多详细信息，我建议阅读更详细的独立文章《 Accelerating Large Language Models with Mixed-Precision Techniques》，其中更深入地探讨了基本概念。

3）低精度训练

还可以更进一步，尝试以“完整”的较低 16 位精度运行（而不是混合精度，它将中间结果转换为 32 位表示。）

可以通过改变来实现较低精度的训练

fabric = Fabric(accelerator="cuda", precision="16-true")

但是，运行此代码时，将在损失中遇到 NaN 值：

Epoch: 0001/0001 | Batch 0000/0703 | Loss: 2.4105
Epoch: 0001/0001 | Batch 0300/0703 | Loss: nan
Epoch: 0001/0001 | Batch 0600/0703 | Loss: nan

这是因为常规 16 位浮点数只能表示 -65,504 到 65,504 之间的数字：

import torch
​
torch.finfo(torch.float16)
finfo(resolution=0.001, min=-65504, max=65504, eps=0.000976562, smallest_normal=6.10352e-05, tiny=6.10352e-05, dtype=float16)

因此，为了避免 NaN 问题，可以使用“bf16-true”设置。

fabric = Fabric(accelerator="cuda", precision="bf16-true")

因此，可以将内存消耗进一步降低至 13.82 GB（同样，不牺牲准确性）：

bf16是谷歌为机器学习和深度学习应用程序开发了这种格式，特别是在他们的张量处理单元（TPU）中。 与传统的 float16 格式相比，Bfloat16 扩展了动态范围，但代价是精度降低。

扩展的动态范围有助于 bfloat16 表示非常大和非常小的数字，使其更适合可能遇到各种值的深度学习应用。 然而，较低的精度可能会影响某些计算的准确性或在某些情况下导致舍入误差。 但在大多数深度学习应用中，这种精度降低对建模性能的影响很小。

虽然 bfloat16 最初是为 TPU 开发的，但现在多种 NVIDIA GPU 也支持这种格式，首先是 A100 Tensor Core GPU，它是 NVIDIA Ampere 架构的一部分。

可以通过以下代码检查的GPU是否支持bfloat16：

>>> import torch
>>> torch.cuda.is_bf16_supported()

4）减少batch size

解决一个大问题：为什么不简单地减少批量大小呢？ 这通常始终是减少内存消耗的一个选项。 然而，它有时会导致更差的预测性能，因为它改变了训练动态。 

不管怎样，减少批量大小，看看这对结果有何影响。 事实证明，我们可以将批量大小降低到 16，从而将内存消耗降低到 5.69 GB，而不会牺牲性能：

5）使用梯度累积创建微批次 

梯度累积是一种在训练期间虚拟增加批大小的方法，当可用 GPU 内存不足以容纳所需的批大小时，这非常有用。 请注意，这仅影响运行时，而不影响建模性能。

在梯度累积中，计算较小批次的梯度，并在多次迭代中累积（通常求和或平均），而不是在每个批次后更新模型权重。 一旦累积的梯度达到目标“虚拟”批量大小，模型权重就会使用累积的梯度进行更新。

为了实现梯度累积，只需要对前向和后向传递进行两个小修改：

使用 16 的有效批量大小和 4 个累积步骤意味着将使用 4 的实际批量大小（因为 16 / 4 = 4）。

此技术的缺点是它将运行时间从 3.96 分钟增加到 12.91 分钟。

当然，甚至可以更小，e 16 个累积步骤。 这将导致微批量大小为 1，进一步减少内存大小（约 75%）

6) 使用更精简的优化器

流行的 Adam 优化器带有附加参？ 例如，Adam 对于每个模型参数都有 2 个额外的优化器参数（均值和方差）。

因此，通过将 Adam 替换为像 SGD 这样的无状态优化器，我们可以将参数数量减少 2/3，这在使用视觉 Transformer 和 LLM 时非常重要。

普通 SGD 的缺点是它通常具有较差的收敛特性。 因此，我们将 Adam 替换为 SGD，并引入余弦衰减学习率调度器来补偿这一点并实现更好的收敛。

简而言之，我们将交换之前使用的 Adam 优化器：

optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
#替换为
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
​
num_steps = NUM_EPOCHS * len(train_loader)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=num_steps)

通过这一更改，我们能够在保持约 97% 的分类准确率的同时实现峰值内存消耗：

7) 以所需的精度在目标设备上创建模型

 当我们在 PyTorch 中实例化模型时，通常会先在 CPU 设备上创建它，然后将其传输到目标设备上并将其转换为所需的精度：

model = vit_l_16(weights=ViT_L_16_Weights.IMAGENET1K_V1)
model.cuda().float16()

考虑到 CPU 上的全精度中间模型表示，这可能效率很低。 相反，我们可以使用 Fabric 中的 init_module 上下文直接在目标设备（例如 GPU）上以所需的精度创建模型：

import lightning as L
​
fabric = Fabric(accelerator="cuda", devices=1, precision="16-true")
​
with fabric.init_module():
    model = vit_l_16(weights=ViT_L_16_Weights.IMAGENET1K_V1)

在这种特定情况（模型）中，前向传递期间的峰值内存大于其全精度表示中的模型大小。 因此，我们将仅针对模型加载本身对 Fabric.init_module 方法进行基准测试。

不带 init_module 的 GPU 峰值内存：1.24 GB (07_01_init-module.py)

带 init_module 的 GPU 峰值内存：0.65 GB (07_03_init-module.py)

从上面的结果我们可以看出，在这种情况下，init_module 将模型加载的峰值内存需求降低了 50%。

8) 分布式训练和张量分片

我们要尝试的下一个修改是多 GPU 训练。 如果我们有多个 GPU 可供使用，这将变得非常有益，因为它可以让我们更快地训练模型。

然而，在这里，我们主要感兴趣的是内存节省。 因此，我们将使用一种更先进的分布式多 GPU 策略，称为完全分片数据并行 (FSDP)，它利用数据并行和张量并行来跨多个设备分片大型权重矩阵。

model = vit_h_14(weights=ViT_H_14_Weights.IMAGENET1K_SWAG_E2E_V1)

#更改为
fabric = Fabric(accelerator="cuda", 
    devices=4, strategy="fsdp", precision="16-mixed")

这会将内存从 22.63 GB 减少到 19.83 GB。

（代码示例 08a_fsdp-defaults.py）

请注意，这些是默认设置。 然而，FSDP 的真正好处是当我们使用具有超过 1 亿个参数的 LLM 和 ViT 时，因为在 GPU 之间分割参数和计算。 为了获得更多控制，我们可以通过自定义自动换行策略来尝试两种修改：

from torchvision.models.vision_transformer import EncoderBlock
auto_wrap_policy = partial(transformer_auto_wrap_policy,
    transformer_layer_cls={EncoderBlock}
)
strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy)
fabric = Fabric(accelerator="cuda", devices=4,
    strategy=strategy, precision="16-mixed"
)

这将内存从 19.83 GB 进一步减少到 17.63 GB

为了更好地控制考虑分片的层的最小大小，我们还可以使用基于大小的自动换行策略：

这会将默认的换行大小从 1 亿降低到 200 万：

auto_wrap_policy = partial(
    size_based_auto_wrap_policy, min_num_params=2_000_000
)
strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy)
fabric = Fabric(accelerator="cuda",
    devices=4, strategy=strategy, precision="16-mixed"
)

了解数据并行性和张量并行性

在数据并行性中，小批量被划分，模型的副本在每个 GPU 上可用。 由于多个 GPU 并行工作，此过程可加快模型训练速度。
简而言之，它是如何工作的：

相同的模型被复制到所有 GPU 上。

然后，每个 GPU 都会被输入不同的输入数据子集（不同的小批量）。

所有 GPU 独立执行模型的前向和后向传递，计算自己的局部梯度。

然后，收集所有 GPU 的梯度并对其进行平均。

然后使用平均梯度来更新模型的参数。

这种方法的主要优点是速度。 由于每个 GPU 都与其他 GPU 同时处理唯一的小批量数据，因此可以在更短的时间内对更多数据进行模型训练。 这可以显着减少训练模型所需的时间，尤其是在处理大型数据集时。

然而，数据并行性有一些局限性。 最重要的是，每个 GPU 必须拥有模型及其参数的完整副本。 这限制了我们可以训练的模型的大小，因为模型必须适合单个 GPU 的内存——这对于现代 ViT 或 LLM 来说是不可行的。

与涉及将小批量拆分到多个设备的数据并行不同，张量并行将模型本身划分到 GPU 上。 在数据并行性中，每个 GPU 都需要适应整个模型，这在训练较大模型时可能是一个限制。 另一方面，张量并行性允许通过分解模型并将其分布在多个设备上来训练对于单个 GPU 来说可能太大的模型。

 

它是如何工作的？ 想想矩阵乘法。 有两种分配方式：按行或按列。 为了简单起见，我们考虑按列分布。 例如，我们可以将大型矩阵乘法运算分解为单独的计算，每个计算都可以在不同的 GPU 上执行，如下图所示。 然后将结果连接起来以获得原始结果，从而有效地分配计算负载。

9) 激活检查点
为了进一步减少神经网络计算期间的内存使用，我们可以添加梯度检查点（也称为激活检查点）。 该方法在前向传递期间选择性地消除某些层激活，然后在后向传递中重新计算它们。 这种方法本质上会牺牲一些计算时间来节省内存。
换句话说，层的输入和输出在前向传递后保留在内存中，但模块内计算涉及的任何中间张量都被释放。 当为这些设置检查点的模块计算向后传递时，将重新计算先前清除的张量。

我们可以通过将“activation_checkpointing=EncoderBlock”添加到我们之前使用的 FSDP 策略中来使用激活检查点：

auto_wrap_policy = partial(
    transformer_auto_wrap_policy, transformer_layer_cls={EncoderBlock}
)
strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy,
    activation_checkpointing=EncoderBlock
)

fabric = Fabric(accelerator="cuda", 
    devices=4, strategy=strategy
)

这将内存消耗从 17.23 GB 降低到 9.03 GB。 然而，这将运行时间从 18.95 分钟略微增加到 22.58 分钟。

10) 参数卸载 

auto_wrap_policy = partial(transformer_auto_wrap_policy, 
    transformer_layer_cls={EncoderBlock}
)
strategy = FSDPStrategy(
    auto_wrap_policy=auto_wrap_policy,
    activation_checkpointing=EncoderBlock,
    cpu_offload=True
)
fabric = Fabric(accelerator="cuda",
    devices=4, strategy=strategy, precision="16-mixed"
) 

这将内存消耗从带有激活检查点的 9.03 GB 减少到带有额外 CPU 卸载的 6.68 GB。 然而，运行时间从 22.58 分钟大幅增加到 101.53 分钟。

结论

本文展示了 10 种减少 PyTorch 模型内存消耗的技术。 当将这些技术应用于Transformer时，我们将单个 GPU 上的内存消耗减少了 20 倍。 我们发现跨 GPU 的张量分片甚至可以降低内存消耗。 同样的优化还可以仅使用 4 GB 峰值 GPU RAM 来训练 BigBird LLM。

这些技术都不是特定于模型的，并且几乎可以与任何 PyTorch 训练脚本一起使用。 使用开源 Fabric 库，大多数优化都可以通过一行代码来启用。