Add custom decompositions for cross entropy loss for the nvfuser executor #2043
Conversation
protonu
requested review from
naoyam,
jjsjann123,
kevinstephano,
kshitij12345,
IvanYashchuk and
mruberry
for more information, see https://pre-commit.ci
protonu
force-pushed
the
pbasu_loss_fwd_bwd
branch
from
a60395a to
b045ccb
Compare
Co-authored-by: Masaki Kozuki <mkozuki@nvidia.com>
There was a problem hiding this comment.
just wondering if you have some numbers of cross entropy loss to compare this with the existing ones e.g. https://github.com/Lightning-AI/lightning-thunder/blob/c6928015914fdbdd708fd8e87fbd9d9c1b4a40ef/thunder/executors/triton_crossentropy.py?
@crcrpar I did compare performance against torchcompile (which uses Triton - but is that the same as the link you sent?) This is a benchmark I was going to add: |
|
@IvanYashchuk , @jjsjann123 would you like to review again? |
|
@beverlylytle Would you like to take a look, too? |
Co-authored-by: beverlylytle <57254617+beverlylytle@users.noreply.github.com>
Co-authored-by: beverlylytle <57254617+beverlylytle@users.noreply.github.com>
|
@IvanYashchuk did you want to take a last look? |
…utor (#2043) This PR adds custom decompositions for Cross-Entropy Loss for the nvFuser executor. Adding these custom decompositions improves performance and allows further optimization in nvFuser. For cross-entropy loss forward: 1. We move the take_along_axis computation before the log softmax is computed. This allows us to reduce memory traffic for the inputs. For cross-entropy loss backward: 1. We replace a scatter-op with a iota and where op as we don't have support for scatter exposed in nvFuser. 2. We can get rid of a reduction that shows up when backward is computed as nll_loss backward followed by log softmax backward.
This PR adds custom decompositions for Cross-Entropy Loss for the nvFuser executor.
Adding these custom decompositions improves performance and allows further optimization in nvFuser.
For cross-entropy loss forward:
For cross-entropy loss backward:
cc @tfogal