使用 PyTorch 扩展视觉模型训练平台

TL;DR：我们演示了如何使用 PyTorch 和 FairScale 的 FullyShardedDataParallel (FSDP) API 编写大型视觉 Transformer 模型。我们讨论了在 GPU 集群上扩展和优化这些模型的技术。此平台扩展工作的目标是实现大规模研究。本博客不讨论模型精度、新模型架构或新训练秘籍。

1. 引言

最新的视觉研究[1, 2] 表明模型扩展是一个很有前途的研究方向。在这个项目中，我们的目标是使我们的平台能够训练大规模视觉 Transformer (ViT) [3] 模型。我们展示了在 FAIR 视觉平台中将可训练的最大 ViT 从 1B 参数扩展到 120B 参数的工作。我们使用 PyTorch 编写了 ViT，并利用其对 GPU 集群上大规模分布式训练的支持。

在本博客的其余部分，我们将首先讨论主要挑战，即可扩展性、优化和数值稳定性。然后，我们将讨论如何使用包括数据和模型并行性、自动混合精度、核融合和bfloat16在内的技术来解决这些挑战。最后，我们将展示我们的结果并得出结论。

2. 主要挑战

2.1 可扩展性

关键的可扩展性挑战是有效地将模型的运算和状态分片到多个 GPU 上。一个 100B 参数的模型仅参数就需要约 200GB 的 RAM，假设采用 fp16 表示。因此，不可能将模型完全放入单个 GPU 中（A100 最多只有 80GB RAM）。因此，我们需要一种有效的方式将模型的数据（输入、参数、激活和优化器状态）分片到多个 GPU 上。

这个问题的另一个方面是，在不显著改变训练秘籍的情况下进行扩展。例如，某些表示学习秘籍使用高达 4096 的全局批量大小，超过这个大小我们就会开始看到精度下降。如果没有使用某种形式的张量或流水线并行性，我们无法扩展到超过 4096 个 GPU。

2.2 优化

关键的优化挑战是，即使我们扩展模型参数和浮点运算的数量，也要保持高 GPU 利用率。当我们将模型扩展到万亿次浮点运算及以上时，我们的软件堆栈开始遇到主要的瓶颈，这些瓶颈会超线性地增加训练时间并降低加速器利用率。我们只需要数百或数千个 GPU 就可以运行单个实验。加速器利用率的提高可以显著降低成本并提高设备利用率。它使我们能够资助更多的项目并并行运行更多的实验。

2.3 数值稳定性

关键的稳定性挑战是在大规模情况下避免数值不稳定和发散。我们在实验中凭经验观察到，当我们扩大模型大小、数据、批量大小、学习率等时，训练不稳定性会变得严重且难以处理。视觉 Transformer 即使在较低的参数阈值下也特别面临训练不稳定性。例如，我们发现即使是 ViT-H（只有 6.3 亿参数）在混合精度模式下不使用强数据增强也很难训练。我们需要研究模型特性和训练秘籍，以确保模型稳定训练并收敛。

3. 我们的解决方案

图 1描绘了我们针对每个挑战的解决方案。

3.1 通过数据并行性和模型并行性解决扩展挑战

我们应用了各种形式的数据并行性和模型并行性，以使非常大的模型能够适应 GPU 内存。

我们使用基于 PyTorch 的 FairScale 的 FullyShardedDataParallel (FSDP) API [4]，将参数、梯度和优化器状态分片到多个 GPU 上，从而减少每个 GPU 的内存占用。此过程包括以下三个步骤

步骤 1：我们将整个模型封装在一个 FSDP 实例中。这会在正向传播结束时对模型参数进行分片，并在正向传播开始时聚合参数。这使我们能够将参数量从 1.5B 扩展到 4.5B，提高了约 3 倍。
步骤 2：我们尝试将单个模型层封装在单独的 FSDP 实例中。这种嵌套封装通过分片和收集单个模型层的参数而不是整个模型的参数，进一步减少了内存占用。在这种模式下，峰值内存由 GPU 内存中单独封装的 Transformer 块决定，而不是整个模型。
步骤 3：我们使用激活检查点来减少激活的内存消耗。它在正向传播期间保存输入张量并丢弃中间激活张量。这些在反向传播期间重新计算。

此外，我们还尝试了模型并行技术，例如流水线并行[5]，这使我们能够在不增加批量大小的情况下扩展到更多 GPU。

3.2 通过高级 AMP 和内核融合解决优化挑战

高级 AMP

