Eager 模式 + Compile API

在本文档中，我们将介绍如何使用 PyTorch/XLA 新的实验性 eager 模式和 compile API。目标是使 PyTorch/XLA 体验更贴近原生 PyTorch，并使开发过程更轻松。

目前，PyTorch/XLA 默认在 LazyTensor 追踪模式下运行。在以下代码中

import torch
import torch_xla
import torchvision

device = torch_xla.device()
model = torchvision.models.resnet18().to(device)
input = torch.randn(64, 3, 224, 224).to(device)

# model tracing
res = model(input)

# model execution, same as `xm.mark_step`
torch_xla.sync()

实际的模型编译和设备执行发生在调用 torch_xla.sync 时。这种方法有多个缺点。

1. 用户常常对框架何时追踪以及框架何时执行感到困惑。

2. 非核心模型代码（例如数据预处理）通常会生成一些小的待处理执行，这些执行会泄漏到主图（步进函数）中，并导致重新编译。整个图的重新编译通常非常昂贵。

3. 很难调试何时/为何发生重新编译。

为了缓解上述问题，我们希望引入具有 eager 和 compile 的新 UX。

基本用法

import torch
import torch_xla
import torchvision

# Run ops eagerly by default
torch_xla.experimental.eager_mode(True)

device = torch_xla.device()
model = torchvision.models.resnet18().to(device)

# Mark the function to be compiled
compiled_model = torch_xla.compile(model)
input = torch.randn(64, 3, 224, 224).to(device)

# Compilation and execution happens right away.
res = compiled_model(input)

请注意：

1. 目前，用户必须通过 torch_xla.experimental.eager_mode(True) 手动启用 eager 模式。

2. 想要编译的代码区域应由 torch_xla.compile 包裹。

torch_xla.compile 的实现实际上非常直接，它在进入目标函数时禁用 eager 模式并开始追踪。当目标函数返回时，它将调用 torch_xla.sync() 并重新启用 eager 模式。与现有的 mark_step/sync 方法相比，使用 eager + compile API 可以获得相同的性能。

推理

torch_xla.experimental.eager_mode(True)
compiled_model = torch.compile(model, backend="openxla")

建议在推理时使用 torch.compile 而不是 torch_xla.compile，以减少追踪开销。

训练

torch_xla.experimental.eager_mode(True)

def step_fn(model, data, target, loss_fn, optimizer):
    optimizer.zero_grad()
    logits = model(data)
    loss = loss_fn(logits, target)
    loss.backward()
    optimizer.step()
    return loss

step_fn = torch_xla.compile(step_fn)

在训练中，我们要求用户重构 step_fn，因为通常最好将模型的前向传播、反向传播和优化器一起编译。长期目标是也使用 torch.compile 进行训练，但目前我们建议用户使用 torch_xla.compile（出于性能原因）。

基准测试

我在 v4-8 单芯片上使用假数据运行了一个 2 层仅解码器模型训练（它几乎就是一个 llama2），进行了 300 步。以下是我观察到的数字。

模式 token/秒

---

追踪模式（基线） 147 Eager 模式 65 Eager + torch_xla compile 147

: Eager 模式基准

对于仅解码器模型，Eager 模式可以达到完全编译模型约 45% 的性能。有关更多信息，请参阅 train_decoder_only_base.py 和 eager example。请注意，eager 模式的性能非常依赖于模型。当我尝试运行 resnet50 时，eager 模式的性能约为编译模式的 1%。我们不期望用户使用 eager 模式来执行主训练循环。Eager 模式旨在用于处理训练/推理逻辑的非核心部分（数据预处理、随机数生成等）或调试。