降龙十八掌：这套优化transformer内存占用的组合技值得收藏

很多时候，内存限制会阻碍 ViT 以及 LLM 的训练，这篇文章介绍了 9 种减少内存消耗峰值的方法。难能可贵的是，这几种方法可以同时使用，就好像降龙十八掌中最后一掌，正是将前几张组合起来就能打出最强大的效果。

峰值内存消耗是训练深度学习模型（如视觉 Transformer 和 LLM）时的常见瓶颈。本文提供了一系列可以在不牺牲建模性能和预测精度的情况下，将 PyTorch 中的内存消耗降低到约 1/20 的技术。

以 PyTorch 的 Torchvision 库中的视觉 transformer 为基础，本文作者编写了大约 100 行代码的训练脚本，并且所有代码示例都可以在 GitHub 上找到。

以下是本文将要介绍的技术名称：

• 微调 vision transformer
• 自动混合精度训练
• 低精度训练
• Reduced Batch Size 训练
• 梯度积累与 Microbatches
• 选择更精简的优化器
• 在目标设备上实例化模型
• 分布式训练与张量共享
• 参数卸载
• 以上九种方法结合起来，就形成了一种可以用于 llm 的综合方法，也可以称之为第十种方法。

这些方法是互相解耦的，可以将它们叠加在一起使用。

本文在实验中使用的 ViT 为 ViT-L-16 模型。在依次将上述方法添加后，研究者将训练 BigBird-Roberta LLM 来执行文本分类任务。这些技术使得在消费类硬件上训练这样的模型成为可能。

微调 vision transformer

为了简化实验中的 PyTorch 代码，本文使用了开源库 ——Fabric，十几行代码就能应用各种先进的 PyTorc 技术（自动混合精度训练、多 GPU 训练、张量分片等）。

原生 PyTorch 代码和修改后的使用 Fabric 的代码之间的区别很微妙，只有较小的修改，如下面的代码所示：

如上所述，改动虽然不大，但是可以方便的使用 PyTorch 中的高级功能，而无需重新构造任何现有代码。

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

1. 导入 Fabric 并实例化 Fabric 对象。
2. 使用 Fabric 设置模型、优化器和数据加载程序。
3. 调用 fabric.backward () 构造损失函数，而不是通常使用的 loss.backward ()

使用普通 PyTorch 和 PyTorch with Fabric 的性能和内存消耗几乎完全相同：

Plain PyTorch (01_pytorch-vit.py):

Time elapsed 17.94 minMemory used: 26.79 GBTest accuracy 95.85%

PyTorch with Fabric (01-2_pytorch-fabric.py)

Time elapsed 17.88 minMemory used: 26.84 GBTest accuracy 96.06%

也可以将下面的代码

model = vit_l_16(weights=ViT_L_16_Weights.IMAGENET1K_V1)

替换为

model = vit_l_16(weights=None)

替换后，将不再是微调，而是从头开始训练相同的 ViT 架构，预测准确率会从 96% 以上下降到约 60%：

自动混合精度

上一节使用 Fabric 修改了 PyTorch 代码，在此基础上，使用混合精度和分布式训练，也只需更改一行代码。

应用混合精度训练

应用混合精度训练，只需一个小的修改，将下面这行代码

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

替换为

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

之后，在不牺牲预测精度的情况下，内存消耗从 26.84GB 减少到 18.21GB，如下所示：

01-2_pytoch-fabric.py 和 02_mixed-precision.py 的结果对比

此外，混合精确训练不仅减少了内存使用，还将运行时间减少了 6 倍（从 17.88 分钟减少到 3.45 分钟），这可以说是意外收获。

什么是混合精度训练？

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

之所以被称为「混合」而不是「低」精度训练的原因是，并不会将所有参数和操作都转移成 16 位浮点数。实际上，在训练期间会在 32 位和 16 位运算之间切换。

如下图所示，混合精度训练可以分解为：将权重转换为较低精度（如 FP16）以实现更快的计算、计算梯度、将梯度转换回较高精度（FP32）以实现数值稳定性，以及用缩放的梯度更新原始权重等几个步骤。

这种方法在保证训练有效的前提下，还能保持神经网络的准确性和稳定性。

感兴趣的读者还可以在本文作者的另一篇文章：《使用混合精度技术加速大型语言模型》中获得更多底层概念。

文章地址：https://lightning.ai/pages/community/tutorial/accelerating-large-language-models-with-mixed-precision-techniques/

低精度训练

还可以更进一步，尝试以「完全」较低的 16 位精度运行，而不是混合精度。

将下面这行代码

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

替换为

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

但需要注意的是，这样会在训练中产生 NaN 值：

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

这是因为常规的 16 位浮点只能表示 - 65504 和 65504 之间的数字：

In [1]: import torch
In [2]: torch.finfo(torch.float16)Out[2]: 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（同样，在不牺牲准确性的情况下）：

将 03_bfloat16.py 与之前的代码的结果进行比较

什么是 Bfloat16？

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

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

虽然 bfloat16 最初是为 TPU 开发的，但这种格式从 A100 Tensor Core GPU 开始，也得到了其之后的 NVIDIA GPU 的支持。

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

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

减少批大小

减少批大小通常是减少内存消耗的一个有效方法。然而，它有时会导致较差的预测性能，因为这样要改变训练动态。

无论哪种方式，需要探讨减少批量大小对结果有何影响。事实证明，可以在不牺牲性能的情况下将批大小降低到 16，从而将内存消耗降至 5.69 GB：

将 04_lower-batchsize.py 与以前的代码进行比较。

梯度积累与微批

