NVIDIA / NVIDIA/Megatron-LM

[ENHANCEMENT] Support for FusedLinearCrossEntropy Loss for memory savings

Open
#1,738 4 comments 0 reactions 0 assignees Claimed by @Jianbing-D View on GitHub
enhancement module: transformer engine
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Is your feature request related to a problem? Please describe.**
We already support FusedCrossEntropy losses, it would be even better to support fusing the last linear layer and the CrossEntropy layer to ensure even more memory savings. See https://github.com/pytorch/pytorch/issues/124480 for more discussion. There is talk of adding this as a native operator in PyTorch to help facilitate this. Maybe this issue would be better open in TransformerEngine though. This would massively reduce the amount of memory needed for finetuning small models with large vocabularies or large vocabularies period.

**Describe the solution you'd like**
Support fusing the final linear layer with the cross entropy layer in a dtype aware way that avoid unnecessary computation based the precision of the floating point.

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Proposed implementation**
If you have a proposed implementation for the feature state it here or link to a PR.

See the linked Liger kernel or Apple's Cut Cross Entropy kernel.

**Additional context**
Add any other context or screenshots about the feature request here.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.