Pytorch/XLA 概览

本节简要概述 PyTorch XLA 的基本细节，这将帮助读者更好地理解所需的代码修改和优化。

与普通 PyTorch 不同，普通 PyTorch 按行执行代码，并且在获取 PyTorch 张量值之前不会阻止执行，PyTorch XLA 的工作方式有所不同。它会遍历 Python 代码，并将 XLA 张量（PyTorch）上的操作记录在中介表示（IR）图表中，直到遇到屏障（下文讨论）。此生成 IR 图的过程称为跟踪（LazyTensor 跟踪或代码跟踪）。然后，PyTorch XLA 将 IR 图转换为称为 HLO（High-Level Opcodes）的较低级机器可读格式。HLO 是 XLA 编译器特有的计算表示，它允许 XLA 为正在运行的硬件生成高效的代码。HLO 被输入到 XLA 编译器进行编译和优化。然后，PyTorch XLA 会缓存编译结果，以便稍后需要时重复使用。图的编译在主机（CPU）上完成，即运行 Python 代码的机器。如果存在多个 XLA 设备，则主机将为每个设备分别编译代码，除非使用 SPMD（单程序多数据）。例如，v4-8 有一台主机和 四个设备。在这种情况下，主机将为这四个设备分别编译代码。对于 pod 切片，当存在多个主机时，每个主机负责编译其连接的 XLA 设备的代码。如果使用 SPMD，则代码（对于给定的形状和计算）将在每个主机上仅编译一次，供所有设备使用。

有关更多详细信息和示例，请参阅 LazyTensor 指南。

IR 图中的操作仅在需要张量值时执行。这称为张量的求值或物化。有时也称为惰性求值，它可以带来显著的 性能提升。

PyTorch XLA 中的同步操作，例如打印、日志记录、检查点或回调会阻塞跟踪并导致执行速度变慢。当操作需要 XLA 张量的特定值时，例如 print(xla_tensor_z)，跟踪会被阻塞，直到该张量的值对主机可用为止。请注意，仅执行负责计算该张量值的图的部分。这些操作不会切断 IR 图，但会通过 TransferFromDevice 触发主机-设备通信，从而导致性能下降。

屏障是一个特殊指令，它告诉 XLA 执行 IR 图并物化张量。这意味着 PyTorch XLA 张量将被求值，并且结果将对主机可用。PyTorch XLA 中用户暴露的屏障是 torch_xla.sync()，它会中断 IR 图，并导致 XLA 设备上的代码执行。`torch_xla.sync()` 的一个关键特性是，与同步操作不同，它不会在设备执行图时阻塞进一步的跟踪。但是，它确实会阻塞对正在物化的张量值的访问。

LazyTensor 指南中的示例说明了两个张量相加的简单情况。现在，假设我们有一个 for 循环，它会添加 XLA 张量并在之后使用该值

for x, y in tensors_on_device:
    z += x + y

没有屏障的情况下，Python 跟踪将生成一个单一的图，该图将 len(tensors_on_device) 次加法运算包装起来。这是因为 for 循环没有被跟踪捕获，因此循环的每次迭代都会创建一个新的子图，对应于 z += x+y 的计算，并将其添加到图中。这里是一个 len(tensors_on_device)=3 的示例。

然而，在循环末尾引入一个屏障将生成一个较小的图，该图将在 `for` 循环内的第一次传递中编译一次，并在接下来的 `len(tensors_on_device)-1` 次迭代中重用。屏障将向跟踪发出信号，表明到目前为止跟踪的图可以提交执行，并且如果该图以前被见过，将重用缓存的编译程序。

for x, y in tensors_on_device:
    z += x + y
    torch_xla.sync()

在这种情况下，将有一个小图被使用 len(tensors_on_device)=3 次。

需要强调的是，在 PyTorch XLA 中，如果 for 循环的末尾有屏障，那么 for 循环内的 Python 代码将被跟踪，并且每次迭代都会构建一个新的图。这可能成为一个重要的性能瓶颈。

