[scan] Avoid re-tracing the combine function on every call

[tengyifei]

🚀 Feature

It should be possible to somehow cache the traced graphs in torch_xla.experimental.scan so we don't trace on every call.

Motivation

Today torch_xla.experimental.scan and scan_layers traces the user function with both AOTAutograd (to get the backward) and with LazyTensor (to lower them to HLO). AOTAutograd is very slow and we can easily become tracing bound. For example, python3 examples/train_decoder_only_base.py takes 1min30s but python3 examples/train_decoder_only_base.py scan.decoder_with_scan.DecoderWithScan takes 4min.

Pitch

We could wait for torch.scan to support autograd (c.f. #7901 (comment)) which will take a long time. In the meantime, we can implement some simple caching based on the id of the input function/module.

The caching should be opt-in because it's only sound if the function is pure. We can add a assume_pure=True argument to scan so that it only uses the caching when the user confirms that their function is pure.

[tengyifei]

A prototype is at master...flash-attention-aot-autograd#diff-540e30d11e43afaaff4e4f89efc8b6ee396b222b8748d9eac70327ba0a7e0517R547

[miladm]

This seems like a feature that's not limited to making scan better - i.e. it can make any code running through the same torch code path batter; correct?

Isn't this related to the cache we talked about with Rochard Zou earlier today @tengyifei?

[tengyifei]

Isn't this related to the cache we talked about with Rochard Zou earlier today

Yes. But if we tried to implement the general case, IMO we'll end up reimplementing the guard system in Dynamo.

[iwknow]

Hi @tengyifei, Are you working on this? if not, i can take a look at this issue

[tengyifei]

Hi @iwknow , thanks for all your contributions. It really helps a lot.

Yes, this one is free for taking too. I linked a somewhat messy prototype in a previous comment thread. I think that one might have some gaps. The idea is that we could cache the XLA computation object based on the shape of the inputs and the function.

I'm working on a xb.call_jax feature that has a similar caching mechanism. It should demonstrate the idea more clearly.

A benefit of this is that once scan supports assume_pure=True, we can directly use that in GRU, because we know that our implementation is pure!

[iwknow]

can you please provide the links for the related xb.call_jax feature so i can keep as a reference? thanks!

[tengyifei]

@iwknow Hi yes, here's it is:

xla/torch_xla/core/xla_builder.py

Line 981 in 4eea5e1

def call_jax(jax_func,

call_jax calls jax_func_to_xla_computation. jax_func_to_xla_computation calls _jax_to_xla_computation_cache_get_or_insert. And _jax_to_xla_computation_cache_get_or_insert calls get_xla_computation when it realizes that it needs to add a new entry to the cache.

This call_jax feature is for calling JAX from PyTorch/XLA. As such, its hashing is specific to JAX functions. For scan, we could add a assume_pure=True kwargs and optionally opt-in to a similar caching. In that case, I think scan could:

• Use a 2-level look up similar to _jax_to_xla_computation_cache_get_or_insert
• Use the fn from the user as the key into the first level dict
• Use the shapes and dtypes of the inputs, as well as their PyTree spec, as the key into the second level dict

As to what/how to cache the XlaComputation object, I think my prototype in master...flash-attention-aot-autograd#diff-540e30d11e43afaaff4e4f89efc8b6ee396b222b8748d9eac70327ba0a7e0517R547 roughly got it right (search for CachingLoweringContext in the collapsed scan.py file), but the hashing strategy in that thing was wrong and should be changed to follow the call_jax style.

We may need to stringify the PyTree spec because the PyTorch PyTree spec is not hashable, unlike JAX.

We could write similar unit tests following

xla/test/test_jax_interop.py

Line 167 in 4eea5e1

def test_call_jax_cache_hlo(self):

Thanks!