Fixes #8805.

We introduce a decorator, @assume_pure, that can be placed on PyTorch/XLA functions and easily eliminate lazy tensor tracing overhead. If you have a pure function that only uses torch upstream ops, that function can be decorated with @assume_pure and will only be traced once for each unique input tensor shape combinations.

Design

@assume_pure brings together three pieces of existing technologies:

• jax.vjp, which takes a JAX function and gives you the autograd forward and backward pass
• torchax, which converts a pure PyTorch function to a JAX function
• xb.call_jax, which can call any JAX function from PyTorch/XLA and integrate it into the HLO graph

It works by:

• Use torchax.interop.jax_view to obtain a JAX function from the input PyTorch function
• Use jax.vjp to get the forward and backward pass
• Return a torch.autograd.Function instance, where the forward implementation is xb.call_jax(forward_pass), and the backward implementation is xb.call_jax(backward_pass), respectively.

The core logic is actually just a single line:

def assume_pure(fn):
  from torchax.interop import jax_view
  return j2t_autograd(jax_view(fn))

How is the HLO cached

xb.call_jax caches the HLO if all the input shapes/dtypes and non-tensor arguments are the same.

Therefore, subsequent xb.call_jax will just reuse the cached HLO instead of retracing.

The same kind of caching happens in both the forward and backward pass.

Different from the jax wrapper we used in splash_attention, the j2t_autograd function saves the residuals (intermediate activations) during the forward pass and reuses them during the backward by plugging those into the vjp_fun again. This means it won't force a rematerialization (rerun the fwd) during the backward.

Alternatives

Instead of jax.vjp we could also use AOTAutograd to get the forward and backward pass. However, AOTAutograd has a number of downsides:

• It does more than just getting the backward. It also forcefully decomposes all operations into the "aten" op set. Decomposing operations will negatively impact performance, especially in the case of einsum.
• There is no straightforward path to support profiler trace spans. In contrast, in the proposed approach we could translate xp.Trace(...) to jax.named_scope(...).
• Supporting custom operations such as pallas kernels will be cumbersome. We'll need to wrap every kernel into a PyTorch custom operator in order for AOTAutograd to not crash on those functions. In contrast, in the proposed approach we could augment our pallas kernels to directly jump into JAX when the input tensor is a torchax tensor.

Instead of assume_pure, we could also use torch.compile to cache the XLA executable of the compiled function and skip the lazy tensor tracing. However, torch.compile has its own downsides:

• torch.compile itself uses AOTAutograd and will suffer from the decomposition and customer operations issues etc.
• torch.compile has a general perception of "either it works, or debugging will be complicated", which has been corroborated by experiments by people in the PyTorch/XLA team. See PyTorch team members' own recommendation 1. In contrast, @assume_pure has very simple rules for determining if it will work: if your function is pure, then it works.
• torch.compile will graph break when entering and leaving the compiled region. In contrast, @assume_pure can avoid tracing overhead without even breaking the graph. The cached HLO is inlined into the overall HLO.

Benchmarks

I tested tracing an example 100 layer decoder-only model:

~/pytorch/xla
❯ TESTBRIDGE_TEST_ONLY=test_trace_transformer_with_spda_attention python3 test/test_assume_pure.py --benchmark_iterations 100
[...]
No `@assume_pure` time: 140.1342 ms
`@assume_pure` time: 24.1658 ms

Importantly, the @assume_pure does not scale with increasing complexity inside the model. That's because we only trace the model once, paying a fixed up-front cost, and then later runs will reuse the cached XLA computation object.

Anecdotally, @bhavya01 reported saving >200ms tracing time in an SDXL experiment. That's very significant since each training step is sub-1 second.