Complete torch.compile executor

[carmocca]

What does this PR do?

Adds an instance of TorchCompileExecutor named torch_compile_Ex that registers all of pytorch_executor's and sdpa_ex's operators.

The rename of the existing executor torch_compile_executor to torch_compile_cat_ex breaks backwards compatibility.

Playground script with litgpt:

from litgpt import GPT
import thunder
import torch

from thunder import pytorch_executor
from thunder.executors.torch_compile import torch_compile_ex

with torch.device("cuda"):
    model = GPT.from_name("Llama-2-7b-hf", n_layer=1)

model = thunder.jit(model, executors=[torch_compile_ex])
x = torch.randint(model.max_seq_length, (2, 5), device="cuda")
y = model(x)

forward_trace = thunder.last_traces(model)[-1].python()
print(forward_trace)
assert "TorchCompile" in str(forward_trace)

Output:

@torch.no_grad()
@no_autocast
def augmented_forward_fn(idx, t_lm_head_weight, t_transformer_h_0_attn_attn_weight, t_transformer_h_0_attn_proj_weight, t_transformer_h_0_mlp_fc_1_weight, t_transformer_h_0_mlp_fc_2_weight, t_transformer_h_0_mlp_proj_weight, t_transformer_h_0_norm_1_weight, t_transformer_h_0_norm_2_weight, t_transformer_ln_f_weight, t_transformer_wte_weight, tos1, t_sin):
  # idx: "cuda:0 i64[2, 5]"
  # t_lm_head_weight: "cuda:0 f32[32000, 4096]"
  # t_transformer_h_0_attn_attn_weight: "cuda:0 f32[12288, 4096]"
  # t_transformer_h_0_attn_proj_weight: "cuda:0 f32[4096, 4096]"
  # t_transformer_h_0_mlp_fc_1_weight: "cuda:0 f32[11008, 4096]"
  # t_transformer_h_0_mlp_fc_2_weight: "cuda:0 f32[11008, 4096]"
  # t_transformer_h_0_mlp_proj_weight: "cuda:0 f32[4096, 11008]"
  # t_transformer_h_0_norm_1_weight: "cuda:0 f32[4096]"
  # t_transformer_h_0_norm_2_weight: "cuda:0 f32[4096]"
  # t_transformer_ln_f_weight: "cuda:0 f32[4096]"
  # t_transformer_wte_weight: "cuda:0 f32[32000, 4096]"
  # tos1: "cuda:0 f32[4096, 128]"
  # t_sin: "cuda:0 f32[4096, 128]"
  [t10, t101, t106, t70, t81] = TorchCompile0(idx, t_lm_head_weight, t_sin, t_transformer_h_0_attn_attn_weight, t_transformer_h_0_attn_proj_weight, t_transformer_h_0_mlp_fc_1_weight, t_transformer_h_0_mlp_fc_2_weight, t_transformer_h_0_mlp_proj_weight, t_transformer_h_0_norm_1_weight, t_transformer_h_0_norm_2_weight, t_transformer_ln_f_weight, t_transformer_wte_weight, tos1)
  return {'output': t106, 'flat_args': [idx, t_lm_head_weight, t_transformer_h_0_attn_attn_weight, t_transformer_h_0_attn_proj_weight, t_transformer_h_0_mlp_fc_1_weight, t_transformer_h_0_mlp_fc_2_weight, t_transformer_h_0_mlp_proj_weight, t_transformer_h_0_norm_1_weight, t_transformer_h_0_norm_2_weight, t_transformer_ln_f_weight, t_transformer_wte_weight, tos1, t_sin], 'flat_output': (t106,)}, ((idx, t10, t101, t70, t81, t_lm_head_weight, t_sin, t_transformer_h_0_attn_attn_weight, t_transformer_h_0_attn_proj_weight, t_transformer_h_0_mlp_fc_1_weight, t_transformer_h_0_mlp_fc_2_weight, t_transformer_h_0_mlp_proj_weight, t_transformer_h_0_norm_1_weight, t_transformer_h_0_norm_2_weight, t_transformer_ln_f_weight, t_transformer_wte_weight, tos1), (False, False, 0.29730177875068026, 0.29730177875068026, 4096.0, 4096.0, 4096.0, 32000, 2, -1))

Fixes https://github.com/Lightning-AI/lit-thunder-LEGACY/issues/2141

cc @Borda @apaz-cli

[IvanYashchuk]

Rename the existing torch_compile executor to torch_compile_partial

I suggest renaming the existing executor to "concat_inductor". This is what it does, uses Inductor to fuse concatenation and surrounding operations.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[carmocca]

Are these known flakes on Windows?

FAILED thunder/tests/test_grad.py::test_vjp_correctness_getitem_torch_cpu_float64 - AssertionError: Scalars are not close!

Expected 75.60970520045251 but got 81.13586235571096.
Absolute difference: 5.526157155258446 (up to 1e-05 allowed)
Relative difference: 0.07308793415617461 (up to 1.3e-06 allowed)
FAILED thunder/tests/test_grad.py::test_phantom_grad_vs_torch_consistency_getitem_torch_cpu_bfloat16 - AssertionError: Tensor-likes are not close!

Mismatched elements: 112 / 140 (80.0%)
Greatest absolute difference: 4.0 at index (0, 0, 0) (up to 1e-05 allowed)
Greatest relative difference: 2.0 at index (0, 0, 0) (up to 0.016 allowed)
FAILED thunder/tests/test_grad.py::test_phantom_grad_vs_torch_consistency_getitem_torch_cpu_float64 - AssertionError: Tensor-likes are not close!

Mismatched elements: 112 / 140 (80.0%)
Greatest absolute difference: 7.0 at index (0, 0, 0) (up to 1e-07 allowed)
Greatest relative difference: 3.5 at index (0, 0, 0) (up to 1e-07 allowed)
FAILED thunder/tests/test_grad.py::test_phantom_grad_vs_torch_consistency_getitem_torch_cpu_float32 - AssertionError: Tensor-likes are not close!

Mismatched elements: 112 / 140 (80.0%)
Greatest absolute difference: 5.0 at index (0, 0, 0) (up to 1e-05 allowed)
Greatest relative difference: 2.5 at index (0, 0, 0) (up to 1.3e-06 allowed)

[apaz-cli]

@carmocca I usually assume that consistency tests are flakes, yeah. Re-run, and it should go away. If it doesn't, then it wasn't a flake :)

[t-vi]

LGTM Awesome work. Thank you @carmocca @IvanYashchuk @apaz-cli

[IvanYashchuk]

I wanted to highlight that our nightly scripts and reporting depend on the current naming convention to monitor performance history. There was no real need to modify these benchmark options names in this PR in a breaking way. In the future, I suggest we explore different alternatives for modifications that alter existing behavior before finalizing the merge. Additionally, after merging, it's important to communicate these changes through various channels, not limited to GitHub, to ensure everyone is informed.