Add a transform to remove param and buffer shape check from prologue

[kshitij12345]

Related #1542 #1531

Adds a Transform which removes check_shape_and_metadata from prologue for any proxy tagged with STATIC_MEMORY_LOCATION (which mainly includes module parameters and buffers currently).

Example (have attached the sample trace)

import torch
from thunder.tests.litgpt_model import Config
from litgpt.model import GPT
import thunder
import thunder.transforms
import thunder.transforms.extraction_only_prologue_transform

model = GPT.from_name("llama1-like", n_layer=1)
x = torch.randint(model.max_seq_length, (2, 5))
cmodel = thunder.jit(model, transforms=[thunder.transforms.extraction_only_prologue_transform.ExtractionOnlyPrologueTransform()])
_ = cmodel(x)
prl_trc = thunder.last_prologue_traces(cmodel)[-1]

with open("prologue_trc.py", "w") as f:
    f.write(str(prl_trc))

prologue_with_checks_trc.py.txt
prologue_without_checks_trc.py.txt

Benchmark

import torch
import torch.utils.benchmark
from thunder.tests.litgpt_model import Config
from litgpt.model import GPT
import thunder
import thunder.transforms
import thunder.transforms.extraction_only_prologue_transform

cfg = Config("Llama3-8b")
with torch.device("cuda"):
    model = GPT(cfg)
    x = torch.randint(cfg.vocab_size, (2, 5))

cmodel = thunder.jit(model, transforms=[thunder.transforms.extraction_only_prologue_transform.ExtractionOnlyPrologueTransform()])
# cmodel = thunder.jit(model)
_ = cmodel(x)

comp_and_input_fn = cmodel._lc_cd._get_computation_and_inputs
cd = cmodel._lc_cd
cs = cmodel._lc_cs
cache_entry, inps, pro_to_epi = comp_and_input_fn(x)
prologue_fn = cache_entry.prologue_fn

# Mock the cache_info (required by prologue)
cache_info = {}
cache_info["default_dtype"] = torch.float32
cache_info["default_device"] = torch.device("cpu")
cache_info["is_autocast_enabled"] = False
cache_info["no_grad_sync"] = False
cache_info["alias_tensor_indices"] = ''
cache_info["is_grad_enabled"] = True

thunder._cache_info_ctx.set(cache_info)

with thunder.compile_data_and_stats(cd, cs):
    prologue_fn(x)

timer = torch.utils.benchmark.Timer("prologue_fn(x)", globals={"prologue_fn": prologue_fn, "x": x, "thunder": thunder, "cd": cd, "cs": cs}, setup="thunder.core.compile_data.set_compile_data_and_stats(cd, cs)")
print(timer.blocked_autorange(min_run_time=1))

Without Transform

<torch.utils.benchmark.utils.common.Measurement object at 0x7fd9fda0f940>
prologue_fn(x)
setup: thunder.core.compile_data.set_compile_data_and_stats(cd, cs)
  Median: 308.20 us
  IQR:    2.99 us (305.91 to 308.90)
  4 measurements, 1000 runs per measurement, 1 thread

With Transform

<torch.utils.benchmark.utils.common.Measurement object at 0x73f72e0a6320>
prologue_fn(x)
setup: thunder.core.compile_data.set_compile_data_and_stats(cd, cs)
  Median: 130.09 us
  IQR:    0.35 us (129.99 to 130.34)
  8 measurements, 1000 runs per measurement, 1 thread

[t-vi]

Thanks for working on this, but no to the implementation. We want this as a transform, please.

[t-vi]

Note that #1542 basically duplicates #1531 except that it does not follow the thunder style of using Transforms in the comments regarding implementation strategy.

[kshitij12345]

@t-vi have updated the implementation, have a look whenever you get a chance.

[t-vi]

Seems good (and quite simple), thank you @kshitij12345