Pytorch/XLA 概述

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

与常规 PyTorch 不同，常规 PyTorch 逐行执行代码，并且在获取 PyTorch 张量的值之前不会阻止执行，而 PyTorch XLA 的工作方式不同。它遍历 python 代码，并将 (PyTorch) XLA 张量上的操作记录在中间表示 (IR) 图中，直到遇到屏障（如下所述）。生成 IR 图的这个过程称为追踪（LazyTensor 追踪或代码追踪）。然后，PyTorch XLA 将 IR 图转换为称为 HLO（High-Level Opcodes，高级操作码）的较低级机器可读格式。HLO 是特定于 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 中用户可见的屏障是 xm.mark_step()，它会中断 IR 图，并导致代码在 XLA 设备上执行。xm.mark_step 的一个关键特性是，与同步操作不同，它不会在设备执行图时阻止进一步的追踪。但是，它会阻止访问正在物化的张量的值。

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
    xm.mark_step()

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

重要的是要强调，在 PyTorch XLA 中，如果循环结束时存在屏障，则会追踪 for 循环内的 Python 代码，并且每次迭代都会构建一个新图。这可能是重要的性能瓶颈。

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

图大小和编译时间之间的权衡也很重要。如果有一个大的 IR 图，XLA 编译器可能会花费大量时间来优化和融合操作。这可能会导致非常长的编译时间。但是，由于在编译期间执行的优化，稍后的执行可能会快得多。

有时值得使用 xm.mark_step() 中断 IR 图。如上所述，这将产生一个更小的图，该图可以在以后重用。但是，使图更小可能会减少 XLA 编译器原本可以完成的优化。

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

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 仓库包含用于训练和 Serving 许多 LLM 和扩散模型的示例。

将代码转换为 PyTorch XLA

修改代码的一般指南

• 将 cuda 替换为 xm.xla_device()

• 移除进度条，打印会访问 XLA 张量值的操作

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

• 使用 MPDeviceLoader 包装数据加载器

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

记住：每种情况都是独特的，因此您可能需要针对每种情况做不同的事情。

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

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

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

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

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

import torch_xla.core.xla_model as xm
self.device = xm.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 = xm.xla_device()
attr = attr.to(torch.device(device))

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

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

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

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

示例 2. HF Stable Diffusion 推理

现在，考虑在 HuggingFace diffusers 库中使用 Stable Diffusion 推理，用于模型的 SD-XL 和 2.1 版本。供您参考，下面描述的更改可以在此 仓库 中找到。您可以克隆仓库并在您的 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 设备上运行

本节介绍需要对 text_to_image 推理示例 代码进行的更改，以使其在 TPU 上运行。

原始代码使用 Lora 进行推理，但本教程不会使用它。相反，在初始化管道时，我们将 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 许可证。接下来，转到 account→settings→access 令牌并生成新令牌。复制令牌并在您的 vm 上使用该特定令牌值运行以下命令

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

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

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

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

如果您尝试在不进行任何额外更改的情况下运行代码，您会注意到编译时间非常长（>6 小时）。这是因为 XLA 编译器尝试一次为所有调度器步骤构建一个图，类似于我们在上一个示例中讨论的内容。为了使代码运行更快，我们需要使用 xm.mark_step() 将图分解成更小的部分，并在下一步中重用它们。这发生在 function 中的 pipe.__call__ 这些行 中。禁用进度条、删除回调并在 for 循环结束时添加 xm.mark_step() 可以显著加快代码速度。更改在此 commit 中提供。

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

性能分析和性能分析

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

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

如果性能分析中的间隙是由于主机上发生的 Python 代码追踪引起的，那么这可能是一个瓶颈，并且没有进一步可以直接完成的优化。否则，应进一步分析代码以了解注意事项并进一步提高性能。请注意，您不能 xp.Trace() 包装调用 xm.mark_step() 的代码部分。

为了说明这一点，我们可以查看已捕获并上传到 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，这表明主机上正在发生某些事情，这些事情正在等待计算完成才能继续进行。查看水印 代码，我们可以看到张量被传输到 cpu 并转换为 numpy 数组，以便稍后使用 cv2 和 pywt 库进行处理。由于这部分不容易优化，我们将保持原样。

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

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

在多个 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 文件进行了修改。

在 Pods 上运行

一旦您有了在单个主机设备上运行的代码，就不需要进一步的更改。您可以创建 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 环境变量禁用此功能。