Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Parallel Cross Entropy using online softmax #1456

Merged
merged 33 commits into from
Feb 26, 2025

Conversation

sanandaraj5597
Copy link
Contributor

Description

This PR implements a parallel cross entropy function using the online technique to calculate softmax. This feature has multiple aspects:

  1. The vocab dimension can be sharded along the TP axis to perform this loss calculation in a distributed fashion.
  2. Online softmax helps us parallelize the softmax calculation giving us more efficiency.
  3. Calculating gradients in the forward itself, so backward step is a no-op.
  4. Storing the gradients in-place of the input tensor saving memory.
  5. OAI Triton implementation helps us integrate GPU kernel level semantics and torch level communication API's together.

[Thanks to Liger kernel implementation for providing the idea about online softmax and in-place gradient calculation.]

Copy link
Collaborator

@timmoon10 timmoon10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add a test in tests/pytorch?

@timmoon10 timmoon10 self-requested a review February 4, 2025 22:38
@timmoon10 timmoon10 added the enhancement New feature or request label Feb 4, 2025
@sanandaraj5597
Copy link
Contributor Author

@timmoon10 Added tests.

Copy link
Member

@ksivaman ksivaman left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@sanandaraj5597 Could you fix the linting errors? You could run it locally using bash qa/L0_pytorch_lint/test.sh

Signed-off-by: Selvaraj Anandaraj <[email protected]>
@sanandaraj5597
Copy link
Contributor Author

Fixed lint errors.

sanandaraj5597 and others added 2 commits February 19, 2025 09:19
Co-authored-by: Kirthi Shankar Sivamani <[email protected]>
Signed-off-by: Selvaraj Anandaraj <[email protected]>
Copy link
Collaborator

@timmoon10 timmoon10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, pending CI

@timmoon10
Copy link
Collaborator

/te-ci pytorch

@timmoon10
Copy link
Collaborator

/te-ci pytorch

Copy link
Member

@ksivaman ksivaman left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@timmoon10
Copy link
Collaborator

/te-ci pytorch

@timmoon10
Copy link
Collaborator

/te-ci pytorch

Selvaraj Anandaraj and others added 4 commits February 20, 2025 20:37
@timmoon10
Copy link
Collaborator

/te-ci pytorch

@timmoon10
Copy link
Collaborator

It seems that the Triton dependency is messing up other tests. The latest upstream Triton (3.2.0) doesn't support Blackwell, so the NVIDIA PyTorch container is using a custom internal build. I've removed Triton as a formal dependency, but we should put it back once Blackwell support is upstreamed.

@timmoon10
Copy link
Collaborator

/te-ci pytorch

@timmoon10 timmoon10 merged commit 8ca2caf into NVIDIA:main Feb 26, 2025
1 of 2 checks passed
@sanandaraj5597 sanandaraj5597 deleted the parallel_cross_entropy branch February 26, 2025 06:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants