在 Cloud TPU 虚拟机上对模型进行性能分析

性能剖析是您在 Cloud TPU 上优化模型性能的主要工具之一。核心性能分析工具称为 XProf，可从 OpenXLA/XProf GitHub 代码库获取。XProf 支持对所有基于 XLA 的框架（包括 JAX、Pytorch XLA 和 Tensorflow/Keras）进行分析。

捕获性能剖析文件

如需使用 XProf，您需要先捕获模型的性能分析文件。您可以通过以下两种方式捕获性能剖析：

以编程方式捕获
按需捕获（又称手动捕获）

使用程序化捕获时，您需要对模型代码进行注释，以指定要在代码中的哪个位置捕获性能分析文件。通常，您会捕获几个训练步骤的性能剖析文件，或者对模型中的特定代码块进行性能剖析。在不同的框架（JAX、Pytorch XLA 和 TensorFlow）中，可以通过多种不同的方式捕获轨迹，包括基于 API 的开始/停止轨迹或基于上下文管理器的轨迹。在 MaxText 等更高级别的框架中，您只需在运行 MaxText 时启用标志 profiler=xplane 即可启用性能分析文件捕获。

如果您想临时捕获性能分析文件，或者未启用程序化性能分析文件捕获，可以使用按需性能分析文件捕获。如果您在运行期间发现模型指标存在问题，并希望在某个时间段内立即捕获性能分析文件以诊断问题，那么此功能会很有用。

如需启用按需捕获性能分析文件，您仍需在代码中启动 XProf 服务器。启用 profiler.start_server 将在机器学习工作负载上启动一个 XProf 服务器，该服务器会监听按需捕获触发器，以开始捕获性能分析文件。您可以通过 TensorBoard 界面或使用 XProfiler 工具通过 CLI 触发按需性能分析文件捕获。

您可以访问以下网址，了解如何为不同框架启用程序化和按需捕获：

无论是通过程序还是按需捕获性能分析文件，您都需要指定要存储捕获的性能分析文件的位置。您可以选择将性能分析文件存储在 TPU 虚拟机上的目录中，该目录的路径类似于 /profiles/run-name；也可以使用 Cloud Storage，其路径类似于 gs://bucket-name/run-name/。

通过上述路径，您的性能分析文件将捕获在以下位置：

/profiles/run-name/plugins/profile/session1/

或

gs://bucket-name/run-name/plugins/profile/session1/。

会话表示在模型运行期间捕获一次性能分析文件。例如，如果您在训练运行中捕获了步骤 1-3 的性能剖析文件，之后又捕获了步骤 8-10 的性能剖析文件，那么这些性能剖析文件是同一运行的性能剖析文件，但第一次捕获的性能剖析文件（来自步骤 1-3）将为 session1，第二次捕获的性能剖析文件（来自步骤 8-10）将为 session2。

每次运行下都会显示不同的日期戳，以表示不同的会话。您可以在不同的会话中以编程方式、按需或以两者混合的方式捕获性能分析文件。如果您想比较模型训练不同部分的性能分析文件，这会非常有用。例如，比较训练开始时的性能分析文件与训练结束时的性能分析文件。

使用 TensorBoard 直观呈现性能分析文件

如需在 Tensorboard 中查看 XProf 性能分析文件，您需要安装 tensorboard-plugin-profile 插件。

为了在 Google Cloud上直观呈现捕获的性能分析文件，我们建议您在 TPU 虚拟机上使用 TensorBoard，或者使用 XProfiler 工具在虚拟机或 GKE Pod 上托管 TensorBoard。

如果您在 TPU 虚拟机上捕获了性能剖析文件，则可以在该 TPU 虚拟机上使用 TensorBoard 查看捕获的性能剖析文件。

如需在 TPU 虚拟机上安装 XProf 和 TensorBoard，您可以运行以下命令：

pip install tensorboard_plugin_profile tensorboard
pip install xprof

如需查看使用 TensorBoard 捕获的性能剖析文件，请运行以下命令：

$ tensorboard --logdir=profiles/run-name

或

$ tensorboard --logdir=profiles

如需从多次运行中加载多个性能分析文件，请将 TensorBoard 指向包含所有运行和性能分析文件的根目录，而不是指向包含特定运行性能分析文件的子目录。

用于在 Google Cloud上直观呈现 Tensorboard 性能分析文件的 Cloud-Diagnostics-XProf 库

借助 cloud-diagnostics-xprof（又称 XProfiler 库），您可以更轻松地托管 TensorBoard 并在 Google Cloud上直观呈现性能分析文件。如需了解详情，请参阅 cloud-diagnostics-xprof GitHub 代码库。

与在本地运行 XProf 和 TensorBoard 相比，cloud-diagnostics-xprof 代码库具有以下改进：