当相同的计算发生在相同形状的张量上时，XLA 图可以被重用。如果输入或中间张量的形状发生变化，XLA 编译器将重新编译一个带有新张量形状的新图。这意味着，如果您有动态形状，或者您的代码不重用张量图，那么在 XLA 上运行模型将不适合该用例。将输入填充到固定形状可以作为一种避免动态形状的选项。否则，编译器将在优化和融合将不再使用的操作上花费大量时间。

图大小和编译时间之间的权衡也很重要。如果有一个大型 IR 图，XLA 编译器可能会花费大量时间来优化和融合操作。这可能导致很长的编译时间。然而，由于编译过程中进行的优化，后续执行可能会快得多。

有时使用 `torch_xla.sync()` 来中断 IR 图是有价值的。如上所述，这将生成一个可以稍后重用的小图。然而，减小图的大小可能会减少 XLA 编译器可以进行的优化。

另一个需要考虑的重要点是 MPDeviceLoader。一旦您的代码在 XLA 设备上运行，请考虑将 torch 数据加载器包装在 XLA MPDeviceLoader 中，它会将数据预加载到设备以提高性能，并且包含 torch_xla.sync()。后者会自动中断数据批次的迭代并将它们发送执行。请注意，如果您不使用 MPDeviceLoader，则在 optimizer_step() 中可能需要将 barrier=True 设置为启用 torch_xla.sync()（如果您正在运行训练作业），或者显式添加 torch_xla.sync()。

TPU 设置

创建带有基础映像的 TPU 以使用 nightly wheels，或通过指定 `RUNTIME_VERSION` 从稳定版本创建。

export ZONE=us-central2-b
export PROJECT_ID=your-project-id
export ACCELERATOR_TYPE=v4-8 # v4-16, v4-32, …
export RUNTIME_VERSION=tpu-vm-v4-pt-2.0 # or tpu-vm-v4-base
export TPU_NAME=your_tpu_name

gcloud compute tpus tpu-vm create ${TPU_NAME} \
--zone=${ZONE} \
--accelerator-type=${ACCELERATOR_TYPE} \
--version=${RUNTIME_VERSION} \
--subnetwork=tpusubnet

如果您有一个单主机 VM（例如 v4-8），您可以 ssh 到您的 vm 并直接从 vm 中运行以下命令。否则，对于 TPU pod，您可以使用 `--worker=all --command=""`，类似于

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone=us-central2-b \
--worker=all \
--command="pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly-cp38-cp38-linux_x86_64.whl"

接下来，如果您使用的是基础映像，请安装 nightly 包和所需的库

pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly-cp38-cp38-linux_x86_64.whl
​​pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly-cp38-cp38-linux_x86_64.whl
sudo apt-get install libopenblas-dev -y

sudo apt-get update && sudo apt-get install libgl1 -y # diffusion specific

参考实现

AI-Hypercomputer/tpu-recipies 存储库包含训练和提供 LLM 和扩散模型的大量示例。

将代码转换为 PyTorch XLA

修改代码的通用指南

• 将 `cuda` 替换为 `torch_xla.device()`

• 删除进度条、会访问 XLA 张量值的打印输出

• 减少会访问 XLA 张量值的日志记录和回调

• 用 MPDeviceLoader 包装数据加载器

• 进行性能分析以进一步优化代码

请记住：每个情况都是独特的，因此您可能需要针对每种情况采取不同的措施。

示例 1. 在 PyTorch Lightning 的单个 TPU 设备上进行 Stable Diffusion 推理

作为第一个例子，我们考虑 PyTorch Lightning 中 Stable Diffusion 模型的推理代码，可以从命令行如下运行：

    python scripts/txt2img.py --prompt "a photograph of an astronaut riding a horse"

供您参考，下面描述的修改的 diff 可以在 此处 找到。让我们一步一步地介绍它们。与上面的通用指南一样，从与 `cuda` 设备相关的更改开始。此推理代码编写用于在 GPU 上运行，并且 `cuda` 可以在多个地方找到。通过从此行删除 `model.cuda()`，并从此处删除 `precision_scope` 来开始进行更改。此外，将此行中的 `cuda` 设备替换为 `xla` 设备，类似于下面的代码：

