性能分析计算#

使用 Perfetto 查看程序跟踪#

我们可以使用 JAX 性能分析器生成 JAX 程序的跟踪，并使用 Perfetto 可视化工具 进行可视化。目前，此方法会阻塞程序，直到点击链接并加载 Perfetto UI 加载跟踪。如果您希望在没有任何交互的情况下获得性能分析信息，请查看下面的 XProf 性能分析器。

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 中的“共享”按钮来创建可与他人共享的跟踪永久链接。

远程性能分析#

在对远程运行的代码（例如托管的虚拟机）进行性能分析时，您需要为端口 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.start_server(<port>) 在感兴趣的脚本中启动性能分析服务器，而不是使用 jax.profiler.trace 以编程方式捕获跟踪。如果您只需要在脚本的某个部分激活性能分析器服务器，可以通过调用 jax.profiler.stop_server() 来关闭它。

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

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

默认情况下，生成的跟踪信息会被转储到临时目录中，但可以通过传递 --log_dir=<directory of choice> 来覆盖。此外，默认情况下，程序会提示您打开一个指向 ui.perfetto.dev 的链接。打开链接后，Perfetto UI 将加载跟踪文件并打开可视化工具。通过将 --no_perfetto_link 传递到命令中可以禁用此功能。或者，您也可以将 Tensorboard 指向 log_dir 来分析跟踪（请参阅下面的“XProf（Tensorboard 性能分析）”部分）。

XProf（TensorBoard 性能分析）#

XProf 可用于对 JAX 程序进行性能分析。XProf 是获取和可视化程序性能跟踪和概要的绝佳方式，包括 GPU 和 TPU 上的活动。最终结果大致如下：

XProf example

安装#

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

pip install xprof

如果您安装了 TensorBoard，xprof pip 包还将安装 TensorBoard 性能分析器插件。请注意，只能安装一个版本的 TensorFlow 或 TensorBoard，否则您可能会遇到“重复插件”错误，该错误在 下面 有描述。有关安装 TensorBoard 的更多信息，请参阅 https://tensorflowcn.cn/guide/profiler。

使用 TensorBoard 的夜间版本进行性能分析需要使用夜间版本的 XProf。

pip install tb-nightly xprof-nightly

程序化捕获#

您可以通过 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() 调用。我们使用它来确保设备上的执行被跟踪捕获。有关为什么需要这样做，请参阅 异步分派。

您也可以使用 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 秒跟踪的说明。

  1. 启动 XProf 服务器

    xprof --logdir /tmp/profile-data/

    您应该能够通过 https://:8791/ 加载 XProf。您可以使用 --port 标志指定不同的端口。如果您在远程服务器上运行 JAX，请参阅下面的 在远程机器上进行性能分析。

  2. 在您想要进行性能分析的 Python 程序或进程中，在开头附近添加以下内容：

    import jax.profiler
jax.profiler.start_server(9999)

    这将启动 XProf 连接到的性能分析器服务器。在进行下一步之前，必须运行性能分析器服务器。完成服务器使用后，您可以调用 jax.profiler.stop_server() 来关闭它。

    如果您想分析一个长时间运行的程序的片段（例如，一个长时间的训练循环），您可以将其放在程序开头，然后像往常一样运行程序。如果您想分析一个简短的程序（例如，一个微基准测试），一种选择是在 IPython shell 中启动性能分析器服务器，并在下一步启动捕获后使用 %run 来运行简短程序。另一种选择是启动程序开头时启动性能分析器服务器，并使用 time.sleep() 来为您提供足够的时间开始捕获。

  3. 打开 https://:8791/，然后点击左上角的“CAPTURE PROFILE”按钮。输入“localhost:9999”作为配置文件服务 URL（这是您在上一步中启动的性能分析器服务器的地址）。输入您想要进行性能分析的毫秒数，然后点击“CAPTURE”。

  4. 如果尚未运行您想要进行性能分析的代码（例如，如果您在 Python shell 中启动了性能分析器服务器），请在捕获运行时运行它。

  5. 捕获完成后，XProf 应该会自动刷新。（并非所有 XProf 性能分析功能都已与 JAX 集成，因此一开始可能看起来没有任何内容被捕获。）在左侧的“Tools”下，选择 trace_viewer。

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

您还可以使用以下工具

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)

添加自定义跟踪事件#

默认情况下，跟踪查看器中的事件大部分是低级别的内部 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()

通用选项#

  1. host_tracer_level：设置主机端活动的跟踪级别。

    支持的值

    0：完全禁用主机（CPU）跟踪。

    1：仅启用用户仪器化的 TraceMe 事件的跟踪（这是默认值）。

    2：包含级别 1 跟踪以及高级程序执行详细信息，例如昂贵的 XLA 操作。

    3：包含级别 2 跟踪以及更详细的低级别程序执行详细信息，例如廉价的 XLA 操作。

  2. python_tracer_level：控制是否启用 Python 跟踪。

    支持的值

    0：禁用 Python 函数调用跟踪。

    1：启用 Python 跟踪（这是默认值）。

高级配置选项#

TPU 选项#

  1. 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。

  2. tpu_num_sparse_cores_to_trace：指定要跟踪的 TPU 上的稀疏核心数量。

  3. tpu_num_sparse_core_tiles_to_trace：指定要跟踪的 TPU 上稀疏核心内的图块数量。

  4. 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：一个逗号分隔的 GPU 性能监视指标字符串，使用 CUPTI 的 PM 采样功能收集（例如，"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_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 文档。