梯度累积是一种在训练过程中虚拟增加批大小的方法，当可用的 GPU 内存不足以容纳所需的批量大小时，这是非常有用的。并且这种方法只会在运行时产生影响，建模性能并不会受到影响。

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

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

05_gradient-acum.py 中的代码修改

本文作者的另一篇文章《使用梯度累积在单个 GPU 上微调 LLM》，更详细地介绍了梯度累积的细节。

文章地址：https://lightning.ai/blog/gradient-accumulation/

有效批大小为 16，并且累积步数为 4，意味着实际批大小为 4（因为 16/4=4）。

05_gradient-acum.py 的结果

这种技术的缺点是运行时间从 3.96 分钟增加到 12.91 分钟。

值得注意的是，批大小最小可以减少到 1，进一步减少 75% 的内存消耗。

使用更精简的优化器

时下流行的 Adam 优化器其实附带了额外的参数，例如，Adam 为每个模型参数提供了 2 个额外的优化器参数（平均值和方差）。

因此，通过将 Adam 与 SGD 等无状态优化器进行交换，可以将参数数量减少 2/3，这在使用 ViT 和 LLM 时非常重要。

普通 SGD 的缺点是收敛性较差。因此，Adam 与 SGD 交换后，需要引入余弦衰减学习速率调度器来进行补偿。

简而言之，通过将以下代码

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% 分类准确率的同时减少峰值内存消耗：

06_sgd-with-scheduler.py 的结果

在目标设备上创建模型

在 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）
• GPU 带 init_module 的峰值内存：0.65 GB（07_03_init-module.py）

可以看到，在这种情况下，init_module 将模型加载的峰值内存需求减少了 50%。

有关 init_module 的更多详细信息，可以参阅这篇关于大型模型的高效初始化的的文章。

文章地址：https://lightning.ai/pages/community/efficient-initialization-of-large-models/

分布式训练与张量共享

下一个修改是多 GPU 训练。多个 GPU 可供使用是有效的，因为这样做可以更快地训练模型。

然而，本文探讨的是内存节省。因此，需要一种更先进的分布式多 GPU 策略，称为完全共享数据并行（FSDP），该策略利用数据并行性和张量并行性在多个设备上共享大权重矩阵。

但是如果模型已经很小了，例如将此技术添加到上面第 7 节的代码中时，是几乎看不到任何效果的。因此，为了纯粹地关注分片的效果，可以与第 1 节中的全精度基线进行比较。

将以下代码

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

替换为

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)

08_fsdp 与 - 01-2.py 的结果

除了手动定义，请也可以使用以下方法，自动确定要分割哪些层：

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

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

在数据并行中，mini-batch 需要继续被划分，并且每个 GPU 上都有一份模型副本。由于多个 GPU 并行工作，能够加快模型训练。

以下是工作原理：

1. 在所有 GPU 中复制相同的模型。
2. 然后，每个 GPU 被馈送输入数据的不同子集（不同的小批量）。
3. 所有 GPU 独立地执行模型的前向和后向传递，计算各自的局部梯度。
4. 然后，收集梯度并对所有 GPU 进行平均。
5. 然后使用平均梯度来更新模型的参数。

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

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

与数据并行不同，张量并行将模型本身划分为 GPU。在数据并行中，每个 GPU 都需要适应整个模型，这在训练更大的模型时可能会成为一个限制。然而，张量并行性允许通过分解模型并将其分布在多个设备上来训练对于单个 GPU 来说可能太大的模型。

具体来说，其原理和矩阵乘法相似。按行或按列都可以对模型进行拆解。简单起见，以按列拆解为例，可以将一个大型矩阵乘法运算分解为单独的计算，每个计算都可以在不同的 GPU 上执行，如下图所示。然后将结果连接起来以获得原始结果，从而有效地分配了计算负载。

参数卸载

除了上一节中解释的 FSDP 策略之外，还可以将优化器参数卸载到 CPU，可以通过将以下代码

strategy = FSDPStrategy(    auto_wrap_policy=auto_wrap_policy,          activation_checkpointing=EncoderBlock,)

替换为

strategy = FSDPStrategy(    auto_wrap_policy=auto_wrap_policy,          activation_checkpointing=EncoderBlock,    cpu_offload=True)

内存消耗从 6.59 GB 减少到 6.03 GB：

09_fsdp-cpu-offload-with-01-2.py 的结果。

美中不足的小缺点是运行时间从 5.5 分钟增加到了 8.3 分钟。

将前面几招连着打出，就成为了最强的降龙十八掌最后一掌！

前几节对优化 ViT 进行了大量介绍，其实这些技术也同样适用于 LLM。

作者在 Lit LLaMA 和 Lit GPT 存储库中使用了许多这些技巧，这些存储库支持 LLaMA、Falcon、Pythia 和其他流行的模型。尽管如此，为了创建一个更通用的例子，作者从流行的 HF transformers 库中微调 LLM，用于对 IMDb 电影评论的情绪进行分类。

使用上述技术，仅使用 1.15 Gb 内存（bonus_DistilBERT-after.py）而不是 3.99 Gb（bonus_bigbird-before.py）就可以训练 DistilBERT 分类器。更令人印象深刻的是，通过将这些技术应用于 transformers 库中的 BigBird 模型，BigBird 仅消耗 4.03 GB（bonus_BigBird-after.py）。

strategy = FSDPStrategy(        cpu_offload=True   )
   fabric = Fabric(        accelerator="cuda",        devices=4,        strategy=strategy,         precision="bf16-true"   )
   with fabric.init_module():       model = AutoModelForSequenceClassification.from_pretrained(            "google/bigbird-roberta-base", num_labels=2)

结论

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

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

参考链接：https://lightning.ai/pages/community/tutorial/pytorch-memory-vit-llm/