理解 LazyTensor 系统在 PyTorch/XLA 和 Cloud TPU 上的性能

作者：Vaibhav Singh

引言

易用性、表达力和可调试性是 PyTorch 的核心原则之一。其易用性的关键驱动因素之一是 PyTorch 默认采用“即时（eager）”执行方式，即按操作（op by op）执行保留了程序的命令式特性。然而，即时执行不提供基于编译器的优化，例如当计算可以表示为图时进行的优化。

LazyTensor [1] 最初随 PyTorch/XLA 一起推出，有助于结合这两种看似不同的方法。虽然 PyTorch 的即时执行方式已被广泛使用、直观且易于理解，但惰性（lazy）执行方式目前尚不那么普遍。

在本文中，我们将探讨 LazyTensor 系统的一些基本概念，旨在应用这些概念来理解和调试 PyTorch 中基于 LazyTensor 的实现的性能。尽管我们将使用 PyTorch/XLA 在 Cloud TPU 上作为探索这些概念的载体，但我们希望这些想法对于理解构建在 LazyTensor 之上的其他系统也会有所帮助。

LazyTensor

对 PyTorch 张量执行的任何操作默认都会作为内核或内核组合分派到底层硬件。这些内核在底层硬件上异步执行。程序执行直到张量值被获取时才会被阻塞。这种方法与 GPU 等大规模并行编程硬件配合得非常好。

LazyTensor 系统的起点是一个自定义张量类型。在 PyTorch/XLA 中，这种类型称为 XLA 张量。与 PyTorch 的原生张量类型不同，对 XLA 张量执行的操作被记录到一个 IR（中间表示）图中。让我们看一个计算两个张量乘积之和的例子。

import torch
import torch_xla
import torch_xla.core.xla_model as xm

dev = xm.xla_device()

x1 = torch.rand((3, 3)).to(dev)
x2 = torch.rand((3, 8)).to(dev)

y1 = torch.einsum('bs,st->bt', x1, x2)
print(torch_xla._XLAC._get_xla_tensors_text([y1]))

您可以运行这个 colab 笔记本，以检查 y1 的结果图。请注意，目前还没有执行任何计算。

y1 = y1 + x2
print(torch_xla._XLAC._get_xla_tensors_text([y1]))

操作将持续进行，直到 PyTorch/XLA 遇到一个 barrier（屏障）。这个 barrier 可以是 mark step() API 调用，或任何其他强制执行迄今为止记录的图的事件。

xm.mark_step()
print(torch_xla._XLAC._get_xla_tensors_text([y1]))

一旦调用 mark_step()，图就会被编译，然后在 TPU 上执行，即张量被具体化（materialized）。因此，图现在被简化为一个单独的 y1 张量，其中包含计算结果。

编译一次，多次执行

XLA 编译过程提供了优化（例如，通过使用暂存存储器（scratch-pad memory）进行多个操作（ops）来减少 HBM 压力的操作融合（op-fusion），参考），并利用较低级别的 XLA 基础设施来优化使用底层硬件。然而，有一个注意事项，编译过程是昂贵的，即会增加训练步骤的时间。因此，只有当我们能够编译一次并多次执行（编译缓存（compilation cache）有所帮助，以确保同一张图不会被多次编译）时，这种方法才能很好地扩展。

在下面的例子中，我们创建一个小的计算图并计时其执行

y1 = torch.rand((3, 8)).to(dev)
def dummy_step() :
  y1 = torch.einsum('bs,st->bt', y1, x)
  xm.mark_step()
  return y1

%timeit dummy_step

The slowest run took 29.74 times longer than the fastest. This could mean that an intermediate result is being cached.
10000000 loops, best of 5: 34.2 ns per loop

您会注意到最慢的步骤比最快的要长得多。这是因为图编译的开销，对于给定的图形状、输入形状和输出形状，这种开销只会发生一次。后续步骤更快是因为无需进行图编译。

这也意味着当“编译一次并多次执行”的假设被打破时，我们预计会看到性能瓶颈。理解这个假设何时被打破是理解和优化 LazyTensor 系统性能的关键。让我们看看是什么触发了编译。

图编译、执行和 LazyTensor Barrier

我们看到计算图在遇到 LazyTensor barrier 时会被编译和执行。有三种情况会自动或手动引入 LazyTensor barrier。第一种是显式调用 mark_step() API，如前例所示。当您使用 MpDeviceLoader 包装数据加载器时，mark_step() 也会在每个步骤中隐式调用（强烈建议这样做以重叠计算和数据上传到 TPU 设备）。xla_model 的 Optimizer step 方法也允许隐式调用 mark_step（当您设置 barrier=True 时）。

引入 barrier 的第二种情况是当 PyTorch/XLA 找到一个没有对应 XLA HLO op 映射（lowering）的操作（op）时。PyTorch 有 2000 多个操作。虽然其中大多数操作是复合的（即可以用其他基本操作表示），但其中一些操作在 XLA 中没有相应的 lowering。