设置和封装 XProf 和 TensorBoard 依赖项
将性能分析文件存储在 Cloud Storage 中，这有助于长期保留数据和进行运行后分析（研究人员完成运行后，系统会删除捕获的本地性能分析文件）
通过在 Compute Engine 虚拟机或 GKE pod 上预配 TensorBoard 来快速加载大型性能分析文件和多个性能分析文件，并可根据用户对加载速度和费用的需求更改机器类型
创建链接以分享性能分析文件，并与团队成员和 Google 工程师协作
更轻松地对 GKE 和 Compute Engine 上的工作负载进行按需分析，以便选择运行工作负载的任何主机来捕获性能分析文件

在使用 XProfiler 库之前，您需要以编程方式或通过启动性能剖析服务器来捕获工作负载代码的性能剖析文件，以便稍后按需捕获性能剖析文件。

如需设置 XProfiler 库，您需要设置 Google Cloud CLI 和 Python 虚拟环境，然后只需运行以下命令：

pip install cloud-diagnostics-xprof

这将安装所有必需的 XProf 和 TensorBoard 依赖项。

接下来，如需创建用于托管 TensorBoard 的虚拟机或 GKE pod，请运行以下命令：

xprofiler create -z $ZONE -l $GCS_PATH

或

xprofiler create --GKE -z $ZONE -l $GCS_PATH

将 $ZONE 替换为任意可用区，并将 $GCS_PATH 替换为您的个性能分析文件轨迹的路径。您可以指定包含多次运行的所有性能分析文件轨迹的根目录，也可以指定一次运行的一组特定性能分析文件轨迹。

例如，如果您在以下目录中捕获性能分析文件：

gs://bucket-name/run1/plugins/profile/session1/<profile.xplane.pb>
gs://bucket-name/run1/plugins/profile/session2/<profile.xplane.pb>
gs://bucket-name/run2/plugins/profile/session1/<profile.xplane.pb>

您可以将 GCS 路径设置为根目录 (gs://bucket-name)：

xprofiler create -z $ZONE -l gs://bucket-name/

在 TensorBoard 界面中，您将看到所有运行和会话的所有性能分析文件，如 run1/session1、run1/session2、run2/session1 所示。

默认情况下，xprofiler create 将创建一个 Compute Engine 虚拟机，具体来说是一个 c4-highmem-8 虚拟机。您可以使用 -m 标志更改机器类型。如果您想在 GKE pod 上（而不是在 Compute Engine 虚拟机上）创建 TensorBoard 实例，可以将 –GKE 标志传递给 xprofiler create。在 GKE pod 上托管 TensorBoard 实例可让您更轻松地管理 TensorBoard 以及部署在 GKE 上的其余工作负载。

与本地托管的 TensorBoard 相比，托管 TensorBoard 的虚拟机或 GKE pod 可在 Google Cloud 上更快地加载大型性能剖析文件和多个性能剖析文件。根据我们的基准比较，对于使用默认 c4-highmem-8 虚拟机的首次加载，您会在几分钟内看到大约 1 GB 的加载性能分析文件。此外，您还可以根据自己的性能和费用需求选择合适的机器类型。

运行 xprofiler create 后，您会看到类似于以下内容的输出：

Instance for gs://<bucket> has been created.
You can access it via the following:
1. https://<id>-dot-us-<region>.notebooks.googleusercontent.com.
2. xprofiler connect -z <zone> -l gs://bucket-name -m ssh
Instance is hosted at xprof-97db0ee6-93f6-46d4-b4c4-6d024b34a99f VM.

第一个选项是一个链接，点击该链接可在 TensorBoard 上查看 XProf 性能分析文件。通过此链接，您可以与团队成员分享性能分析文件，还可以与帮助您在 Google Cloud上优化性能的 Google 工程师分享性能分析文件。

您可以根据存储性能分析文件数据的 Cloud Storage 存储桶的权限设置来控制谁有权访问该链接。如果该链接因某种原因无法正常运行，我们还提供了一种使用 SSH 连接到 TensorBoard 实例的方法，以便使用 xprofiler connect 命令查看您的性能分析文件。

如果您在工作负载代码中启用了 Cloud Profiler 服务器（这是以程序化方式捕获和按需捕获所必需的），并且想要执行按需分析，则可以通过以下两种方式实现：

a. 点击 TensorBoard 上的捕获性能剖析文件按钮。现在，您可以选择工作负载运行的任何设备主机进行分析。我们支持对在 GKE 或 Compute Engine 上运行的工作负载进行按需捕获。

b. 使用 XProfiler 捕获命令，提供可用区、Cloud Storage 存储桶、框架、宿主虚拟机/Pod 名称以及捕获时长（以毫秒为单位）等信息。此信息与您需要在 TensorBoard 界面中输入的信息相同。

如需详细了解如何使用 cloud-diagnostics-xprof 库，请参阅此 GitHub 页面。