使用 TensorFlow Profiler 优化 TensorFlow GPU 性能

概述

本指南将向您展示如何使用 TensorFlow Profiler 和 TensorBoard 来深入了解您的 GPU 并获得最佳性能,以及在发现一个或多个 GPU 未充分利用时进行调试。

如果您是 Profiler 新手

请记住,将计算卸载到 GPU 并不总是有利,特别是对于小型模型而言。由于以下原因,可能会出现开销:

  • 主机(CPU)和设备(GPU)之间的数据传输;以及
  • 主机启动 GPU 内核时涉及的延迟。

性能优化工作流程

本指南概述了如何从单个 GPU 开始调试性能问题,然后扩展到单个主机上的多个 GPU。

建议按以下顺序调试性能问题:

  1. 优化和调试单个 GPU 上的性能
    1. 检查输入管道是否为瓶颈。
    2. 调试单个 GPU 的性能。
    3. 启用混合精度(使用 fp16(float16))并可选地启用 XLA
  2. 优化和调试多 GPU 单主机上的性能。

例如,如果您使用 TensorFlow 分布式策略 在单个主机上的多个 GPU 上训练模型,并注意到 GPU 利用率不佳,则应首先优化和调试单个 GPU 的性能,然后再调试多 GPU 系统。

作为在 GPU 上获得高性能代码的基线,本指南假设您已经在使用 tf.function。Keras Model.compileModel.fit API 将在后台自动使用 tf.function。在使用 tf.GradientTape 编写自定义训练循环时,请参考 使用 tf.function 提高性能,了解如何启用 tf.function

接下来的部分将讨论针对上述每种情况的建议方法,以帮助您识别和解决性能瓶颈。

1. 优化单个 GPU 上的性能

在理想情况下,您的程序应该具有高 GPU 利用率,最小的 CPU(主机)到 GPU(设备)通信,以及来自输入管道的零开销。

分析性能的第一步是获取使用单个 GPU 运行的模型的配置文件。

TensorBoard 的 Profiler 概述页面(显示模型在配置文件运行期间性能的顶层视图)可以提供您的程序距离理想情况有多远的线索。

TensorFlow Profiler Overview Page

概述页面需要注意的关键数字是

  1. 实际设备执行占步骤时间的比例
  2. 放置在设备上的操作与主机上的操作的百分比
  3. 使用 fp16 的内核数量

实现最佳性能意味着在这三种情况下最大化这些数字。要深入了解您的程序,您需要熟悉 TensorBoard 的 Profiler 跟踪查看器。以下部分展示了一些常见的跟踪查看器模式,在诊断性能瓶颈时应该注意这些模式。

以下是一张在单个 GPU 上运行的模型跟踪视图的图像。从TensorFlow 名称范围TensorFlow 操作部分,您可以识别模型的不同部分,例如前向传递、损失函数、反向传递/梯度计算和优化器权重更新。您还可以看到在旁边的 GPU 上运行的操作,这些操作指的是 CUDA 流。每个流用于特定任务。在此跟踪中,Stream#118 用于启动计算内核和设备到设备的复制。Stream#119 用于主机到设备的复制,Stream#120 用于设备到主机的复制。

下面的跟踪显示了高性能模型的常见特征。

image