当使用没有 XLA lowering 的操作时会发生什么？PyTorch XLA 停止操作记录，并切断通向未 lowering 操作输入的图。然后编译并分派这个被切断的图进行执行。执行结果（具体化的张量）从设备发送回主机，然后未 lowering 的操作在主机（CPU）上执行，接着下游的 LazyTensor 操作创建新的图，直到再次遇到 barrier。

导致 LazyTensor barrier 的第三种也是最后一种情况是存在需要张量值的控制结构/语句或其他方法。这种语句至少会触发通向该张量的计算图的执行（如果该图已被看到），或者导致图的编译和执行。

此类方法的其他示例包括 .item()、isEqual()。一般来说，任何将张量映射到标量的操作都会导致这种行为。

动态图

如前所述，如果相同形状的图被多次执行，则图编译成本会被分摊。这是因为编译后的图会使用从图形状、输入形状和输出形状派生的哈希值进行缓存。如果这些形状发生变化，就会触发重新编译，而过于频繁的编译会导致训练时间性能下降。

让我们考虑以下示例

def dummy_step(x, y, loss, acc=False):
  z = torch.einsum('bs,st->bt', y, x)
  step_loss = z.sum().view(1,)
  if acc:
    loss = torch.cat((loss, step_loss))
  else:
    loss = step_loss
  xm.mark_step()
  return loss


import time
def measure_time(acc=False):
  exec_times = []
  iter_count = 100
  x = torch.rand((512, 8)).to(dev)
  y = torch.rand((512, 512)).to(dev)
  loss = torch.zeros(1).to(dev)
  for i in range(iter_count):
    tic = time.time()
    loss = dummy_step(x, y, loss, acc=acc)
    toc = time.time()
    exec_times.append(toc - tic)
  return exec_times

dyn = measure_time(acc=True) # acc= True Results in dynamic graph
st = measure_time(acc=False) # Static graph, computation shape, inputs and output shapes don't change

import matplotlib.pyplot as plt
plt.plot(st, label = 'static graph')
plt.plot(dyn, label = 'dynamic graph')
plt.legend()
plt.title('Execution time in seconds')

请注意，静态和动态情况具有相同的计算，但动态图每次都会编译，导致总体运行时间更长。实际上，包含重新编译的训练步骤有时会慢一个数量级或更多。在下一节中，我们将讨论一些用于调试训练性能下降的 PyTorch/XLA 工具。

使用 PyTorch/XLA 分析训练性能

PyTorch/XLA 性能分析包含两个主要组成部分。首先是客户端性能分析（client side profiling）。只需设置环境变量 PT_XLA_DEBUG 为 1 即可开启此功能。客户端性能分析会指出源代码中未 lowering 的操作或设备到主机的数据传输（device-to-host transfer）。客户端性能分析还会报告训练期间是否发生过于频繁的编译。您可以结合这个笔记本中的分析器，探索 PyTorch/XLA 提供的一些指标和计数器。

PyTorch/XLA 分析器提供的第二个组成部分是内联跟踪注解（inline trace annotation）。例如：

import torch_xla.debug.profiler as xp

def train_imagenet():
  print('==> Preparing data..')
  img_dim = get_model_property('img_dim')
  ....
  server = xp.start_server(3294)
  def train_loop_fn(loader, epoch):
    ....
    model.train()
    for step, (data, target) in enumerate(loader):
      with xp.StepTrace('Train_Step', step_num=step):
        ....
        if FLAGS.amp:
        ....
        else:
          with xp.Trace('build_graph'):
            output = model(data)
            loss = loss_fn(output, target)
            loss.backward()
          xm.optimizer_step(optimizer)

注意 start_server API 调用。您在此处使用的端口号与您将用于 tensorboard 分析器的端口号相同，以便查看类似于以下的 op 跟踪（op trace）：

Op 跟踪结合客户端调试功能是调试和优化 PyTorch/XLA 训练性能的一套强大工具。有关分析器使用的更详细说明，建议读者查阅 PyTorch/XLA 性能调试系列博客的第一部分、第二部分和第三部分。

总结

在本文中，我们回顾了 LazyTensor 系统的基础知识。我们以这些基础知识为基础，结合 PyTorch/XLA 来理解训练性能下降的潜在原因。我们讨论了为什么“编译一次并多次执行”有助于在 LazyTensor 系统上获得最佳性能，以及当这个假设被打破时训练为什么会变慢。

我们希望这些见解能帮助 PyTorch 用户在使用 LazyTensor 系统进行创新工作时有所裨益。

致谢

非常感谢我的杰出同事 Jack Cao, Milad Mohammedi, Karl Weinmeister, Rajesh Thallam, Jordan Tottan (Google) 和 Geeta Chauhan (Meta) 给予他们的细致审阅和反馈。感谢来自 Google, Meta 和开源社区的 PyTorch/XLA 扩展开发团队，让 PyTorch 在 TPU 上成为可能。最后，感谢 LazyTensor 论文的作者，他们不仅开发了 LazyTensor，还撰写了这样一篇易于理解的论文。

参考文献

[1] LazyTensor: 将即时执行与领域特定编译器相结合