PyTorch 自动混合精度库背后的谜团

原文towardsdatascience.com/the-mystery-behind-the-pytorch-automatic-mixed-precision-library-d9386e4b787e?sourcecollection_archive---------4-----------------------#2024-09-17如何通过三行代码实现 2 倍速度提升的模型训练https://mengliuz.medium.com/?sourcepost_page---byline--d9386e4b787e--------------------------------https://towardsdatascience.com/?sourcepost_page---byline--d9386e4b787e-------------------------------- 孟流赵·发布于Towards Data Science ·阅读时长 8 分钟·2024 年 9 月 17 日–你是否曾希望你的深度学习模型能够更快运行GPU 昂贵数据集庞大训练过程似乎没有尽头你有一百万个实验要运行还有一个截止日期要赶——所有这些都是你期望某种形式的训练加速的充分理由。但该选择哪个呢已经有许多关于模型训练性能调优的优秀参考资料来自于PyTorch、HuggingFace和Nvidia包括异步数据加载、缓冲区检查点、分布式数据并行和自动混合精度。在这篇文章中我将介绍自动混合精度技术。我将从简要介绍 Nvidia 的张量核心设计开始然后介绍 2018 年在 ICLR 上发表的开创性论文《混合精度训练》最后通过一个简单的示例展示如何在 FashionMNIST 上训练 ResNet50并通过加载 2 倍批量大小实现 2 倍的训练速度提升同时只需增加三行代码。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/c4fce0287d310bd6839f85ae33144b2a.png图片来源pxhere.com/en/photo/872846硬件基础 — Nvidia 张量核心首先让我们回顾一下 GPU 设计的一些基本知识。Nvidia GPU 中最受欢迎的商业产品之一是 Volta 系列例如 V100 GPU基于 GV100 GPU 设计。因此我们将基于 GV100 架构展开讨论。对于 GV100**流多处理器SM**是计算的核心设计。每个 GPU 包含 6 个 GPU 处理集群GPC和 84 个 SMV100 为 80 个 SM。整体设计如下图所示。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/dbf7020857a40f17d2576b5ced5b5123.pngVolta GV100 GPU 设计。每个 GPU 包含 6 个 GPC图形处理集群每个 GPC 包含 14 个 SM流多处理器。图像来源arxiv.org/pdf/1803.04014对于每个 SM它包含两种类型的核心CUDA 核心和 Tensor 核心。CUDA 核心是 Nvidia 在 2006 年推出的原始设计是 CUDA 平台的核心部分。CUDA 核心可以分为三种类型FP64 核心/单元、FP32 核心/单元和 Int32 核心/单元。每个 GV100 SM 包含 32 个 FP64 核心64 个 FP32 核心和 64 个 Int32 核心。Tensor 核心是在 Volta/Turing2017系列 GPU 中引入的与之前的 Pascal2016系列 GPU 区别开来。每个 GV100 SM 包含 8 个 Tensor 核心。V100 GPU 的详细规格请参见这里。下面是 SM 设计的详细图示。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/e200f32cb2b5c8c316fddbb749d81e98.png流处理器SM的子波段。每个 SM 包含四个子波段。图像来源arxiv.org/pdf/1903.03640为什么是 Tensor 核心Nvidia 的 Tensor 核心专门用于执行通用矩阵乘法GEMM和半精度矩阵乘法与累加HMMA操作。简而言之GEMM 执行 A*B C 的矩阵运算而 HMMA 将操作转换为半精度格式。详细讨论请参见这里。由于深度学习在训练过程中大量涉及 MMATensor 核心在当今的模型训练和加速中至关重要。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/79b6bd3fc6116e9b9ae38ea680529295.pngGEMM 操作示例。对于 HMMAA 和 B 通常转换为 FP16而 C 和 D 可以是 FP16 或 FP32。图像来源arxiv.org/pdf/1811.08309当然在切换到混合精度训练时始终检查您使用的 GPU 的规格。只有最新的 GPU 系列支持 Tensor 核心并且混合精度训练仅能在这些机器上使用。数据格式基础 — 单精度FP32与半精度FP16现在让我们更仔细地了解 FP32 和 FP16 格式。FP32 和 FP16 是 IEEE 格式分别使用 32 位和 16 位二进制存储表示浮动数值。这两种格式都包含三个部分a符号位b指数位c尾数位。FP32 和 FP16 的区别在于分配给指数和尾数的位数这导致了不同的值范围和精度。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/215ad72e4c2afa8517c0686351fff244.pngFP16IEEE 标准、BF16Google Brain 标准、FP32IEEE 标准和 TF32Nvidia 标准之间的差异。图片来源en.wikipedia.org/wiki/Bfloat16_floating-point_format如何将 FP16 和 FP32 转换为实际值根据 IEEE-754 标准FP32 的十进制值 (-1)^(符号位) × 2^(十进制指数 —127) ×隐含的前导 1 十进制尾数其中 127 是偏置指数值。对于 FP16公式变为(-1)^(符号位) × 2^(十进制指数 — 15) ×隐含的前导 1 十进制尾数其中 15 是相应的偏置指数值。更多关于偏置指数值的详细信息请参见此处。在这个意义上FP32 的数值范围大约是[-2¹²⁷, 2¹²⁷] ~[-1.71e38, 1.71e38]而 FP16 的数值范围大约是[-2¹⁵, 2¹⁵][-32768, 32768]。注意FP32 的十进制指数范围在 0 到 255 之间我们排除了最大值 0xFF因为它表示 NAN。因此最大的十进制指数是 254–127 127。FP16 也适用类似的规则。对于精度注意到指数和尾数都对精度限制有贡献这也叫做非标准化参见详细讨论因此 FP32 可以表示精度达到 2(-23)*2(-126)2^(-149)而 FP16 的精度可以达到 2¹⁰*2(-14)2(-24)。FP32 和 FP16 表示的差异带来了混合精度训练的关键问题因为深度学习模型的不同层/操作对数值范围和精度的敏感度不同需要分别处理。混合精度训练现在我们已经了解了 MMA 的硬件基础、Tensor 核心的概念以及 FP32 与 FP16 之间的关键区别接下来可以进一步讨论混合精度训练的细节。混合精度训练的概念最早出现在 2018 年 ICLR 论文“混合精度训练”中该论文在训练过程中将深度学习模型转换为半精度浮点数而不会损失模型精度或修改超参数。如前所述由于 FP32 和 FP16 之间的主要区别是数值范围和精度论文详细讨论了FP16 如何导致梯度消失以及如何通过损失缩放解决该问题。此外论文还提出了使用 FP32 主权重副本和在特定操作如归约和向量点积累加中使用 FP32 等技巧。损失缩放。论文中给出了使用 FP32 精度训练 Multibox SSD 检测器网络的示例如下所示。如果不进行任何缩放FP16 梯度的指数范围将≥2^(-24)低于此值的部分将变为零这相较于 FP32 来说是不够的。然而通过实验发现只需将梯度缩放 2³8 倍即可将半精度训练的准确度恢复至与 FP32 相匹配。从这个角度看作者认为[2^(-27), 2^(-24)]之间的少数梯度仍然在训练过程中很重要而低于 2^(-27)的值则不重要。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/7a45fe75b18cd7c86e9d63a193f9f30d.png在 Multibox SSD 训练示例中使用 FP32 精度时梯度值的范围。请注意值范围在[2^(-27), 2^(-24)]之间超出了 FP16 的非正规化范围仅占总梯度的几百分比但在整体训练中仍然很重要。图片来源arxiv.org/pdf/1710.03740解决这一缩放差异的方法是应用损失缩放。根据链式法则缩放损失将确保所有梯度都按相同的比例进行缩放。在最终的权重更新之前需要将梯度恢复为原始尺度。自动混合精度训练Nvidia 最早开发了自动混合精度训练作为一个名为 APEX 的 PyTorch 扩展后来被 PyTorch、TensorFlow、MXNet 等主流框架广泛采用。有关详细信息请参阅 Nvidia 的文档。为简便起见我们将仅介绍 PyTorch 的自动混合精度库pytorch.org/docs/stable/amp.html。amp 库可以自动处理大多数混合精度训练技术例如 FP32 主权重复制。用户主要接触到操作自动转换和梯度/损失缩放。操作自动转换。尽管我们提到过张量核心可以大幅提升 GEMM 操作的性能但某些操作并不适合使用半精度表示。amp 库提供了一个CUDA 操作列表列出了可用于半精度的操作。大多数矩阵乘法、卷积和线性激活都可以通过 amp.autocast 完全覆盖但是对于归约/求和、softmax 和损失计算这些计算仍然在 FP32 精度下进行因为它们对数据范围和精度更为敏感。梯度/损失缩放。amp 库提供了自动梯度缩放技术因此用户在训练过程中无需手动调整缩放。有关缩放因子的更详细算法可以在这里找到。一旦梯度被缩放必须在梯度裁剪和正则化之前进行缩放回原值。更多细节可以在这里找到。一个 FashionMNIST 训练示例torch.amp 库相对容易使用只需要三行代码就能将训练速度提升 2 倍。我们从一个非常简单的任务开始在 FashionMNIST 数据集上训练 ResNet50 模型MIT 许可证使用 FP32我们可以看到十个 Epoch 的训练时间为 333 秒https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/aa7d6413a871ff3c3b806bff458fcddd.pngResNet50 在 FashionMNIST 上的训练。图像由作者提供。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/35c410a68cae78f6b83badcb647fa49e.png小于 2**(-24) 的梯度与总梯度的比值。我们可以看到FP16 会将几乎四分之一的总梯度归零。图像由作者提供。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/aa65d416e3eb94c05cc56ef30eab8e1e.png评估结果。图像由作者提供。现在我们使用了 amp 库。amp 库只需要额外三行代码来进行混合精度训练。我们可以看到训练在 141 秒内完成比 FP32 训练速度提升了 2.36 倍同时精度、召回率和 F1-score 保持不变。scalertorch.cuda.amp.GradScaler()# start your training code# ...withtorch.autocast(device_typecuda):# training code# wrapping loss and optimizerscaler.scale(loss).backward()scaler.step(optimizer)scaler.update()https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/130adb13d32d608790dbc8b828d9ea5f.png使用 amp 的训练代码。图像由作者提供。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/d6afd6d8ec261110cf753fb45cfdbea1.png训练中的缩放因子。缩放因子仅在第一步发生变化并保持不变。图像由作者提供。https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/60a8df649cc72aa64cdad00e012f89f2.png最终结果与 FP32 训练结果可比。图像由作者提供。上面的代码的 GitHub 链接在这里。总结混合精度训练是加速深度学习模型训练的一个宝贵技术。它不仅加速了浮点运算还节省了 GPU 内存因为训练批次可以转换为 FP16从而节省了一半的 GPU 内存。通过 PyTorch 的 amp 库额外的代码可以减少为三行因为权重复制、损失缩放、操作类型转换等都由库内部处理。然而混合精度训练并没有真正解决 GPU 内存问题特别是在模型权重大小远大于数据批量时。一方面只有模型的某些层被转换为 FP16其余部分仍以 FP32 计算其次权重更新仍然需要 FP32 副本这仍然占用大量 GPU 内存第三像 Adam 这样的优化器的参数在训练过程中会占用大量 GPU 内存而混合精度训练保持优化器参数不变。从这个角度来看更先进的技术如 DeepSpeed 的 ZERO 算法是必要的。参考文献Micikevicius 等人《混合精度训练》。ICLR 2018PyTorch AMP 库:pytorch.org/tutorials/recipes/recipes/amp_recipe.htmlNvidia CUDA 浮点运算:docs.nvidia.com/cuda/floating-point/index.html