接下来，此特定模型配置正在使用 `FrozenCLIPEmbedder`，因此我们将此行也进行修改。为了简单起见，我们将在本教程中直接定义 `device`，但您也可以将 `device` 值传递给函数。

    import torch_xla.core.xla_model as xm
    self.device = torch_xla.device()

代码中另一个包含 cuda 特定代码的地方是 DDIM 调度器。在文件顶部添加 `import torch_xla.core.xla_model as xm`，然后替换这些行。

if attr.device != torch.device("cuda"):
   attr = attr.to(torch.device("cuda"))

替换

device = torch_xla.device()
attr = attr.to(torch.device(device))

接下来，您可以通过删除打印语句、禁用进度条以及减少或删除会访问 XLA 张量值的回调和日志记录来减少设备（TPU）和主机（CPU）之间的通信。这些操作需要设备停止执行，回退到 CPU，执行日志记录/回调，然后返回到设备。这可能是一个重要的性能瓶颈，尤其是在大型模型上。

进行这些更改后，代码将在 TPU 上运行。但是，性能会非常慢。这是因为 XLA 编译器尝试构建一个单一的（巨大的）图，该图包装了推理步数（在本例中为 50），因为 for 循环内部没有屏障。编译器很难优化图，这会导致严重的性能下降。如上所述，使用屏障（torch_xla.sync()）中断 for 循环将生成一个更小的图，该图更容易由编译器优化。这还将允许编译器重用上一步的图，从而提高性能。

现在，代码已准备好在合理的时间内针对 TPU 运行。可以通过捕获性能分析并进一步研究来获得更多优化和分析。但是，这在此处不作涵盖。

注意：如果您在 v4-8 TPU 上运行，那么您有 4 个可用的 XLA (TPU) 设备。按上述方式运行代码将仅使用一个 XLA 设备。为了在所有 4 个设备上运行，您需要使用 `torch_xla.launch()` 函数在所有设备上生成代码。我们将在下一个示例中讨论 `torch_xla.launch`。

示例 2. HF Stable Diffusion 推理

现在，我们考虑使用 HuggingFace diffusers 库中的文本到图像推理，用于模型的 SD-XL 和 2.1 版本。供您参考，下面描述的更改可以在此 repo 中找到。您可以克隆该 repo 并在 TPU VM 上使用以下命令运行推理：

(vm)$ git clone https://github.com/pytorch-tpu/diffusers.git
(vm)$ cd diffusers/examples/text_to_image/
(vm)$ python3 inference_tpu_single_device.py

在单个 TPU 设备上运行

本节介绍需要对文本到图像推理示例代码进行的更改，以便在 TPU 上运行它。

原始代码使用 Lora 进行推理，但本教程不使用它。相反，在初始化 pipeline 时，我们将 `model_id` 参数设置为 `stabilityai/stable-diffusion-xl-base-0.9`。我们还将使用默认的调度器（DPMSolverMultistepScheduler）。但是，对其他调度器也可以进行类似的更改。

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install . # pip install -e .

cd examples/text_to_image/
pip install -r requirements.txt
pip install invisible_watermark transformers accelerate safetensors

(如果找不到 `accelerate`，请注销，然后重新登录。)

登录到 HF 并同意模型卡上的 sd-xl 0.9 许可证。接下来，转到 帐户→设置→访问令牌 并生成一个新令牌。复制令牌，并在您的 VM 上使用该特定令牌值运行以下命令：

(vm)$ huggingface-cli login --token _your_copied_token__

HuggingFace 的 README 提供了为在 GPU 上运行而编写的 PyTorch 代码。要在 TPU 上运行它，第一步是将 CUDA 设备更改为 XLA 设备。这可以通过将 `pipe.to("cuda")` 行替换为以下行来完成：

