近日，深度学习领域知名研究者、lightning ai 的首席人工智能教育者 sebastian raschka 在 cvpr 2023 上发表了主题演讲「scaling pytorch model training with minimal code changes」。
为了能与更多人分享研究成果，sebastian raschka 将演讲整理成一篇文章。文章探讨了如何在最小代码更改的情况下扩展 pytorch 模型训练，并表明重点是利用混合精度（mixed-precision）方法和多 gpu 训练模式，而不是低级机器优化。
文章使用视觉 transformer（vit）作为基础模型，vit 模型在一个基本数据集上从头开始，经过约 60 分钟的训练，在测试集上取得了 62% 的准确率。
github 地址：https://github.com/rasbt/cvpr2023
以下是文章原文：
构建基准在接下来的部分中，sebastian 将探讨如何在不进行大量代码重构的情况下改善训练时间和准确率。
想要注意的是，模型和数据集的详细信息并不是这里的主要关注点（它们只是为了尽可能简单，以便读者可以在自己的机器上复现，而不需要下载和安装太多的依赖）。所有在这里分享的示例都可以在 github 找到，读者可以探索和重用完整的代码。
脚本 00_pytorch-vit-random-init.py 的输出。
不要从头开始训练现如今，从头开始训练文本或图像的深度学习模型通常是低效的。我们通常会利用预训练模型，并对模型进行微调，以节省时间和计算资源，同时获得更好的建模效果。
如果考虑上面使用的相同 vit 架构，在另一个数据集（imagenet）上进行预训练，并对其进行微调，就可以在更短的时间内实现更好的预测性能：20 分钟（3 个训练 epoch）内达到 95% 的测试准确率。
00_pytorch-vit-random-init.py 和 01_pytorch-vit.py 的对比。
提升计算性能我们可以看到，相对于从零开始训练，微调可以大大提升模型性能。下面的柱状图总结了这一点。
00_pytorch-vit-random-init.py 和 01_pytorch-vit.py 的对比柱状图。
当然，模型效果可能因数据集或任务的不同而有所差异。但对于许多文本和图像任务来说，从一个在通用公共数据集上预训练的模型开始是值得的。
接下来的部分将探索各种技巧，以加快训练时间，同时又不牺牲预测准确性。
开源库 fabric在 pytorch 中以最小代码更改来高效扩展训练的一种方法是使用开源 fabric 库，它可以看作是 pytorch 的一个轻量级包装库 / 接口。通过 pip 安装。
pip install lightning
下面探索的所有技术也可以在纯 pytorch 中实现。fabric 的目标是使这一过程更加便利。
在探索「加速代码的高级技术」之前，先介绍一下将 fabric 集成到 pytorch 代码中需要进行的小改动。一旦完成这些改动，只需要改变一行代码，就可以轻松地使用高级 pytorch 功能。
pytorch 代码和修改后使用 fabric 的代码之间的区别是微小的，只涉及到一些细微的修改，如下面的代码所示：
普通 pytorch 代码（左）和使用 fabric 的 pytorch 代码
总结一下上图，就可以得到普通的 pytorch 代码转换为 pytorch+fabric 的三个步骤：
导入 fabric 并实例化一个 fabric 对象。使用 fabric 设置模型、优化器和 data loader。损失函数使用 fabric.backward ()，而不是 loss.backward ()。
这些微小的改动提供了一种利用 pytorch 高级特性的途径，而无需对现有代码进行进一步重构。
深入探讨下面的「高级特性」之前，要确保模型的训练运行时间、预测性能与之前相同。
01_pytorch-vit.py 和 03_fabric-vit.py 的比较结果。
正如前面柱状图中所看到的，训练运行时间、准确率与之前完全相同，正如预期的那样。其中，任何波动都可以归因于随机性。
在前面的部分中，我们使用 fabric 修改了 pytorch 代码。为什么要费这么大的劲呢？接下来将尝试高级技术，比如混合精度和分布式训练，只需更改一行代码，把下面的代码
fabric = fabric(accelerator=cuda)
改为
fabric = fabric(accelerator=cuda, precisinotallow=bf16-mixed)
04_fabric-vit-mixed-precision.py 脚本的比较结果。脚本地址：https://github.com/rasbt/cvpr2023/blob/main/04_fabric-vit-mixed-precision.py
通过混合精度训练，我们将训练时间从 18 分钟左右缩短到 6 分钟，同时保持相同的预测性能。这种训练时间的缩短只需在实例化 fabric 对象时添加参数「precisinotallow=bf16-mixed」即可实现。
理解混合精度机制混合精度训练实质上使用了 16 位和 32 位精度，以确保不会损失准确性。16 位表示中的计算梯度比 32 位格式快得多，并且还节省了大量内存。这种策略在内存或计算受限的情况下非常有益。
之所以称为「混合」而不是「低」精度训练，是因为不是将所有参数和操作转换为 16 位浮点数。相反，在训练过程中 32 位和 16 位操作之间切换，因此称为「混合」精度。
如下图所示，混合精度训练涉及步骤如下：
将权重转换为较低精度（fp16）以加快计算速度；计算梯度；将梯度转换回较高精度（fp32）以保持数值稳定性；使用缩放后的梯度更新原始权重。这种方法在保持神经网络准确性和稳定性的同时，实现了高效的训练。
更详细的步骤如下：
将权重转换为 fp16：在这一步中，神经网络的权重（或参数）初始时用 fp32 格式表示，将其转换为较低精度的 fp16 格式。这样可以减少内存占用，并且由于 fp16 操作所需的内存较少，可以更快地被硬件处理。计算梯度：使用较低精度的 fp16 权重进行神经网络的前向传播和反向传播。这一步计算损失函数相对于网络权重的梯度（偏导数），这些梯度用于在优化过程中更新权重。将梯度转换回 fp32：在计算得到 fp16 格式的梯度后，将其转换回较高精度的 fp32 格式。这种转换对于保持数值稳定性非常重要，避免使用较低精度算术时可能出现的梯度消失或梯度爆炸等问题。乘学习率并更新权重：以 fp32 格式表示的梯度乘以学习率将用于更新权重（标量值，用于确定优化过程中的步长）。步骤 4 中的乘积用于更新原始的 fp32 神经网络权重。学习率有助于控制优化过程的收敛性，对于实现良好的性能非常重要。
brain float 16前面谈到了「float 16-bit」精度训练。需要注意的是，在之前的代码中，指定了 precisinotallow=bf16-mixed，而不是 precisinotallow=16-mixed。这两个都是有效的选项。
在这里，bf16-mixed 中的「bf16」表示 brain floating point（bfloat16）。谷歌开发了这种格式，用于机器学习和深度学习应用，尤其是在张量处理单元（tpu）中。bfloat16 相比传统的 float16 格式扩展了动态范围，但牺牲了一定的精度。
扩展的动态范围使得 bfloat16 能够表示非常大和非常小的数字，使其更适用于深度学习应用中可能遇到的数值范围。然而，较低的精度可能会影响某些计算的准确性，或在某些情况下导致舍入误差。但在大多数深度学习应用中，这种降低的精度对建模性能的影响很小。
虽然 bfloat16 最初是为 tpu 开发的，但从 nvidia ampere 架构的 a100 tensor core gpu 开始，已经有几种 nvidia gpu 开始支持 bfloat16。
我们可以使用下面的代码检查 gpu 是否支持 bfloat16：
>>> torch.cuda.is_bf16_supported()true
如果你的 gpu 不支持 bfloat16，可以将 precisinotallow=bf16-mixed 更改为 precisinotallow=16-mixed。
多 gpu 训练和完全分片数据并行接下来要尝试修改多 gpu 训练。如果我们有多个 gpu 可供使用，这会带来好处，因为它可以让我们的模型训练速度更快。
这里介绍一种更先进的技术 — 完全分片数据并行（fully sharded data parallelism (fsdp)），它同时利用了数据并行性和张量并行性。
在 fabric 中，我们可以通过下面的方式利用 fsdp 添加设备数量和多 gpu 训练策略：
fabric = fabric(accelerator=cuda, precisinotallow=bf16-mixed,devices=4, strategy=fsdp# new!)
06_fabric-vit-mixed-fsdp.py 脚本的输出。
现在使用 4 个 gpu，我们的代码运行时间大约为 2 分钟，是之前仅使用混合精度训练时的近 3 倍。
理解数据并行和张量并行在数据并行中，小批量数据被分割，并且每个 gpu 上都有模型的副本。这个过程通过多个 gpu 的并行工作来加速模型的训练速度。
如下简要概述了数据并行的工作原理：
同一个模型被复制到所有的 gpu 上。每个 gpu 分别接收不同的输入数据子集（不同的小批量数据）。所有的 gpu 独立地对模型进行前向传播和反向传播，计算各自的局部梯度。收集并对所有 gpu 的梯度求平均值。平均梯度被用于更新模型的参数。每个 gpu 都在并行地处理不同的数据子集，通过梯度的平均化和参数的更新，整个模型的训练过程得以加速。
这种方法的主要优势是速度。由于每个 gpu 同时处理不同的小批量数据，模型可以在更短的时间内处理更多的数据。这可以显著减少训练模型所需的时间，特别是在处理大型数据集时。
然而，数据并行也有一些限制。最重要的是，每个 gpu 必须具有完整的模型和参数副本。这限制了可以训练的模型大小，因为模型必须适应单个 gpu 的内存。这对于现代的 vits 或 llms 来说这是不可行的。
与数据并行不同，张量并行将模型本身划分到多个 gpu 上。并且在数据并行中，每个 gpu 都需要适 应整个模型，这在训练较大的模型时可能成为一个限制。而张量并行允许训练那些对单个 gpu 而言可能过大的模型，通过将模型分解并分布到多个设备上进行训练。
张量并行是如何工作的呢？想象一下矩阵乘法，有两种方式可以进行分布计算 —— 按行或按列。为了简单起见，考虑按列进行分布计算。例如，我们可以将一个大型矩阵乘法操作分解为多个独立的计算，每个计算可以在不同的 gpu 上进行，如下图所示。然后将结果连接起来以获取结果，这有效地分摊了计算负载。
以上就是改动一行代码，pytorch训练三倍提速，这些「高级技术」是关键的详细内容。