计算性能分析

分析计算#

使用 Perfetto 查看程序追踪#

我们可以使用 JAX 分析器来生成 JAX 程序的追踪信息，并使用 Perfetto 可视化工具进行可视化。目前，此方法会阻塞程序，直到点击链接并加载 Perfetto UI 追踪信息为止。

如果您希望在无需任何交互的情况下获取分析信息，请查看下方的 XProf 分析器。XProf 提供了更详细的追踪查看器，并提供如内存分析工具 (Memory Profile Tool) 和图查看器 (Graph Viewer) 等工具来分析内存使用情况和 XLA 计算图结构。

with jax.profiler.trace("/tmp/jax-trace", create_perfetto_link=True):
  # Run the operations to be profiled
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

计算完成后，程序会提示您打开一个指向 ui.perfetto.dev 的链接。打开链接时，Perfetto UI 将加载追踪文件并打开可视化界面。

Perfetto trace viewer

程序执行将在加载链接后继续。该链接在首次打开后即失效，但它会重定向到一个保持有效的新 URL。之后，您可以点击 Perfetto UI 中的“Share”按钮来创建可分享给其他人的追踪永久链接。

远程分析#

当分析远程运行的代码（例如在托管的虚拟机上）时，需要建立端口 9001 的 SSH 隧道才能使链接生效。您可以使用此命令实现：

$ ssh -L 9001:127.0.0.1:9001 <user>@<host>

或者如果您使用的是 Google Cloud：

$ gcloud compute ssh <machine-name> -- -L 9001:127.0.0.1:9001

手动捕获#

除了使用 jax.profiler.trace 以程序化方式捕获追踪外，您还可以通过在相关脚本中调用 jax.profiler.start_server(<port>) 来启动分析服务器。如果您只需要在脚本的特定部分激活分析服务器，可以通过调用 jax.profiler.stop_server() 将其关闭。

脚本运行且分析服务器启动后，我们可以通过运行以下命令手动捕获并进行追踪：

$ python -m jax.collect_profile <port> <duration_in_ms>

默认情况下，生成的追踪信息会转储到临时目录，但这可以通过传入 --log_dir=<选择的目录> 来覆盖。此外，默认情况下，程序会提示您打开指向 ui.perfetto.dev 的链接。打开链接时，Perfetto UI 将加载追踪文件并打开可视化界面。通过在命令中传入 --no_perfetto_link 可禁用此功能。或者，您也可以将 Tensorboard 指向 log_dir 来分析追踪（请参阅下方的“XProf (Tensorboard Profiling)”部分）。

XProf (TensorBoard 分析)#

XProf 可用于分析 JAX 程序。XProf 是获取和可视化程序性能追踪及分析报告（包括 GPU 和 TPU 上的活动）的绝佳方式。最终结果大致如下：

XProf example

安装#

XProf 既可以作为 TensorBoard 的插件使用，也可以作为独立运行的程序使用。

pip install xprof

如果您安装了 TensorBoard，xprof pip 包也会安装 TensorBoard Profiler 插件。请注意仅安装一个版本的 TensorFlow 或 TensorBoard，否则您可能会遇到下文所述的“重复插件”错误。有关安装 TensorBoard 的更多信息，请参阅 https://tensorflowcn.cn/guide/profiler。

使用 nightly 版本 TensorBoard 进行分析需要 nightly 版本的 XProf。

pip install tb-nightly xprof-nightly

XProf 和 Tensorboard#

XProf 是支持 Tensorboard 分析和追踪捕获功能的基础工具。只要安装了 xprof，Tensorboard 中就会出现“Profile”选项卡。只要启动时指向相同的日志目录，使用此功能与独立运行 XProf 的效果相同。这包括分析捕获、分析报告和查看功能。XProf 取代了之前推荐的 tensorboard_plugin_profile 功能。

$ tensorboard --logdir=/tmp/profile-data
[...]
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.19.0 at https://:6006/ (Press CTRL+C to quit)