例如,GPU 计算时间线(Stream#118)看起来“繁忙”,几乎没有间隙。主机到设备(Stream #119)和设备到主机(Stream #120)的复制最少,步骤之间的间隙也最少。当您为您的程序运行 Profiler 时,您可能无法在跟踪视图中识别出这些理想特征。本指南的其余部分涵盖了常见场景以及如何修复它们。

1. 调试输入管道

GPU 性能调试的第一步是确定您的程序是否受输入限制。最简单的方法是使用 Profiler 的 输入管道分析器(在 TensorBoard 上),它提供了输入管道中花费时间的概述。

image

如果您的输入管道对步骤时间有重大贡献,您可以采取以下潜在措施

  • 您可以使用 tf.data 特定的 指南 来学习如何调试您的输入管道。
  • 另一种快速检查输入管道是否为瓶颈的方法是使用不需要任何预处理的随机生成输入数据。 这是一个示例,说明了如何将此技术用于 ResNet 模型。如果输入管道是最佳的,您应该在使用真实数据和生成随机/合成数据时体验到类似的性能。合成数据情况下的唯一开销将是由于输入数据复制造成的,而输入数据复制也可以预取和优化。

此外,请参考 优化输入数据管道的最佳实践

2. 调试单个 GPU 的性能

有几个因素会导致 GPU 利用率低。以下是查看 跟踪查看器 时通常观察到的几种场景以及潜在的解决方案。

1. 分析步骤之间的间隙

当您的程序运行不佳时,一个常见的观察结果是训练步骤之间的间隙。在下面的跟踪视图图像中,步骤 8 和 9 之间存在一个大的间隙,这意味着 GPU 在这段时间内处于空闲状态。

image

如果您的跟踪查看器显示步骤之间存在大的间隙,这可能表明您的程序受输入限制。在这种情况下,如果您尚未完成,则应参考上一节关于调试输入管道的部分。

但是,即使输入管道经过优化,您仍然可能在一步结束和另一步开始之间存在间隙,这是由于 CPU 线程争用造成的。 tf.data 利用后台线程来并行化管道处理。这些线程可能会干扰每一步开始时发生的 GPU 主机端活动,例如复制数据或调度 GPU 操作。

如果您在主机端(它在 GPU 上调度这些操作)注意到大的间隙,您可以设置环境变量 TF_GPU_THREAD_MODE=gpu_private。这将确保 GPU 内核从它们自己的专用线程启动,并且不会排队到 tf.data 工作后面。

步骤之间的间隙也可能是由指标计算、Keras 回调或主机上运行的操作(在 tf.function 之外)造成的。这些操作的性能不如 TensorFlow 图中的操作。此外,其中一些操作在 CPU 上运行,并将张量从 GPU 来回复制。

如果在优化输入管道后,您仍然在跟踪查看器中注意到步骤之间的间隙,您应该查看步骤之间的模型代码,并检查禁用回调/指标是否会提高性能。跟踪查看器中也有一些关于这些操作的详细信息(设备端和主机端)。在这种情况下,建议是通过在固定数量的步骤后执行这些操作(而不是每一步都执行)来分摊这些操作的开销。在使用 Model.compile 方法(在 tf.keras API 中)时,设置 steps_per_execution 标志会自动执行此操作。对于自定义训练循环,请使用 tf.while_loop

2. 实现更高的设备利用率

1. 小型 GPU 内核和主机内核启动延迟

主机将内核排队到 GPU 上运行,但在内核实际在 GPU 上执行之前存在延迟(大约 20-40 μs)。在理想情况下,主机在 GPU 上排队足够的内核,以便 GPU 将大部分时间用于执行,而不是等待主机排队更多内核。

Profiler 的 概述页面(在 TensorBoard 上)显示了 GPU 由于等待主机启动内核而空闲的时间。在下图中,GPU 由于等待内核启动而空闲了大约 10% 的步骤时间。

image

同一程序的 跟踪查看器 显示了内核之间的小间隙,主机在这些间隙中忙于在 GPU 上启动内核。

image

通过在 GPU 上启动大量小型操作(例如标量加法),主机可能无法跟上 GPU。TensorBoard 上的同一配置文件的 TensorFlow 统计信息 工具显示 126,224 个 Mul 操作花费了 2.77 秒。因此,每个内核大约为 21.9 μs,这非常小(与启动延迟大致相同),并且可能导致主机内核启动延迟。

image

如果您的 跟踪查看器 显示 GPU 上的操作之间存在许多小间隙(如上图所示),您可以

  • 连接小型张量并使用矢量化操作,或者使用更大的批次大小,使每个启动的内核执行更多工作,这将使 GPU 繁忙更长时间。
  • 确保您使用 tf.function 创建 TensorFlow 图,这样您就不会以纯急切模式运行操作。如果您使用的是 Model.fit(与使用 tf.GradientTape 的自定义训练循环相反),那么 tf.keras.Model.compile 会自动为您执行此操作。
  • 使用 tf.function(jit_compile=True) 或自动聚类融合内核,使用 XLA。有关更多详细信息,请转到下面的 启用混合精度和 XLA 部分,了解如何启用 XLA 以获得更高的性能。此功能可以导致高设备利用率。
2. TensorFlow 操作放置

Profiler 的 概述页面 显示了放置在主机上的操作与设备上的操作的百分比(您还可以通过查看 跟踪查看器 来验证特定操作的放置。如以下图像所示,您希望主机上的操作百分比与设备上的操作百分比相比非常小。

image

理想情况下,大多数计算密集型操作应该放置在 GPU 上。

要找出模型中的操作和张量分配到的设备,请将 tf.debugging.set_log_device_placement(True) 设置为程序的第一条语句。

请注意,在某些情况下,即使您指定将操作放置在特定设备上,其实现也可能会覆盖此条件(例如:tf.unique)。即使对于单个 GPU 训练,指定分布策略(例如 tf.distribute.OneDeviceStrategy)也会导致在您的设备上更确定地放置操作。

将大多数操作放置在 GPU 上的一个原因是防止主机和设备之间过度内存复制(主机和设备之间模型输入/输出数据的内存复制是预期的)。下面的跟踪视图中展示了 GPU 流#167#168#169 上的过度复制示例。

image

如果这些复制阻塞了 GPU 内核的执行,它们有时会损害性能。 跟踪查看器 中的内存复制操作包含有关复制张量的源操作的更多信息,但可能并不总是容易将 memCopy 与操作相关联。在这些情况下,查看附近的操作以检查内存复制是否在每一步的相同位置发生会很有帮助。

3. GPU 上更高效的内核

一旦您的程序的 GPU 利用率可以接受,下一步就是通过利用 Tensor Core 或融合操作来提高 GPU 内核的效率。

1. 利用 Tensor Core

现代 NVIDIA® GPU 具有专门的 Tensor Core,可以显着提高合格内核的性能。

您可以使用 TensorBoard 的 GPU 内核统计信息 来可视化哪些 GPU 内核是 Tensor Core 适用的,以及哪些内核正在使用 Tensor Core。启用 fp16(请参阅下面的启用混合精度部分)是使您的程序的通用矩阵乘法 (GEMM) 内核(matmul 操作)利用 Tensor Core 的一种方法。当精度为 fp16 且输入/输出张量维度可被 8 或 16 整除(对于 int8)时,GPU 内核可以有效地利用 Tensor Core。

有关如何使内核对 GPU 效率更高的其他详细建议,请参考 NVIDIA® 深度学习性能 指南。

2. 融合操作

使用 tf.function(jit_compile=True) 将较小的操作融合成更大的内核,从而导致性能显着提升。要了解更多信息,请参考 XLA 指南。

3. 启用混合精度和 XLA

在完成上述步骤后,启用混合精度和 XLA 是两个可选步骤,可以进一步提高性能。建议的方法是逐一启用它们,并验证性能提升是否符合预期。

1. 启用混合精度

TensorFlow 混合精度 指南展示了如何在 GPU 上启用 fp16 精度。在 NVIDIA® GPU 上启用 AMP 以使用 Tensor Core,与仅在 Volta 和更新的 GPU 架构上使用 fp32 (float32) 精度相比,可实现高达 3 倍的整体加速。

确保矩阵/张量维度满足调用使用 Tensor Core 的内核的要求。当精度为 fp16 且输入/输出维度可被 8 或 16 整除(对于 int8)时,GPU 内核可以有效地使用 Tensor Core。

请注意,使用 cuDNN v7.6.3 及更高版本时,卷积维度将在必要时自动填充以利用 Tensor Core。

遵循以下最佳实践以最大限度地提高 fp16 精度的性能优势。

1. 使用最佳 fp16 内核

启用 fp16 后,程序的矩阵乘法 (GEMM) 内核应使用相应的 fp16 版本,该版本利用了 Tensor Core。但是,在某些情况下,这不会发生,并且您不会体验到启用 fp16 预期的加速,因为您的程序回退到低效的实现。

image

GPU 内核 统计页面显示哪些操作是 Tensor Core 适用的,以及哪些内核实际上正在使用高效的 Tensor Core。NVIDIA® 深度学习性能指南 包含有关如何利用 Tensor Core 的更多建议。此外,使用 fp16 的好处也会在以前受内存限制的内核中显示,因为现在操作将花费一半的时间。

2. 动态与静态损失缩放

使用 fp16 时,损失缩放是必要的,以防止由于低精度导致的下溢。损失缩放有两种类型,动态和静态,两者在 混合精度指南 中有更详细的解释。您可以使用 mixed_float16 策略在 Keras 优化器中自动启用损失缩放。

在尝试优化性能时,重要的是要记住,动态损失缩放会引入额外的条件操作,这些操作在主机上运行,并导致在跟踪查看器中步骤之间可见的间隙。另一方面,静态损失缩放没有这样的开销,并且在性能方面可能是一个更好的选择,但您需要指定正确的静态损失缩放值。

2. 使用 tf.function(jit_compile=True) 或自动聚类启用 XLA

作为使用单个 GPU 获得最佳性能的最后一步,您可以尝试启用 XLA,它将融合操作并导致更好的设备利用率和更低的内存占用。有关如何在程序中使用 tf.function(jit_compile=True) 或自动聚类启用 XLA 的详细信息,请参阅 XLA 指南。

您可以将全局 JIT 级别设置为 -1 (关闭)、12。较高的级别更激进,可能会降低并行性并使用更多内存。如果您有内存限制,请将值设置为 1。请注意,对于具有可变输入张量形状的模型,XLA 性能不佳,因为 XLA 编译器必须在遇到新形状时不断编译内核。

2. 优化多 GPU 单主机上的性能

tf.distribute.MirroredStrategy API 可用于将模型训练从一个 GPU 扩展到单个主机上的多个 GPU。(要了解有关如何使用 TensorFlow 进行分布式训练的更多信息,请参阅 使用 TensorFlow 进行分布式训练使用 GPU使用 TPU 指南以及 使用 Keras 进行分布式训练 教程。)

虽然从一个 GPU 到多个 GPU 的转换理想情况下应该开箱即用地可扩展,但有时您可能会遇到性能问题。

从使用单个 GPU 训练到在同一主机上的多个 GPU 训练时,理想情况下您应该体验到性能扩展,只有梯度通信和主机线程利用率增加的额外开销。由于这种开销,例如,如果您从 1 个 GPU 迁移到 2 个 GPU,您将不会获得确切的 2 倍加速。

下面的跟踪视图显示了在多个 GPU 上训练时额外通信开销的示例。在连接梯度、跨副本进行通信以及在进行权重更新之前进行拆分时,会有一些开销。

image

以下清单将帮助您在优化多 GPU 场景中的性能时获得更好的性能

  1. 尝试最大限度地提高批次大小,这将导致更高的设备利用率并将通信成本分摊到多个 GPU 上。使用 内存分析器 有助于了解程序离峰值内存利用率有多近。请注意,虽然较高的批次大小会影响收敛,但这通常会被性能优势所抵消。
  2. 从单个 GPU 迁移到多个 GPU 时,同一主机现在必须处理更多输入数据。因此,在 (1) 之后,建议重新检查输入管道性能,并确保它不是瓶颈。
  3. 在程序的跟踪视图中检查 GPU 时间线以查找任何不必要的 AllReduce 调用,因为这会导致跨所有设备的同步。在上面显示的跟踪视图中,AllReduce 是通过 NCCL 内核完成的,并且每个 GPU 上只有一个 NCCL 调用用于每个步骤上的梯度。
  4. 检查不必要的 D2H、H2D 和 D2D 复制操作,这些操作可以最小化。
  5. 检查步长时间以确保每个副本都在执行相同的工作。例如,可能会发生一个 GPU(通常是 GPU0)被过度订阅,因为主机错误地最终将更多工作放在它上面。
  6. 最后,在跟踪视图中检查所有 GPU 上的训练步骤以查找任何按顺序执行的操作。这通常发生在程序包含来自一个 GPU 到另一个 GPU 的控制依赖项时。过去,在这种情况下调试性能是逐案解决的。如果您在程序中观察到这种行为,请 提交一个 GitHub 问题,其中包含跟踪视图的图像。

1. 优化梯度 AllReduce

使用同步策略进行训练时,每个设备都会收到一部分输入数据。

在计算模型的前向和反向传递后,需要聚合和减少每个设备上计算的梯度。此梯度 AllReduce 发生在每个设备上的梯度计算之后,以及优化器更新模型权重之前。

每个 GPU 首先连接跨模型层的梯度,使用 tf.distribute.CrossDeviceOps (tf.distribute.NcclAllReduce 是默认值) 在 GPU 之间进行通信,然后在每个层进行减少后返回梯度。

优化器将使用这些减少的梯度来更新模型的权重。理想情况下,此过程应该在所有 GPU 上同时发生,以防止任何开销。

AllReduce 的时间应该与以下时间大致相同

(number of parameters * 4bytes)/ (communication bandwidth)

此计算对于快速检查以了解运行分布式训练作业时的性能是否符合预期,或者您是否需要进行进一步的性能调试非常有用。您可以从 Model.summary 获取模型中的参数数量。

请注意,每个模型参数的大小为 4 字节,因为 TensorFlow 使用 fp32 (float32) 来通信梯度。即使您启用了 fp16,NCCL AllReduce 也会使用 fp32 参数。

为了获得扩展的好处,步长时间需要比这些开销高得多。实现此目标的一种方法是使用更高的批次大小,因为批次大小会影响步长时间,但不会影响通信开销。

2. GPU 主机线程争用

在运行多个 GPU 时,CPU 的工作是通过有效地在设备之间启动 GPU 内核来保持所有设备繁忙。

但是,当 CPU 可以在一个 GPU 上调度大量独立操作时,CPU 可以决定使用大量主机线程来保持一个 GPU 繁忙,然后以非确定性顺序在另一个 GPU 上启动内核。这会导致倾斜或负面扩展,从而会对性能产生负面影响。

下面的 跟踪查看器 显示了当 CPU 无效地交错 GPU 内核启动时的开销,因为 GPU1 处于空闲状态,然后在 GPU2 开始运行操作后开始运行操作。

image

主机的跟踪视图显示主机在 GPU2 上启动内核,然后在 GPU1 上启动内核(请注意,下面的 tf_Compute* 操作不代表 CPU 线程)。

image

如果您在程序的跟踪视图中遇到这种 GPU 内核交错,建议的操作是

  • 将 TensorFlow 环境变量 TF_GPU_THREAD_MODE 设置为 gpu_private。此环境变量将告诉主机将 GPU 的线程保持为私有。
  • 默认情况下,TF_GPU_THREAD_MODE=gpu_private 将线程数设置为 2,这在大多数情况下已经足够了。但是,可以通过将 TensorFlow 环境变量 TF_GPU_THREAD_COUNT 设置为所需的线程数来更改该数字。