Pytorch/XLA 概述

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

与常规的 PyTorch 不同，常规 PyTorch 逐行执行代码，并且在获取 PyTorch 张量的值之前不会阻塞执行，而 PyTorch XLA 的工作方式有所不同。它遍历 Python 代码并将对 (PyTorch) XLA 张量的操作记录在中间表示 (IR) 图中，直到遇到屏障（下文将讨论）。生成 IR 图的过程称为追踪（LazyTensor 追踪或代码追踪）。然后，PyTorch XLA 将 IR 图转换为称为 HLO（高级操作码）的低级机器可读格式。HLO 是特定于 XLA 编译器的计算表示形式，它允许编译器为其运行的硬件生成高效代码。HLO 被输入到 XLA 编译器进行编译和优化。PyTorch XLA 随后会缓存编译结果，以便以后在需要时重用。图的编译在主机（CPU）上进行，主机是运行 Python 代码的机器。如果存在多个 XLA 设备，主机将为每个设备单独编译代码，除非使用 SPMD（单程序，多数据）。例如，v4-8 有一台主机和 四个设备。在这种情况下，主机将为这四个设备中的每一个单独编译代码。对于 Pod Slice（包含多个主机的情况），每个主机都会为其连接的 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 仓库包含许多 LLM 和扩散模型的训练和服务示例。

将代码转换为 PyTorch XLA

修改代码的一般准则

• 将 cuda 替换为 xm.xla_device()

• 移除进度条、会访问 XLA 张量值的打印语句

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

• 使用 MPDeviceLoader 包装数据加载器

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

记住：每个案例都是独特的，因此您可能需要针对不同情况采取不同的做法。

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

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

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 许可证。接下来，前往 账户→设置→访问令牌 并生成一个新令牌。复制该令牌并在您的 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 = 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() 将图分解成更小的部分并在后续步骤中重用它们。这发生在 pipe.__call__ 这些行中的函数内部。禁用进度条、移除回调并在 for 循环末尾添加 xm.mark_step() 会显著加快代码速度。更改在此提交中提供。

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

性能分析和性能评估

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

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

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

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

从 Stable Diffusion 模型版本 2.1 开始

如果我们不插入任何追踪就捕获配置文件，我们将看到以下内容

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

如果我们添加跟踪并使用设备上能容纳的最大批量大小（在本例中为 32）重新捕获配置文件，我们将看到设备上的空隙是由主机上运行的 Python 进程引起的。

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

现在，让我们检查模型的 XL 版本，并做同样的事情。我们将像处理 2.1 版本一样，将跟踪添加到 管线文件 中并捕获配置文件。

这次，除了中间由 pipe_watermark 跟踪引起的较大空隙之外，在 这个循环 中，推理步骤之间还有许多小的空隙。

首先仔细查看由 pipe_watermark 引起的较大空隙。空隙之前是 TransferFromDevice 操作，这表明主机上正在进行某些操作，这些操作在继续之前正在等待计算完成。查看水印 代码，我们可以看到张量被传输到 CPU 并转换为 numpy 数组，以便稍后使用 cv2 和 pywt 库进行处理。由于这部分优化起来并不简单，我们将暂时保持原样。

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

如果我们研究 U-Net 函数和调度器，我们可以看到 U-Net 代码不包含任何 PyTorch/XLA 的优化目标。然而，在 scheduler.step 内部有 .item() 和 .nonzero() 调用。我们可以 重写 该函数以避免这些调用。如果我们修复这个问题并重新运行配置文件，差异不会很大。但是，由于我们减少了引入较小图的设备-主机通信，这使得编译器能够更好地优化代码。函数 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 版本。对于此版本，对 管线 文件进行了与上述类似的更改。

在 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 环境变量来禁用。