Skip to content

[PyTorch] Use same API in optimizer zero_grad as PyTorch optimizers#1466

Merged
timmoon10 merged 2 commits into
NVIDIA:mainfrom
timmoon10:optim-set-to-none
Feb 22, 2025
Merged

[PyTorch] Use same API in optimizer zero_grad as PyTorch optimizers#1466
timmoon10 merged 2 commits into
NVIDIA:mainfrom
timmoon10:optim-set-to-none

Conversation

@timmoon10
Copy link
Copy Markdown
Member

@timmoon10 timmoon10 commented Feb 8, 2025

Description

The fused optimizers (copied from Apex in #867) have a non-standard API that makes it annoying to swap with the vanilla PyTorch optimizers. This PR deprecates the set_grad_none kwarg in FusedAdam and FusedSGD in favor of a set_to_none kwarg in zero_grad, similar to the vanilla PyTorch API. I've also modified some of the constructor kwargs to be more consistent with PyTorch.

Closes #1453.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

  • Deprecate set_grad_none kwarg in FusedAdam and FusedSGD in favor of set_to_none kwarg in zero_grad
  • Reorder kwargs in FusedAdam and FusedSGD constructors

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@timmoon10
Use same API in optimizer zero_grad as PyT optimizers
b25612b 
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10 timmoon10 added the enhancement New feature or request label Feb 8, 2025
@timmoon10
Copy link
Copy Markdown
Member Author

timmoon10 commented Feb 8, 2025

/te-ci pytorch

@ptrendx ptrendx added the 2.1.0 label Feb 15, 2025
ptrendx
ptrendx reviewed Feb 20, 2025
View reviewed changes
Comment thread transformer_engine/pytorch/optimizers/fused_sgd.py
weight_decay: float = 0.0,
nesterov: bool = False,
*,
wd_after_momentum=False,
Copy link
Copy Markdown
Member

@ptrendx ptrendx Feb 20, 2025

Choose a reason for hiding this comment

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

Why not add the type info to those as well?

Copy link
Copy Markdown
Member Author

@timmoon10 timmoon10 Feb 21, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

I focused on matching the PyTorch API: https://pytorch.org/docs/stable/generated/torch.optim.SGD.html#torch.optim.SGD
I considered other options out of scope.

@timmoon10
Copy link
Copy Markdown
Member Author

timmoon10 commented Feb 21, 2025

/te-ci pytorch

@timmoon10 timmoon10 merged commit b4fbc2b into NVIDIA:main Feb 22, 2025
@timmoon10 timmoon10 deleted the optim-set-to-none branch February 22, 2025 02:13
timmoon10 added a commit that referenced this pull request Feb 26, 2025
@timmoon10
[PyTorch] Use same API in optimizer zero_grad as PyTorch optimizers (…
1b384b9 
…#1466)

Use same API in optimizer zero_grad as PyT optimizers

Signed-off-by: Tim Moon <tmoon@nvidia.com>
@hungryGeek16 hungryGeek16 mentioned this pull request May 31, 2026
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ptrendx ptrendx ptrendx approved these changes

Assignees

No one assigned

Labels

2.1.0 enhancement New feature or request

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

FusedAdam optimizer doesn't have set_to_none keyword argument

2 participants

@timmoon10 @ptrendx