在 Google Cloud 上，我们建议使用 cloud-diagnostics-xprof，以便轻松设置和托管 TensorBoard 和 XProf，并存储运行后的分析结果。

程序化捕获#

您可以通过 jax.profiler.start_trace() 和 jax.profiler.stop_trace() 方法对代码进行插桩以捕获分析追踪。调用带有写入追踪文件目录路径的 start_trace()。这应该是启动 XProf 时使用的相同 --logdir 目录。然后，您可以使用 XProf 查看追踪。

例如，获取分析追踪：

import jax

jax.profiler.start_trace("/tmp/profile-data")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

注意 block_until_ready() 调用。我们使用它来确保设备上的执行被追踪捕获。有关为何需要此操作的详细信息，请参阅异步调度 (Asynchronous dispatch)。

您还可以使用 jax.profiler.trace() 上下文管理器来替代 start_trace 和 stop_trace。

import jax

with jax.profiler.trace("/tmp/profile-data"):
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

查看追踪#

捕获追踪后，您可以使用 XProf UI 查看它。

您可以通过指向您的日志目录，使用独立的 XProf 命令直接启动分析器 UI：

$ xprof --port 8791 /tmp/profile-data
Attempting to start XProf server:
  Log Directory: /tmp/profile-data
  Port: 8791
XProf at https://:8791/ (Press CTRL+C to quit)

在浏览器中导航至提供的 URL（例如 https://:8791/）以查看分析结果。

可用的追踪会显示在左侧的“Runs”下拉菜单中。选择您感兴趣的运行，然后在“Tools”下拉菜单下，选择 trace_viewer。现在您应该能看到执行的时间线。您可以使用 WASD 键导航追踪，并点击或拖动以选择事件获取详细信息。有关使用追踪查看器的更多详细信息，请参阅这些 TensorFlow 文档。

通过 XProf 进行手动捕获#

以下是从正在运行的程序中手动触发 N 秒追踪的说明。

启动 XProf 服务器：
```
xprof --logdir /tmp/profile-data/
```
您应该能够通过 https://:8791/ 加载 XProf。您可以使用 --port 标志指定不同的端口。如果 JAX 在远程服务器上运行，请参阅下方的在远程机器上进行分析。
在您想要分析的 Python 程序或进程中，在开头附近添加以下代码：
```
import jax.profiler
jax.profiler.start_server(9999)
```
这将启动 XProf 连接到的分析服务器。在执行下一步之前，分析服务器必须处于运行状态。完成使用服务器后，您可以调用 jax.profiler.stop_server() 将其关闭。

如果您想分析长时间运行程序（例如长训练循环）的片段，可以在程序开头放入此代码并正常启动程序。如果您想分析短程序（例如微基准测试），一种方法是在 IPython shell 中启动分析服务器，并在下一步启动捕获后使用 %run 运行短程序。另一种方法是在程序开头启动分析服务器，并使用 time.sleep() 为您留出足够的启动捕获时间。
打开 https://:8791/，点击左上角的“CAPTURE PROFILE”按钮。输入“localhost:9999”作为分析服务 URL（这是您在上一步中启动的分析服务器地址）。输入您希望分析的毫秒数，然后点击“CAPTURE”。
如果想要分析的代码尚未运行（例如，如果您在 Python shell 中启动了分析服务器），请在捕获运行期间运行它。
捕获完成后，XProf 应该会自动刷新。（并非所有的 XProf 分析功能都已与 JAX 对接，因此起初看起来可能什么都没捕获到。）在左侧“Tools”下方，选择 trace_viewer。

现在您应该能看到执行的时间线。您可以使用 WASD 键导航追踪，点击或拖动以选择事件并在下方查看详细信息。有关使用追踪查看器的更多详细信息，请参阅这些 XProf 文档。

您还可以使用以下工具：

添加自定义追踪事件#

