Lightning-AI / Lightning-AI/lightning-thunder
Perf is not great on HF Transformers Llama 3.2 1B
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
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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