Add a benchmark for portions of LitGPT model other than SDPA

[IvanYashchuk]

This PR adds a new benchmark, run it with

pytest thunder/benchmarks/litgpt_chunks.py --benchmark-group-by='group,param:info' --benchmark-columns='min,max,mean,stddev,median'

The intent is to be able to compare performance on sections of the GPT network that are not covered by a FlashAttention kernel.

Constructing benchmark cases is slow and -s gives a progress bar:

Constructing benchmark cases for config: CodeLlama-13b-hf (1/39)
Constructing benchmark cases for config: CodeLlama-34b-hf (2/39)
Constructing benchmark cases for config: CodeLlama-70b-hf (3/39)
Constructing benchmark cases for config: CodeLlama-7b-hf (4/39)
Constructing benchmark cases for config: Gemma-2b-it (5/39)
Constructing benchmark cases for config: Gemma-7b-it (6/39)
Constructing benchmark cases for config: Mistral-7B-v0.1 (7/39)
Constructing benchmark cases for config: Nous-Hermes-13b (8/39)
Constructing benchmark cases for config: Nous-Hermes-Llama2-13b (9/39)
Constructing benchmark cases for config: Platypus2-70B (10/39)
Constructing benchmark cases for config: Platypus2-70B-instruct (11/39)
Constructing benchmark cases for config: RedPajama-INCITE-Instruct-3B-v1 (12/39)
Constructing benchmark cases for config: dolly-v2-12b (13/39)
Constructing benchmark cases for config: dolly-v2-3b (14/39)
Constructing benchmark cases for config: dolly-v2-7b (15/39)
Constructing benchmark cases for config: falcon-180B-chat (16/39)
Constructing benchmark cases for config: falcon-40b-instruct (17/39)
Constructing benchmark cases for config: falcon-7b-instruct (18/39)
Constructing benchmark cases for config: open_llama_3b (19/39)
Constructing benchmark cases for config: phi-1_5 (20/39)
Constructing benchmark cases for config: phi-2 (21/39)
Constructing benchmark cases for config: pythia-1.4b-deduped (22/39)
Constructing benchmark cases for config: pythia-12b-deduped (23/39)
Constructing benchmark cases for config: pythia-14m (24/39)
Constructing benchmark cases for config: pythia-160m-deduped (25/39)
Constructing benchmark cases for config: pythia-1b-deduped (26/39)
Constructing benchmark cases for config: pythia-2.8b-deduped (27/39)
Constructing benchmark cases for config: pythia-31m (28/39)
Constructing benchmark cases for config: pythia-410m-deduped (29/39)
Constructing benchmark cases for config: pythia-6.9b-deduped (30/39)
Constructing benchmark cases for config: pythia-70m-deduped (31/39)
Constructing benchmark cases for config: stablecode-instruct-alpha-3b (32/39)
Constructing benchmark cases for config: stablelm-tuned-alpha-3b (33/39)
Constructing benchmark cases for config: stablelm-tuned-alpha-7b (34/39)
Constructing benchmark cases for config: stablelm-zephyr-3b (35/39)
Constructing benchmark cases for config: tiny-llama-1.1b-chat (36/39)
Constructing benchmark cases for config: vicuna-13b-v1.5-16k (37/39)
Constructing benchmark cases for config: vicuna-33b-v1.3 (38/39)
Constructing benchmark cases for config: vicuna-7b-v1.5-16k (39/39)

and it takes 6 minutes to generate the test cases:

468 tests collected in 377.53s (0:06:17)

It's possible to control the batch size by modifying BATCH_SIZE in litgpt_chunks.py and which configs to benchmark can be controlled by modifying the CONFIG_NAMES list.

Thunder is used for tracing the litgpt code and then the trace is split into chunks with the SDPA call as a delimiter. Since the GPT model has a for-loop structure it's enough to trace a model with just two transformer blocks. It gives the following program chunks:

• Embedding (https://github.com/Lightning-AI/litgpt/blob/78bd4cae1e655359e92e7c7b830fcbaa4c15c152/litgpt/model.py#L89-L91) -> Beginning of a transformer block until SDPA
• From one SDPA call to another (so merging the end and start of a transformer block)
• End of a transformer block -> ln_f+lm_head (https://github.com/Lightning-AI/litgpt/blob/78bd4cae1e655359e92e7c7b830fcbaa4c15c152/litgpt/model.py#L95-L96)

We could save the result of generating test cases to a disk, but that's left as an exercise for the future.

TODO:

• Run the benchmark on H100

cc @crcrpar @kevinstephano

[IvanYashchuk]

Uncommenting this would force generating benchmark cases just for this Llama 2 7B config.

[IvanYashchuk]

• Maybe "ATEN" should be removed here forcing Inductor to use cutlass or triton for gemms. I should try if it works without any errors.

[riccardofelluga]

Is this still needed?

[IvanYashchuk]

List comprehensions are easier to read for me than the for-loop below.
I'll remove this of course.

[riccardofelluga]

Looks good to me! I've played a bit with the make_torch_traces_for_config and it does the job as long as the part that we are interested in comes after sdpa.

[lantiga]

Hey @IvanYashchuk should we revive this or close for now? We can add a label for PRs we close that could be potentially of interest for the future.

[IvanYashchuk]

I've put it to draft to prevent merging because I need more time to think about it and convince myself again that it's something we need in the project.
I prefer to keep it in the draft stage to remind myself about it every day.

[lantiga]

"draft" is the new browser tab haha