性能分析计算#
使用 Perfetto 查看程序跟踪#
我们可以使用 JAX 性能分析器生成 JAX 程序的跟踪,这些跟踪可以使用 Perfetto 可视化工具 进行可视化。目前,此方法会阻止程序,直到单击链接并且 Perfetto UI 加载跟踪。如果您希望在没有任何交互的情况下获取性能分析信息,请查看下面的 Tensorboard 性能分析器。
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 将加载跟踪文件并打开可视化工具。
加载链接后,程序执行将继续。该链接在打开一次后不再有效,但它将重定向到一个新的 URL,该 URL 仍然有效。然后,您可以单击 Perfetto UI 中的“共享”按钮来创建指向跟踪的永久链接,该链接可以与他人共享。
远程性能分析#
当分析远程运行的代码(例如在托管 VM 上)时,您需要为链接工作在端口 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=<directory of choice>
来覆盖。此外,默认情况下,程序将提示您打开指向 ui.perfetto.dev
的链接。当您打开链接时,Perfetto UI 将加载跟踪文件并打开可视化工具。通过将 --no_perfetto_link
传入命令,可以禁用此功能。或者,您也可以将 Tensorboard 指向 log_dir
以分析跟踪(请参阅下面的“Tensorboard 性能分析”部分)。
TensorBoard 性能分析#
TensorBoard 的性能分析器 可用于分析 JAX 程序。Tensorboard 是获取和可视化程序性能跟踪和配置文件的好方法,包括 GPU 和 TPU 上的活动。最终结果看起来像这样
安装#
TensorBoard 性能分析器仅适用于与 TensorFlow 捆绑的 TensorBoard 版本。
pip install tensorflow tensorboard-plugin-profile
如果您已经安装了 TensorFlow,您只需要安装 tensorboard-plugin-profile
pip 包。请注意只安装一个版本的 TensorFlow 或 TensorBoard,否则您可能会遇到 下面 描述的“重复插件”错误。有关安装 TensorBoard 的更多信息,请参阅 https://tensorflowcn.cn/guide/profiler。
TensorBoard 性能分析器的 nightly 版本需要 nightly tensorflow 和 tensorboard
pip install tf-nightly tb-nightly tbp-nightly
程序化捕获#
您可以使用 jax.profiler.start_trace()
和 jax.profiler.stop_trace()
方法来检测您的代码以捕获性能分析器跟踪。调用 start_trace()
并指定要将跟踪文件写入到的目录。这应该是用于启动 TensorBoard 的相同 --logdir
目录。然后,您可以使用 TensorBoard 查看跟踪。
例如,要进行性能分析器跟踪
import jax
jax.profiler.start_trace("/tmp/tensorboard")
# 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/tensorboard"):
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()
要查看跟踪,请首先启动 TensorBoard(如果您尚未启动)
$ tensorboard --logdir=/tmp/tensorboard
[...]
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.5.0 at http://localhost:6006/ (Press CTRL+C to quit)
您应该能够在此示例中在 http://localhost:6006/ 加载 TensorBoard。您可以使用 --port
标志指定不同的端口。如果要在远程服务器上运行 JAX,请参阅下面的 在远程机器上进行性能分析。
然后,在右上角的下拉菜单中选择“Profile”,或直接转到 http://localhost:6006/#profile。可用的跟踪会出现在左侧的“Runs”下拉菜单中。选择您感兴趣的运行,然后在“Tools”下,选择 trace_viewer
。您现在应该看到执行的时间线。您可以使用 WASD 键导航跟踪,并单击或拖动以选择事件,以在底部查看更多详细信息。有关使用跟踪查看器的更多详细信息,请参阅 这些 TensorFlow 文档。
您还可以使用 memory_viewer
、op_profile
和 graph_viewer
工具。
通过 TensorBoard 手动捕获#
以下是从正在运行的程序中手动触发 N 秒跟踪的说明。
启动 TensorBoard 服务器
tensorboard --logdir /tmp/tensorboard/
您应该能够在 http://localhost:6006/ 加载 TensorBoard。您可以使用
--port
标志指定不同的端口。如果要在远程服务器上运行 JAX,请参阅下面的 在远程机器上进行性能分析。在您要分析的 Python 程序或进程中,在开头附近添加以下内容
import jax.profiler jax.profiler.start_server(9999)
这将启动 TensorBoard 连接到的性能分析器服务器。性能分析器服务器必须在您继续下一步之前运行。当您完成使用服务器后,可以调用
jax.profiler.stop_server()
将其关闭。如果您想分析长时间运行的程序片段(例如,长时间的训练循环),您可以将此代码放在程序的开头,并像往常一样启动程序。如果您想分析短程序(例如,微基准测试),一种选择是在 IPython shell 中启动性能分析器服务器,并在下一步启动捕获后使用
%run
运行短程序。另一种选择是在程序开始时启动性能分析器服务器,并使用time.sleep()
给您足够的时间来启动捕获。打开 http://localhost:6006/#profile,然后单击左上角的“CAPTURE PROFILE”按钮。输入 “localhost:9999” 作为性能分析服务 URL(这是您在上一步中启动的性能分析器服务器的地址)。输入您要分析的毫秒数,然后单击“CAPTURE”。
如果您要分析的代码尚未运行(例如,如果您在 Python shell 中启动了性能分析器服务器),请在捕获运行时运行它。
捕获完成后,TensorBoard 应该会自动刷新。(并非所有 TensorBoard 性能分析功能都与 JAX 挂钩,因此最初可能看起来像没有捕获任何内容。)在左侧的“Tools”下,选择
trace_viewer
。您现在应该看到执行的时间线。您可以使用 WASD 键导航跟踪,并单击或拖动以选择事件,以在底部查看更多详细信息。有关使用跟踪查看器的更多详细信息,请参阅 这些 TensorFlow 文档。
您还可以使用
memory_viewer
、op_profile
和graph_viewer
工具。
添加自定义跟踪事件#
默认情况下,跟踪查看器中的事件主要是低级内部 JAX 函数。您可以使用代码中的 jax.profiler.TraceAnnotation
和 jax.profiler.annotate_function()
添加您自己的事件和函数。
故障排除#
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
pip install tensorflow
Nsight#
NVIDIA 的 Nsight
工具可用于跟踪和分析 GPU 上 JAX 代码的性能。有关详细信息,请参阅 Nsight
文档。