默认情况下，追踪查看器中的事件主要是底层的内部 JAX 函数。您可以通过在代码中使用 jax.profiler.TraceAnnotation 和 jax.profiler.annotate_function() 来添加您自己的事件和函数。

配置分析器选项#

start_trace 方法接受一个可选的 profiler_options 参数，允许对分析器的行为进行细粒度控制。此参数应该是 jax.profiler.ProfileOptions 的一个实例。

例如，要禁用所有 Python 和宿主追踪：

import jax

options = jax.profiler.ProfileOptions()
options.python_tracer_level = 0
options.host_tracer_level = 0
jax.profiler.start_trace("/tmp/profile-data", profiler_options=options)

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

通用选项#

host_tracer_level：设置宿主端活动的追踪级别。

支持的值：

0：完全禁用宿主 (CPU) 追踪。

1：仅启用用户插桩的 TraceMe 事件追踪。

2：包含 1 级追踪以及高级程序执行细节，如昂贵的 XLA 操作（默认）。

3：包含 2 级追踪以及更详细的、底层的程序执行细节，如廉价的 XLA 操作。
device_tracer_level：控制是否启用设备追踪。

支持的值：

0：禁用设备追踪。

1：启用设备追踪（默认）。
python_tracer_level：控制是否启用 Python 追踪。

支持的值：

0：禁用 Python 函数调用追踪（默认）。

1：启用 Python 追踪。

高级配置选项#

TPU 选项#

tpu_trace_mode：指定 TPU 追踪模式。

支持的值：

TRACE_ONLY_HOST：这意味着仅追踪宿主端 (CPU) 活动，不收集设备 (TPU/GPU) 追踪。

TRACE_ONLY_XLA：这意味着仅追踪设备上的 XLA 级操作。

TRACE_COMPUTE：这会追踪设备上的计算操作。

TRACE_COMPUTE_AND_SYNC：这会追踪设备上的计算操作和同步事件。

如果未提供“tpu_trace_mode”，trace_mode 默认为 TRACE_ONLY_XLA。
tpu_num_sparse_cores_to_trace：指定要追踪的 TPU 稀疏核心数量。
tpu_num_sparse_core_tiles_to_trace：指定要追踪的每个稀疏核心内的瓦片数量。
tpu_num_chips_to_profile_per_task：指定每个任务要分析的 TPU 芯片数量。

GPU 选项#

GPU 分析提供以下选项：

gpu_max_callback_api_events：设置 CUPTI 回调 API 收集的最大事件数。默认为 2*1024*1024。
gpu_max_activity_api_events：设置 CUPTI 活动 API 收集的最大事件数。默认为 2*1024*1024。
gpu_max_annotation_strings：设置可收集的最大注释字符串数量。默认为 1024*1024。
gpu_enable_nvtx_tracking：启用 CUPTI 中的 NVTX 跟踪。默认为 False。
gpu_enable_cupti_activity_graph_trace：为 CUDA 图启用 CUPTI 活动图追踪。默认为 False。
gpu_pm_sample_counters：一个以逗号分隔的字符串，包含要使用 CUPTI PM 采样功能收集的 GPU 性能监控指标（例如 "sm__cycles_active.avg.pct_of_peak_sustained_elapsed"）。PM 采样默认禁用。有关可用指标，请参阅 NVIDIA 的 CUPTI 文档。
gpu_pm_sample_interval_us：设置 CUPTI PM 采样的采样间隔（以微秒为单位）。默认为 500。
gpu_pm_sample_buffer_size_per_gpu_mb：设置 CUPTI PM 采样每个设备的系统内存缓冲区大小（以 MB 为单位）。默认为 64MB。支持的最大值为 4GB。
gpu_num_chips_to_profile_per_task：指定每个任务要分析的 GPU 设备数量。如果未指定、设置为 0 或设置为无效值，则将分析所有可用的 GPU。这可用于减少追踪收集的大小。
gpu_dump_graph_node_mapping：如果启用，会将 CUDA 图节点映射信息转储到追踪中。默认为 False。

