Lightning-AI / Lightning-AI/lightning-thunder
Low perf for backward pass of dropout fused with normalization
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
The fusion definition generated by Thunder for the backward pass of dropout fused with normalization recomputes the dropout mask, resulting in a 1.4× performance degradation compared to a [previous version](https://github.com/NVIDIA/Fuser/blob/df4e918730d77a7f88c1bb3b24625cb2445f58a0/benchmarks/python/test_dropout_rmsnorm_bwd.py#L20) that reused saved masks
To reproduce, use nvFuser main branch run test `test_dropout_rmsnorm_bwd.py` with:
```
pytest /opt/pytorch/nvfuser/benchmarks/python/test_dropout_rmsnorm_bwd.py -vvvs -m 'not skip' -k "float16 and 16384_4096" --benchmark-thunder
```
On B200, output should be:
```
-------------------------------------------------------------------------------------------------------------------------------- benchmark: 2 tests -------------------------------------------------------------------------------------------------------------------------------
Name (time in us) Min Max Mean StdDev Median IQR Outliers OPS (Kops/s) Rounds Iterations
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
test_dropout_rmsnorm_bwd_nvf_benchmark[dtype=torch.bfloat16-size=[16384_4096]] 151.7090 (1.0) 155.0710 (1.0) 152.6755 (1.0) 1.0656 (6.80) 152.1580 (1.0) 1.2790 (5.63) 2;0 6.5498 (1.0) 10 1
test_dropout_rmsnorm_bwd_baseline_benchmark[dtype=torch.bfloat16-size=[16384_4096]-executor='thunder'] 209.3080 (1.38) 209.7900 (1.35) 209.5808 (1.37) 0.1567 (1.0) 209.6120 (1.38) 0.2270 (1.0) 3;0 4.7714 (0.73) 10 1
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
```
`test_dropout_rmsnorm_bwd_nvf_benchmark` uses saved mask instead of computing dropout achieving about 63% SOL
`test_dropout_rmsnorm_bwd_baseline_benchmark` still computes dropout achieving about 46% SOL
cc @tfogal @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.
Research direction
Start with benchmarks/python/test_dropout_rmsnorm_bwd.py and run the reported pytest command for the float16, 16384_4096 case. Compare test_dropout_rmsnorm_bwd_nvf_benchmark with test_dropout_rmsnorm_bwd_baseline_benchmark and inspect the Thunder-generated backward fusion definition. Done means the backward path reuses the saved mask and approaches the faster benchmark result rather than recomputing dropout.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100