Lightning-AI / Lightning-AI/lightning-thunder

Perf is not great on HF Transformers Llama 3.2 1B

Open
#2,655 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

*Note*: If you have a model or program that is not supported yet but should be, please use the program coverage template.

## 🐛 Bug

Perf is not great, on H100. Using the quickstart example.

```
Transformers full options: 482.82ms
Transformers overhead: 649.67ms
transformers plain eager: 1197.29ms
Thunder: 1775.08ms
```

### To reproduce

```python
import torch
import transformers
import transformers.generation
import thunder
import thunder.recipes
import thunder.plugins

from thunder.dev_utils.benchmark import benchmark_n

model_name = "meta-llama/Llama-3.2-1B"

device = "cuda:0" if torch.cuda.is_available() else "cpu"

tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)

with torch.device(device):
model = transformers.AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
model.requires_grad_(False)
model.eval()
# apparently, Transformers 4.51.3 does not instantiate models on the default device
model.to(device)

inp = tokenizer(["Hello world! Here's a long story"], return_tensors="pt")

def generate(model, inp, transformers_compile='full'):
genconf = model._prepare_generation_config(None)[0]
if transformers_compile == 'none': # if you enable this, no torch compile
genconf.disable_compile = True
elif transformers_compile == 'overhead':
# if you enable the line below (but don't disable compile above) you get torch.compile with "default" mode rather than reduce overhead
genconf.compile_config = transformers.generation.CompileConfig(mode="default")
elif transformers_compile != 'full':
raise NotImplementedError(f"unsupported {transformers_compile=}")
out = model.generate(**inp, do_sample=False, generation_config=genconf, cache_implementation="static", max_new_tokens=100)
print(tokenizer.decode(out[0].tolist()))

print("\nGenerating with PyTorch eager:")

transformers_full = benchmark_n(2, generate, model, inp, device=device)
transformers_overhead = benchmark_n(2, generate, model, inp, transformers_compile='overhead', device=device)
transformers_no_compile = benchmark_n(2, generate, model, inp, transformers_compile='none', device=device)

recipe = thunder.recipes.HFTransformers()
thunder_model = thunder.compile(
model,
recipe=recipe,
# plugins=thunder.plugins.ReduceOverhead(), # CUDAGraphs will produce garbage output on main.
)

print("\nGenerating with Thunder:")
thunder_time = benchmark_n(2, generate, thunder_model, inp, device=device)

print(f"Transformers full options: {transformers_full:.2f}ms")
print(f"Transformers overhead: {transformers_overhead:.2f}ms")
print(f"Transformers plain eager: {transformers_no_compile:.2f}ms")
print(f"Thunder: {thunder_time:.2f}ms")
```

To get a first look at what's going on:
```python
with torch.profiler.profile(with_stack=True) as prof:
out = thunder_model.generate(**inp, do_sample=False, cache_implementation="static", max_new_tokens=5)
prof.export_chrome_trace('thunder.json')
print(prof.key_averages().table(sort_by="self_device_time_total"))
```
### Expected behavior

Get within ~20% between thunder and transformers overhead (no cudagraphs).
Ideally get CUDAGraphs-based also within ~20% and working generation.

### Environment

PyTorch 2.8.0 cuda 12.8, nvfuser-cu128-torch2.8, H100

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the quickstart reproduction using meta-llama/Llama-3.2-1B, the thunder.recipes.HFTransformers() compilation path, and thunder_model.generate. Run the profiler snippet and inspect thunder.json and the self-device-time table to compare the generation path with Transformers overhead. Done means Thunder is within about 20% of Transformers overhead without CUDA graphs, with CUDA-graph generation also working ideally.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.