Lightning-AI / Lightning-AI/lightning-thunder
High memory consumption without dataflow-based fusion
@riccardofelluga is already working on this.
Since Feb 11, 2025.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
⏱️ Perf regression
Disabling dataflow-based fusion logic introduced a substantial performance regression. In the models we are tracking, we can see an increase in memory usage.
To Reproduce
There are a few steps to follow:
pip install peft
pip uninstall bitandbytes
pytest thunder/benchmarks/targets.py::test_hf_transformers[mistralai/Mistral-Nemo-Base-2407-BS1-4096-PEFT-forward-thunderfx] --benchmark-json "out.json"
cat out.json | grep max_allocated_memory_MB
git revert f2d715240555ca787fbdbba6f42c6f9d422ae0c3
pytest thunder/benchmarks/targets.py::test_hf_transformers[mistralai/Mistral-Nemo-Base-2407-BS1-4096-PEFT-forward-thunderfx] --benchmark-json "out.json"
cat out.json | grep max_allocated_memory_MB
This will print the peak memory consumption for both runs, on H100 this prints ~71GB and ~67GB after reverting.
Additional info
With this issue, I am not looking to revert the commit, but to track the process of bringing back the dataflow fusion and taking advantage of the moment to brainstorm a bit on it and improve it from the state that it was before being deleted.
Mentioning PR for tracking #1763
cc @riccardofelluga
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.
Assessment
This issue has not been assessed yet.