原文地址：Accelerating PyTorch Model Training, by Sebastian Raschka, on 2023-06-3
借助混合精度与全分片数据并行加速PyTorch模型训练
上周，我在温哥华举办的2023年计算机视觉与模式识别会议（CVPR 2023）上做了题为《以最少代码改动扩展PyTorch模型训练规模》的演讲。在活动现场能和这么多人交流，感觉非常愉快。
为方便日后查阅，也为了没能到场的朋友，我想做个小尝试，把这次演讲整理成一篇短文，内容就在下方。（注：下周我会再分享更多CVPR的参会心得！）
因此，本文将深入探讨如何以最少的代码改动实现PyTorch模型训练的规模扩展。本文的重点是利用混合精度技术与多GPU训练范式，而非底层硬件优化。我们将以一个用于图像分类的简易视觉Transformer（ViT）作为基础模型展开讲解。
从基线模型开始
我们的ViT模型在基础数据集上从头开始训练，耗时约60分钟，测试集准确率达到62%。
在接下来的章节中，我们将探索无需大幅重构代码，就能提升训练速度与准确率的方法。
请注意，模型与数据集的细节并非本文重点（选择它们只是为了尽可能简化，方便你在自己的硬件上复现代码，无需下载安装过多依赖项）。本文分享的所有示例都可在 https://github.com/rasbt/cvpr2023 获取，你可以查看并复用完整代码。
https://github.com/rasbt/cvpr2023/blob/main/00_pytorch-vit-random-init.py 脚本的运行输出。
告别从头训练
在当今时代，针对文本或视觉任务从头开始训练深度学习模型往往性价比不高。我们通常会利用预训练模型，针对自身任务进行微调，以此节省时间与计算资源，同时获得更优的模型性能。
如果沿用上述相同的ViT架构，采用在其他数据集（ImageNet）上预训练的权重进行微调，我们就能在更短的时间内获得大幅提升的预测效果：仅需20分钟（3个训练轮次），测试准确率就能达到95%。
https://github.com/rasbt/cvpr2023/blob/main/00_pytorch-vit-random-init.py 与 https://github.com/rasbt/cvpr2023/blob/main/01_pytorch-vit.py 的效果对比
提升计算性能
上文我们已经看到，相比从头训练，微调能带来巨大的性能提升，下图用更简洁的柱状图做了总结。
https://github.com/rasbt/cvpr2023/blob/main/00_pytorch-vit-random-init.py 与 https://github.com/rasbt/cvpr2023/blob/main/01_pytorch-vit.py 的效果对比
当然，实际效果会因数据集或任务的不同而有所差异，但对于绝大多数文本与视觉任务而言，基于通用公开数据集预训练的模型作为起点都是值得的。
接下来的章节将探索多种在不损失预测准确率的前提下，进一步加快训练速度的技巧。
开源Fabric库介绍
想要以最少的代码改动高效扩展PyTorch训练规模，方法之一是使用 https://lightning.ai/docs/fabric/stable/ 。你可以把它看作是PyTorch的轻量封装库/应用程序接口。我们可以通过以下命令安装：
pip install lightning
当然，下文要介绍的所有技术也都可以用原生PyTorch实现——Fabric的目标只是让这些操作更便捷。
在探索这些“加速代码的高级技巧”之前，我们先看看将Fabric集成到PyTorch代码中需要做的小幅改动。（完成下文这些改动后，只需修改一行代码就能启用PyTorch的高级特性。）
原生PyTorch代码与修改后适配Fabric的代码差异很小，仅需少量修改，如下方代码高亮所示：
左侧：原生PyTorch代码；右侧：接入Fabric的PyTorch代码
总结上图，将原生PyTorch代码转换为PyTorch+Fabric版本主要分为3步：
- 导入Fabric并实例化一个Fabric对象。
- 使用Fabric配置模型、优化器与数据加载器。
- 对损失值调用
fabric.backward()，替代常规的loss.backward()。
这些微小的改动，为我们打开了使用PyTorch高级特性的大门，且无需再重构其余现有代码。
但在深入介绍这些“高级特性”之前，我们先确认模型的训练耗时与预测性能和之前保持一致。
https://github.com/rasbt/cvpr2023/blob/main/01_pytorch-vit.py 与 https://github.com/rasbt/cvpr2023/blob/main/03_fabric-vit.py 的效果对比
从上方柱状图可以看到，训练耗时与准确率和之前完全一致，符合预期。（所有波动都可归因于随机性。）
在上一节中，我们用Fabric修改了PyTorch代码。为什么要费这些功夫呢？正如下文所示，现在我们只需改动一行代码，就能尝试混合精度、分布式训练等高级技术，只需把
fabric = Fabric(accelerator="cuda")
修改为
fabric = Fabric(accelerator="cuda", precision="bf16-mixed")
https://github.com/rasbt/cvpr2023/blob/main/04_fabric-vit-mixed-precision.py 脚本的运行结果
借助混合精度训练，我们将训练时间从约18分钟缩短到了6分钟，同时保持了预测性能不变。仅需在实例化Fabric对象时添加 precision="bf16-mixed" 参数，就实现了训练时间的大幅缩减。
理解混合精度的原理
混合精度训练本质上是同时使用16位与32位精度，确保准确率不会下降。用16位格式计算梯度比32位格式快得多，还能节省大量显存。这种策略尤其在显存或计算资源受限的场景下优势明显。
它之所以被称为“混合”精度训练，而非“低”精度训练，是因为我们不会把所有参数与运算都转为16位浮点数。相反，训练过程中我们会在32位与16位运算之间切换，因此得名“混合”精度。
如下图所示，混合精度训练的流程包括：将权重转为低精度（FP16）以加快计算、计算梯度、将梯度转回高精度（FP32）以保证数值稳定性，再用缩放后的梯度更新原始权重。
这种方式既能实现高效训练，又能维持神经网络的准确率与稳定性。
具体步骤如下：
- 将权重转为FP16：在这一步中，神经网络原本为FP32格式的权重（参数）会被转换为精度更低的FP16格式。这能降低显存占用，同时加快计算速度——因为FP16运算所需显存更少，硬件处理速度更快。
- 计算梯度：神经网络的前向传播与反向传播都使用低精度的FP16权重执行。这一步会计算损失函数相对于网络权重的梯度（偏导数），用于优化过程中更新权重。
- 将梯度转为FP32：用FP16计算完梯度后，会将其转换回精度更高的FP32格式。这一转换对维持数值稳定性至关重要，能避免低精度运算可能出现的梯度消失或梯度爆炸等问题。
- 乘以学习率并更新权重：转为FP32格式的梯度会乘以学习率（一个标量值，决定优化过程中的步长）。
第4步得到的结果会用于更新原始的FP32神经网络权重。学习率用于控制优化过程的收敛速度，对取得良好性能至关重要。
脑浮点16（bfloat16）
上文我们提到了“16位浮点数”精度训练。但请注意，之前代码中我们指定的是 precision="bf16-mixed"，而非 precision="16-mixed"——两者都是合法选项。
这里，“bf16-mixed”中的“bf16”指的是 https://cloud.google.com/tpu/docs/bfloat16<（bfloat16，即脑浮点16）。该格式由谷歌专为机器学习与深度学习应用开发，尤其适配其张量处理单元（TPU）。与传统的float16格式相比，bfloat16以降低精度为代价，扩展了动态范围。/p>
扩展的动态范围让bfloat16可以表示极大与极小的数值，更适合深度学习这类会遇到大范围数值的应用场景。不过，精度降低可能会影响部分计算的准确率，或在某些情况下导致舍入误差。但在绝大多数深度学习应用中，这种精度降低对模型性能的影响微乎其微。
虽然bfloat16最初是为TPU开发的，但如今多款NVIDIA GPU也已支持该格式，首批支持的是NVIDIA Ampere架构中的A100张量核心GPU。
你可以通过以下代码检查你的GPU是否支持bfloat16：
>>> torch.cuda.is_bf16_supported()
True
如果你的GPU不支持脑浮点16，可以将 precision="bf16-mixed" 改为 precision="16-mixed"。
多GPU训练与全分片数据并行
我们接下来要尝试的改进是多GPU训练。如果我们有多块GPU可用，多GPU训练就能发挥作用，让模型训练速度更快。
利用这一思路的一项更高级技术是全分片数据并行（Fully Sharded Data Parallelism，FSDP），它同时结合了数据并行与张量并行的优势。
在Fabric中，我们可以通过添加设备数量与多GPU训练策略来启用FSDP，如下所示：
fabric = Fabric(
accelerator="cuda",
precision="bf16-mixed",
devices=4,
strategy="FSDP" # 新增！
)
https://github.com/rasbt/cvpr2023/blob/main/06_fabric-vit-mixed-fsdp.py 脚本的运行结果
因为我们现在用上了4块GPU，结果其实也在意料之中：代码现在仅需约2分钟就能运行完成，比仅使用混合精度训练时快了约3倍。
理解数据并行与张量并行
在数据并行中，小批次数据会被拆分，每块GPU上都有一份完整的模型副本。多块GPU并行工作，从而加快模型训练速度。
简而言之，其工作原理如下：
- 相同的模型会被复制到所有GPU上。
- 每块GPU会接收输入数据的不同子集（不同的小批次数据）。
- 所有GPU独立执行模型的前向与反向传播，计算各自的本地梯度。
- 随后，所有GPU的梯度会被收集并取平均值。
- 平均后的梯度会用于更新模型的参数。
这种方法的核心优势是速度。由于每块GPU都和其他GPU同时处理一批独立的小批次数据，模型能在更短时间内完成更多数据的训练。尤其在处理大型数据集时，能大幅缩短训练时长。
不过，数据并行也存在局限性。最关键的一点是，每块GPU都必须存储完整的模型及其参数。这就限制了可训练模型的大小——模型必须能放进单块GPU的显存中，而这对于现代的ViT或大语言模型来说往往无法实现。
与数据并行将小批次数据拆分到多台设备不同，张量并行是将模型本身拆分到多块GPU上。数据并行要求每块GPU都能装下整个模型，在训练更大模型时会成为瓶颈；而张量并行则通过拆分模型并分发到多台设备，让单块GPU装不下的大模型也能完成训练。
它是如何工作的？以矩阵乘法为例，有两种分发方式——按行拆分或按列拆分。为简单起见，我们以按列拆分为例。比如，我们可以把一个大型矩阵乘法运算拆分成多个独立计算，每个计算都可以在不同的GPU上执行，如下图所示。最后将结果拼接起来，得到原始运算的结果，从而实现计算负载的分布式处理。
总结
在本文中，我们看到了只需修改几行代码，就能运用混合精度训练、分布式多GPU训练等技术对训练进行优化。在视觉Transformer的示例中，使用混合精度训练后，运行时间从约18分钟缩短到约6分钟，减少了三分之二。再增加3块GPU后，训练时间从约6分钟进一步缩短到约2分钟。在不损失预测准确率的前提下，整体提速接近10倍。
当然，实际效果会因模型、数据集和任务的不同而有所差异。不过，本文的核心目的是介绍如何借助 https://pytorch.org/ 与 https://lightning.ai/docs/fabric/stable/ 便捷地实验这些技术。
（参考：所有代码示例均可在 https://github.com/rasbt/cvpr2023 获取）
本专栏是我的个人兴趣项目，不提供直接报酬。不过，如果您愿意支持我，可以考虑购买 https://sebastianraschka.com/books 上的书籍。如果您觉得这些书有启发、有帮助，也欢迎推荐给您的朋友和同事。
https://www.amazon.com/Machine-Learning-PyTorch-Scikit-Learn-scikit-learn-ebook-dp-B09NW48MR1/dp/B09NW48MR1/ 、https://nostarch.com/machine-learning-and-ai-beyond-basics 与 http://mng.bz/M96o
您的支持对我意义重大！非常感谢！