自动混合精度 (AMP) [6] 训练是指使用低于 FP32 或默认精度的比特精度训练模型，但仍能保持精度。我们尝试了以下三种级别的 AMP：

AMP O1：这指的是混合精度训练，其中权重采用 FP32 格式，部分操作采用 FP16 格式。使用 AMP O1 时，可能影响精度的操作仍保留在 FP32 中，不会自动转换为 FP16。
AMP O2：这指的是混合精度训练，但与 O1 相比，更多的权重和操作采用 FP16 格式。权重不会隐式保留在 FP32 中，而是转换为 FP16。主权重的一个副本以 FP32 精度保存，供优化器使用。如果希望归一化层权重保留在 FP32 中，则需要明确使用层封装来确保这一点。
全 FP16：这指的是全 FP16 训练，其中权重和操作都采用 FP16 格式。由于收敛问题，FP16 很难用于训练。

我们发现，带有 FP32 LayerNorm 封装的 AMP O2 在不牺牲精度的情况下实现了最佳性能。

内核融合

为了减少 GPU 内核启动开销并增加 GPU 工作粒度，我们使用 xformers 库 [7] 尝试了内核融合，包括融合 Dropout 和融合 Layer-norm。

3.3 通过研究运算数值稳定性和训练秘籍解决稳定性挑战

BFloat16 通常与 FP32 中的 LayerNorm 结合使用

bfloat16 (BF16) [8] 浮点格式提供与 FP32 相同的动态范围，同时内存占用与 FP16 相同。我们发现，我们可以使用与 FP32 相同的超参数集以 BF16 格式训练模型，而无需特殊参数调整。然而，我们发现为了使训练收敛，我们需要将 LayerNorm 保留在 FP32 模式下。

3.4 最终训练秘籍

最终训练秘籍摘要。

将外部模型封装在一个 FSDP 实例中。在正向传播后启用参数分片。
使用激活检查点、嵌套 FSDP 封装和参数扁平化来封装单个 ViT 块。
启用混合精度模式 (AMP O2)，采用 bfloat16 表示。将优化器状态保持在 FP32 精度以增强数值稳定性。
将 LayerNorm 等归一化层封装在 FP32 中以获得更好的数值稳定性。
通过将矩阵维度保持为 8 的倍数来最大化 Nvidia TensorCore 利用率。有关更多详细信息，请查看 Nvidia Tensor Core 性能指南。

4. 结果

在本节中，我们展示了 ViT 在三类任务上的扩展结果：（1）图像分类，（2）目标检测，（3）视频理解。我们的关键结果是，在应用了所讨论的扩展和优化技术后，我们能够跨这些视觉任务训练大规模 ViT 主干网络。这使得大规模视觉研究成为可能。我们训练模型直到收敛，以验证即使进行了所有优化，我们也能保持当前的基线。图 2、3、4 的共同趋势是，我们能够在 128 个 A100 GPU 上，在不到 4 小时的每个 epoch 时间内训练高达 25B 参数的模型。60B 和 120B 模型训练速度相对较慢。

图 2显示了图像分类扩展结果。它绘制了使用 128 个 A100-80GB GPU 训练 ImageNet 上的 ViT 在不同模型大小下的 epoch 时间。

图 2：图像分类扩展结果。

图 3显示了目标检测扩展结果。它绘制了使用 128 个 A100-80GB GPU 训练 COCO 上带有不同 ViT 主干的 ViTDet [9] 的 epoch 时间。

图 3：目标检测扩展结果。

图 4显示了视频理解扩展结果。它绘制了使用 128 个 V100 (32 GB) GPU 在 FP32 中训练 MViTv2 [10] 模型在 Kinetics 400 [11] 上的 epoch 时间。

图 4：视频理解扩展结果。

图 5显示了图 2 中 ViT-H 模型在 8 个 A100-40GB GPU 上的优化结果。使用了三个版本：（1）基线使用 PyTorch 的 DDP[12] 和 AMP O1，（2）FSDP + AMP-O2 + 其他优化，以及（3）FSDP + FP16 + 其他优化。这些优化总共将训练速度提高了 2.2 倍。

图 5：各种优化带来的训练加速。

5. 总结

我们展示了如何使用 PyTorch 和 FairScale 的 FullyShardedDataParallel (FSDP) API 来编写大型视觉 Transformer 模型。我们讨论了在 GPU 集群上扩展和优化这些模型的技术。我们希望本文能激励其他人使用 PyTorch 及其生态系统开发大规模机器学习模型。