import torch_xla.core.xla_model as xm
device = torch_xla.device()
pipe.to(device)

此外，需要注意的是，第一次使用 XLA 进行推理时，编译时间会很长。例如，HuggingFace 的 Stable Diffusion XL 模型推理编译时间可能需要大约一个小时，而实际推理可能只需要 5 秒，具体取决于批量大小。同样，GPT-2 模型可能需要 10-15 分钟才能编译，之后训练 epoch 的时间会快得多。这是因为 XLA 会构建一个将要执行的计算图，然后针对它正在运行的特定硬件优化该图。但是，一旦图被编译，它就可以为后续的推理重用，这将快得多。因此，如果您只运行一次推理，您可能不会从使用 XLA 中受益。但是，如果您多次运行推理，或者对一系列提示运行推理，您将在前几次推理后开始看到 XLA 的优势。例如，如果您对 10 个提示的列表进行推理，第一次推理（可能一两次[^1]）可能需要很长时间来编译，但剩余的推理步骤将快得多。这是因为 XLA 会重用它为第一次推理编译的图。

如果您尝试在不进行任何额外更改的情况下运行代码，您会发现编译时间非常长（> 6 小时）。这是因为 XLA 编译器尝试一次为所有调度器步骤构建一个图，就像我们在上一个示例中讨论的那样。为了使代码运行得更快，我们需要使用 `torch_xla.sync()` 将图分解成更小的部分，并在后续步骤中重用它们。这发生在 `pipe.__call__` 函数的这些行中。禁用进度条、删除回调并在 for 循环末尾添加 `torch_xla.sync()` 可以显著加快代码速度。更改在此 commit 中提供。

此外，`self.scheduler.step()` 函数（默认使用 `DPMSolverMultistepScheduler` 调度器）存在一些问题，这些问题在 PyTorch XLA 警告中进行了描述。该函数中的 `.nonzero()` 和 `.item()` 调用会向 CPU 发送张量求值请求，这会触发设备-主机通信。这是不理想的，因为它会减慢代码速度。在这种特定情况下，我们可以通过直接将索引传递给函数来避免这些调用。这将防止函数向 CPU 发送请求，并提高代码性能。更改可在此 commit 中找到。代码现在已准备好在 TPU 上运行。

性能分析和性能分析

为了进一步研究模型的性能，我们可以使用性能分析 指南来对其进行分析。作为经验法则，应使用适合内存的最大批量大小来运行分析脚本，以实现最佳内存使用。重叠代码跟踪与设备执行也有助于更优地利用设备。分析的持续时间应足够长，以便捕获至少一个步骤。TPU 上模型的良好性能意味着设备-主机通信被最小化，并且设备正在持续运行进程，没有空闲时间。

在 `inference_tpu_*.py` 文件中启动服务器，并按照指南所述运行 `capture_profile.py` 脚本，将为我们提供设备上运行的进程的信息。目前，仅分析了一个 XLA 设备。为了更好地理解 TPU 的空闲时间（配置文件中的间隙），应将分析跟踪 (`xp.Trace()`) 添加到代码中。`xp.Trace()` 测量在主机机器上由跟踪包装的代码执行所需的时间。在此示例中，`xp.Trace()` 已添加到pipeline 和 U-net 模型内部，以测量在主机（CPU）上运行特定代码段的时间。

如果配置文件中的间隙是由于在主机上运行的 Python 代码跟踪引起的，那么这可能是一个瓶颈，并且当前没有进一步的直接优化可以进行。否则，应进一步分析代码以了解警告并进一步提高性能。请注意，您无法将调用 `torch_xla.sync()` 的代码部分包装在 `xp.Trace()` 中。

为了说明这一点，我们可以查看已上传到 tensorboard 的已捕获配置文件，遵循分析指南。

从 Stable Diffusion 模型版本 2.1 开始

如果我们捕获一个配置文件而不插入任何跟踪，我们将看到以下内容：

