High memory consumption without dataflow-based fusion

[riccardofelluga]

⏱️ 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

[IvanYashchuk]

How do fusions change for mistralai/Mistral-Nemo-Base with and without f2d7152?