Eager Mode + Compile API¶
In this doc we will go over how to use PyTorch/XLA’s new experimental eager
mode with the compile
API. The goal is to make PyTorch/XLA experience more aligned with the native PyTorch and make development process easier.
Background¶
Currently PyTorch/XLA runs on the LazyTensor tracing mode by default. In the following code
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()
The actual model compilation and device execution happens when torch_xla.sync
is called. There are multiple drawback of this approach.
Users are often confused about when the framework is tracing and when the framework is executing.
Non-core model code(data preprocessing for example) often generates some small pending execution that gets leaked into the main graph(step function) and causes recompilation. The recompilation of the whole graph is usually very expensive.
It is hard to debug when/why recompilation happens.
To mitigate above issues we want to introduce the new UX with eager and compile.
Basic Usage¶
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.experimental.compile(model)
input = torch.randn(64, 3, 224, 224).to(device)
# Compilation and execution happens right away.
res = compiled_model(input)
Note that
Currently user has to manually enable the eager mode by
torch_xla.experimental.eager_mode(True)
.The region of the code that wants to be compiled should be wrapped by
torch_xla.experimental.compile
.
The implementation of the torch_xla.experimental.compile
is actually pretty straight forward, it disable the eager mode when entering the target function and start tracing. It will call the torch_xla.sync()
when target function returns and reenable the eager mode. You can expect the same perfomrance by using the eager
+ compile
API compared to the existing mark_step/sync
approach.
Inference¶
torch_xla.experimental.eager_mode(True)
compiled_model = torch.compile(model, backend="openxla")
It is recommened to use the torch.compile
instead of torch_xla.experimental.compile
for inference to reduce the tracing overhad.
Training¶
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.experimental.compile(step_fn)
In training we asked user to refactor the step_fn
out because it is usually better to compile the model’s forward, backward and optimizer together. The long term goal is to also use torch.compile
for training but right now we recommend user to use torch_xla.experimental.compile
(for perfomrance reason).
Benchmark¶
I run a 2 layer decoder only model training(it is pretty much just a llama2) with fake data on a single chip of v4-8 for 300 steps. Below is the number I observed.
token/s | |
Tracing mode(base line) | 147 |
Eager mode | 65 |
Eager + torch_xla compile | 147 |
Eager mode can achieve ~45% performance of the fully compiled model for the decoder only model. The trainer I used to test can be found here and here. Note that perfomrane of the eager mode is very model dependent. When I tried to run the resnet50, the eager mode perfomrance is ~1% of the compiled mode. We don’t exepct user to use eager mode to execute the main training loop. Eager mode is meant to be used to handle non-core part of the training/inference logic(Data preprocessing, random number generations etc) or debug.