v4-8 上的单个 TPU 设备有两个核心，看起来是忙碌的。除了中间有一个小间隙外，它们的利用率没有明显间隙。如果向上滚动以查找占用主机进程的进程，我们将找不到任何信息。因此，我们将 `xp.traces` 添加到 pipeline 文件以及 U-net 函数。后者对于此特定用例可能没有用，但它确实演示了如何在不同位置添加跟踪以及它们的信息如何在 TensorBoard 中显示。

如果我们添加跟踪并重新捕获最大批量大小（在此情况下为 32）的配置文件，我们将看到设备中的间隙是由在主机上运行的 Python 进程引起的。

我们可以使用适当的工具放大时间线，查看在此期间运行的是哪个进程。此时，Python 代码跟踪在主机上进行，我们此时无法进一步改进跟踪。

现在，让我们检查模型的 XL 版本并执行相同的操作。我们将以与 2.1 版本相同的方式将跟踪添加到 pipeline 文件中，并捕获一个配置文件。

这次，除了中间的巨大间隙（由 `pipe_watermark` 跟踪引起）之外，在此循环内的推理步骤之间有许多小间隙。

首先仔细查看由 `pipe_watermark` 引起的巨大间隙。该间隙之前是 `TransferFromDevice`，这表明主机上正在发生某些事情，它正在等待计算完成后再继续。查看 watermark 代码，我们可以看到张量被传输到 CPU 并转换为 numpy 数组，以便稍后用 `cv2` 和 `pywt` 库进行处理。由于这部分不容易优化，我们将保持原样。

现在，如果我们放大循环，我们可以看到循环内的图被分解成更小的部分，因为发生了 `TransferFromDevice` 操作。

如果我们检查 U-Net 函数和调度器，我们可以看到 U-Net 代码不包含 PyTorch/XLA 的任何优化目标。但是，scheduler.step 中有 `.item()` 和 `.nonzero()` 调用。我们可以重写该函数以避免这些调用。如果我们修复此问题并重新运行性能分析，我们不会看到太大差异。但是，由于我们减少了引入较小图的设备-主机通信，我们允许编译器更好地优化代码。`scale_model_input` 函数具有类似问题，我们可以通过对 `step` 函数进行我们上面所做的更改来修复这些问题。总的来说，由于许多间隙是由 python 级别的代码跟踪和图构建引起的，因此使用当前版本的 PyTorch XLA 无法优化这些间隙，但当 dynamo 在 PyTorch XLA 中启用时，我们可能会看到改进。

在多个 TPU 设备上运行

要使用多个 TPU 设备，您可以使用 `torch_xla.launch` 函数来生成您在单个设备上运行的函数到多个设备。`torch_xla.launch` 函数将在多个 TPU 设备上启动进程并在需要时同步它们。这可以通过将 `index` 参数传递给在单个设备上运行的函数来实现。例如，

import torch_xla

def my_function(index):
  # function that runs on a single device

torch_xla.launch(my_function, args=(0,))

在此示例中，`my_function` 函数将在 v4-8 上的 4 个 TPU 设备上生成，每个设备被分配一个从 0 到 3 的索引。请注意，默认情况下，launch() 函数将在所有 TPU 设备上生成进程。如果您只想运行单个进程，请将参数设置为 `launch(..., debug_single_process=True)`。

此文件说明了如何使用 xmp.spawn 在多个 TPU 设备上运行 stable diffusion 2.1 版本。对于此版本，对 pipeline 文件进行了类似的更改。

在 Pod 上运行

一旦您有了在单主机设备上运行的代码，就不需要进一步更改了。您可以创建 TPU pod，例如，按照这些说明。然后使用以下命令运行您的脚本：

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
  --zone=${ZONE} \
  --worker=all \
  --command="python3 your_script.py"

注意

0 和 1 是 XLA 中的魔术数字，在 HLO 中被视为常量。因此，如果代码中有一个随机数生成器可以生成这些值，代码将为每个值单独编译。这可以通过 `XLA_NO_SPECIAL_SCALARS=1` 环境变量禁用。