Torch.compile seems to hang

4

Apologies for the delay but there are a few issues

  1. AOT eager won’t have substantial speedups relative to inductor
  2. Your model is dynamic so the recompilations are what is giving the appearance of a hang
  3. For some reason dynamic=True has no impact
  4. After it recompiles your model OOMs
  5. Your benchmark suite doesn’t factor in warmup
    6 torch.compile speedups will be most significant on newer GPUs like A100

For now I’m running experiments on an A10G which means I default to inductor which is the compiler with the most significant speedups

For the recompiles I ran TORCH_LOGS="recompiles" python kf_linreg_jax_vs_pt.py and it showed that the culprit was the t variable here https://github.com/probml/dynamax/blob/main/dynamax/linear_gaussian_ssm/demos/kf_linreg_jax_vs_pt.py#L86 if I fix it to 1 then the hang goes away

To handle dynamic models we have an argument to torch.compile(m, dynamic=True) but for some reason I don’t see its impact on your model

Even with t fixed your model is OOM’ing, I was able to make the OOM go away with this line and I’ve opened an issue about this on github since pattern matcher is on by default OOM in fuse_attention inductor pass · Issue #99084 · pytorch/pytorch · GitHub

import torch._inductor.config
torch._inductor.config.pattern_matcher = False

As far as benchmarking is concerned you call torch.compile(m) the compilation only happens on the first inference so for benchmarks you need to remove the first inference time since the assumption is you’ll amortize it over a long enough experiment

Finally I added a torch.set_default_device('cuda') in your script which defaults all tensors to GPU on an A100 and the benchmarks are not anything to write home about quite yet but the speedups are there

sm/demos$ python kf_linreg_jax_vs_pt.py 
torch, time=0.604 compile False N 100 D 500
[2023-04-13 21:36:10,010] torch._inductor.utils: [WARNING] make_fallback(aten.cumprod): a decomposition exists, we should switch to it
torch, time=4.884 compile True N 100 D 500
torch, time=0.552 compile True N 100 D 500