示例#

options = ProfileOptions()
options.advanced_configuration = {"tpu_trace_mode" : "TRACE_ONLY_HOST", "tpu_num_sparse_cores_to_trace" : 2}

如果发现任何无法识别的键或选项值，将返回 InvalidArgumentError。

故障排除#

GPU 分析#

在 GPU 上运行的程序应在追踪查看器的顶部附近生成 GPU 流的追踪信息。如果您只看到宿主追踪，请检查程序日志和/或输出中是否存在以下错误消息。

如果您收到类似以下的错误：Could not load dynamic library 'libcupti.so.10.1'
完整错误：

W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libcupti.so.10.1'; dlerror: libcupti.so.10.1: cannot open shared object file: No such file or directory
2020-06-12 13:19:59.822799: E external/org_tensorflow/tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1422] function cupti_interface_->Subscribe( &subscriber_, (CUpti_CallbackFunc)ApiCallback, this)failed with error CUPTI could not be loaded or symbol could not be found.

将 libcupti.so 的路径添加到环境变量 LD_LIBRARY_PATH 中。（尝试使用 locate libcupti.so 查找路径。）例如：

export LD_LIBRARY_PATH=/usr/local/cuda-10.1/extras/CUPTI/lib64/:$LD_LIBRARY_PATH

如果您执行此操作后仍然收到 Could not load dynamic library 消息，请检查 GPU 追踪是否仍然显示在追踪查看器中。此消息有时即使在一切正常时也会出现，因为它会在多个位置查找 libcupti 库。

如果您收到类似以下的错误：failed with error CUPTI_ERROR_INSUFFICIENT_PRIVILEGES
完整错误：

E external/org_tensorflow/tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1445] function cupti_interface_->EnableCallback( 0 , subscriber_, CUPTI_CB_DOMAIN_DRIVER_API, cbid)failed with error CUPTI_ERROR_INSUFFICIENT_PRIVILEGES
2020-06-12 14:31:54.097791: E external/org_tensorflow/tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1487] function cupti_interface_->ActivityDisable(activity)failed with error CUPTI_ERROR_NOT_INITIALIZED

运行以下命令（请注意这需要重启）：

echo 'options nvidia "NVreg_RestrictProfilingToAdminUsers=0"' | sudo tee -a /etc/modprobe.d/nvidia-kernel-common.conf
sudo update-initramfs -u
sudo reboot now

有关更多信息，请参阅 NVIDIA 关于此错误的文档。

在远程机器上进行分析#

如果想要分析的 JAX 程序在远程机器上运行，一种选择是在远程机器上执行上述所有说明（特别是要在远程机器上启动 TensorBoard 服务器），然后使用 SSH 本地端口转发从本地机器访问 TensorBoard Web UI。使用以下 SSH 命令将默认的 TensorBoard 端口 6006 从本地转发到远程机器：

ssh -L 6006:localhost:6006 <remote server address>

或者如果您使用的是 Google Cloud：

$ gcloud compute ssh <machine-name> -- -L 6006:localhost:6006

多个 TensorBoard 安装#

如果启动 TensorBoard 时出现类似以下错误：ValueError: Duplicate plugins for name projector

通常是因为安装了两个版本的 TensorBoard 和/或 TensorFlow（例如 tensorflow、tf-nightly、tensorboard 和 tb-nightly pip 包都包含 TensorBoard）。卸载单个 pip 包可能导致 tensorboard 可执行文件被删除且难以替换，因此可能有必要卸载所有相关包并重新安装单个版本：

pip uninstall tensorflow tf-nightly tensorboard tb-nightly xprof xprof-nightly tensorboard-plugin-profile tbp-nightly
pip install tensorboard xprof

Nsight#

NVIDIA 的 Nsight 工具可用于在 GPU 上追踪和分析 JAX 代码。有关详细信息，请参阅 Nsight 文档。