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.

Sign up for GitHub

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

Jump to bottom

Decouple optimizer state and grad dtypes in distributed Adam optimizer #1575

Merged
merged 4 commits into from
Feb 25, 2023
Merged

Decouple optimizer state and grad dtypes in distributed Adam optimizer #1575

merged 4 commits into from
Feb 25, 2023

Conversation

timmoon10
Copy link
Contributor

@timmoon10 timmoon10 commented Feb 4, 2023 •
edited
Loading

This PR is intended to add support for BF16 grad reductions in NeMo-Megatron. Main changes:

  • Generalizes the supported dtypes for the distributed Adam optimizer. It can now handle FP32, FP16, or BF16 gradients, independent of the dtypes for the optimizer state or params.
  • The Transformer layer automatically detects the grad dtype for wgrad fusion. It can now handle the case where the main_grad tensor is FP32, FP16, or BF16.

@timmoon10 timmoon10 mentioned this pull request Feb 9, 2023
Merged
8 tasks
@timmoon10 timmoon10 force-pushed the dist-adam-bf16-grad-reductions branch from 993d399 to 62847a1 Compare February 9, 2023 02:42
@timmoon10 timmoon10 marked this pull request as ready for review February 9, 2023 02:49
crcrpar
crcrpar reviewed Feb 24, 2023
View reviewed changes
Copy link
Collaborator

@crcrpar crcrpar left a comment

Choose a reason for hiding this comment

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

it's totally optional to access all the comments

apex/contrib/csrc/optimizers/multi_tensor_distopt_adam_kernel.cu Outdated Show resolved Hide resolved
apex/contrib/csrc/optimizers/multi_tensor_distopt_adam_kernel.cu Outdated Show resolved Hide resolved
apex/transformer/tensor_parallel/layers.py Outdated Show resolved Hide resolved
apex/transformer/tensor_parallel/layers.py Outdated Show resolved Hide resolved
@crcrpar crcrpar added this to the 23.03 milestone Feb 25, 2023
@crcrpar crcrpar merged commit 03c9d80 into NVIDIA:master Feb 25, 2023
This was referenced Mar 27, 2023
Merged
Merged
yuanzhedong pushed a commit to yuanzhedong/apex that referenced this pull request Jul 14, 2023
@timmoon10
Decouple optimizer state and grad dtypes in distributed Adam optimizer (
f94cc04 
NVIDIA#1575)

* Decouple distopt dtypes for grads and optim state

* Automatically detect grad dtype for Transformer layer wgrad fusion

* Review suggestions from @crcrpar
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@crcrpar crcrpar crcrpar left review comments

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
23.03
Development

Successfully merging this pull request may close these issues.
2 participants
@timmoon10 @crcrpar