Pytorch XLA Optimization Guide

Training large-scale deep learning models on TPUs can be a game-changer, but it often comes with its own set of challenges. One of the most common hurdles is achieving optimal performance with PyTorch/XLA. In this blog post, we’ll walk through a real-world debugging session to uncover the secrets of PyTorch/XLA performance optimization.

The Problem: Graph Recompilation

The key to getting good performance with PyTorch/XLA is to have a static computational graph. This means that the structure of the graph (the layers and their connections) should not change between training steps. When the graph is static, XLA can compile and optimize it once, and then reuse the compiled graph for subsequent steps. This is where the real speed-up comes from.

However, if the graph changes between steps, XLA has to recompile it, which is a very expensive operation. This is known as graph recompilation, and it’s a major performance bottleneck.

The Detective Work: Identifying Recompilations

So, how do you know if your graph is being recompiled? PyTorch/XLA provides some excellent tools for this.

Metrics Report

The first tool in our arsenal is the metrics report. You can get a high-level overview of the performance by printing the metrics report at each step:

import torch_xla.debug.metrics as met

print(met.metrics_report())

The key metric to look for is UncachedCompile. If this counter is greater than 1 after the first few steps, it’s a clear sign that your graph is being recompiled.

torch._dynamo.explain

Once you’ve confirmed that your graph is being recompiled, the next step is to pinpoint the exact source of the recompilation. This is where torch._dynamo.explain comes in. This powerful tool can give you a detailed explanation of why and where the graph is being broken.

import torch._dynamo as dynamo

explanation = dynamo.explain(model, input_ids, src_key_padding_mask)
print(explanation)

The output of explain will show you the number of graph breaks and the reasons for them. This is the information you need to identify the source of the dynamism in your code.

The Solutions: A Checklist for Static Graphs

Once you’ve identified the source of the recompilations, the next step is to fix it. Here’s a checklist of common strategies for making your graph static:

• Separate Training and Evaluation: Avoid alternating between model.train() and model.eval() in the same training loop. Perform all the training epochs first, and then perform a final evaluation.

• Implement Sequence Bucketing: If you’re working with variable-length sequences, implement a form of bucketing in your data collator. This will reduce the number of unique sequence lengths and thus the number of recompilations.

• Add a Warm-up Phase: Add a warm-up phase to your training script to pre-compile the graphs for the different bucket sizes before the actual training starts.

• Avoid Data-Dependent Control Flow: Avoid using if statements or loops that depend on the values of tensors.

• Use Supported Operations: Make sure you’re using PyTorch operations that are supported by the openxla backend.

Conclusion

Optimizing PyTorch/XLA performance can be a bit of a detective story, but with the right tools and strategies, you can unlock the full power of your TPUs. By understanding the importance of static graphs and by using tools like the metrics report and torch._dynamo.explain, you can identify and fix the performance bottlenecks in your code and get your models training at lightning speed.