Run LitGPT benchmarking with a custom Attention implementation priority.

PyTorch has recently changed the priority order of the attention implementations. We would like to benchmark (through benchmark_litgpt.py) the best performance of each compilation backend to ensure we have a clear picture how Thunder, and ThunderFX are doing.

There exists a context manager, which can be used to select a given sdp backend. For the sake of our script, we call:

benchmark.train()

at one point in benchmark_litgpt.py

My idea would be to wrap this call in the following manner:

from torch.nn.attention import SDPBackend, sdpa_kernel
...
with sdpa_kernel([SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH], set_priority=True):
    benchmark.train()

Torch should then follow the order in the list in calling the sdp backends.

In theory it should work, however, the below happens.

🐛 Bug

When applying the above snippet, and running:

python thunder/benchmarks/benchmark_litgpt.py     --max_iters 10     --warmup_iters 5     --model_name Phi-3-mini-4k-instruct     --compile inductor

We receive the error:

benchmark_litgpt.py", line 765, in train
    loss.backward()
  File "/usr/local/lib/python3.12/dist-packages/torch/_tensor.py", line 648, in backward
    torch.autograd.backward(
  File "/usr/local/lib/python3.12/dist-packages/torch/autograd/__init__.py", line 347, in backward
    _engine_run_backward(
  File "/usr/local/lib/python3.12/dist-packages/torch/autograd/graph.py", line 823, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/autograd/function.py", line 307, in apply
    return user_fn(self, *args)
           ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 1958, in backward
    return impl_fn()
           ^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 1944, in impl_fn
    out = CompiledFunction._backward_impl(ctx, all_args)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 2079, in _backward_impl
    out = call_func_at_runtime_with_args(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/utils.py", line 126, in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
                            ^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 755, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_inductor/output_code.py", line 464, in __call__
    return self.current_callable(inputs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_inductor/utils.py", line 2203, in run
    return model(new_inputs)
           ^^^^^^^^^^^^^^^^^
  File "/tmp/torchinductor/cx/ccxhth63u7h6hk5gqnevy66wnvou4lzluuzelpihkdaczuns2cn3.py", line 1156, in call
AssertionError: expected size 32==32, stride 96==393216 at dim=1; expected size 4096==4096, stride 3072==96 at dim=2

Please note the error occuring at the loss.backward(). What I think happens is the forward call runs smoothly, but due to reshaping, the shapes of the expected tensors change, and SDPA fails.

NOTE: The above issue does not happen for eager, so maybe this is just torch.compile issue.

I would like to open a discussion and answer the following questions:

1. Is this change necessary? In my opinion it is necessary, as we want to benchmark the best performance of each backend, and this performance is usually achieved through cuDNN SDPA. When running torch.compile with cuDNN SDPA, we could achieve higher performance than ThunderFX, which is not the case when executed with Flash Attention.
2. Is the implementation correct or should it be applied in a different way?
3. Is the loss reshaping really the problem here?

cc @crcrpar