docs: update mcore optimizer docstrings to google style by Akshat8510 · Pull Request #2799 · NVIDIA/Megatron-LM

Akshat8510 · 2026-01-04T04:41:37Z

Description

This PR updates and standardizes the docstrings for the megatron/core/optimizer module to follow the Google Style Guide. This is part of the documentation overhaul tracked in issue #2653.

Changes

Updated optimizer.py, distrib_optimizer.py, grad_scaler.py, and clip_grads.py with structured Args, Returns, and Attributes blocks.
Formatted OptimizerConfig dataclass attributes for better compatibility with auto-generated API guides.
Refined technical descriptions and updated documentation references to use relative paths.

Files Updated:

megatron/core/optimizer/clip_grads.py
megatron/core/optimizer/distrib_optimizer.py
megatron/core/optimizer/grad_scaler.py
megatron/core/optimizer/optimizer.py
megatron/core/optimizer/optimizer_config.py
megatron/core/optimizer/qk_clip.py

Mentions: @Phlip79 @sbhavani

copy-pr-bot · 2026-01-04T04:41:41Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Akshat8510 · 2026-01-04T04:43:17Z

Hi @Phlip79, I have submitted a PR for the optimizer and clip_grads modules as part of this docstring update.

Phlip79 · 2026-01-09T00:58:40Z

Thank you Akshat! Could you please ensure that your changes do not reduce verbosity?

Akshat8510 · 2026-01-09T04:27:50Z

Thank you for the feedback, @Phlip79! I will do a thorough pass over all the changed files to ensure that all technical details and explanations from the original comments are preserved within the new Google Style format. I'll update the PR shortly.

Akshat8510 · 2026-01-09T05:09:35Z

Hi @Phlip79, I have updated the docstrings across all 6 files to restore the original technical verbosity and detailed explanations while maintaining the Google Style structure. This includes restoring the core mapping logic in distrib_optimizer.py, the FSDP/filtering details in clip_grads.py, and the in-place modification warnings. Thank you for the feedback!

Akshat8510 · 2026-01-23T18:06:55Z

Hi @Phlip79 and @chtruong814, I wanted to follow up on this PR. I have addressed the feedback regarding technical verbosity across all 6 files.

Akshat8510 · 2026-01-23T18:07:15Z

I noticed the needs-follow-up label is still active, please let me know if there are any further changes required on my end to clear this and move toward a merge. Thank you!

Akshat8510 · 2026-01-29T14:45:00Z

Hi @Phlip79 @chtruong814, I've squashed the commits and restored the technical verbosity as requested. Should I update the branch whenever it is out-of-date with main, or only if merge conflicts occur?

Akshat8510 · 2026-02-09T07:54:57Z

Hi @Phlip79 and @chtruong814, it's been two weeks since I restored the technical verbosity and squashed the history. The PR is ready for a final look.

Akshat8510 · 2026-02-09T07:55:22Z

Could you please trigger the internal tests (approve the workflow) so we can see if any further adjustments are needed? Thank you!

Akshat8510 · 2026-02-14T18:05:13Z

Hi @Phlip79 and @chtruong814,
I noticed the needs follow-up label was added again. Could you please let me know what specific changes are required from my side? I’m happy to address any feedback to help move this forward. Thank you!

Akshat8510 · 2026-03-05T04:02:28Z

Hi @jaredcasper, I sincerely apologize for these unintended logical changes. During a rebase and squash to clean up my Git history, I must have incorrectly resolved some conflicts or applied auto-formatting that altered the code logic.

I will reset the logic in these files to match main exactly and re-apply only the docstring updates. I will also restore the missing technical explanations you pointed out. Thank you for your patience!"

Signed-off-by: Akshat Kumar <akshat230405@gmail.com>

Updated docstrings to improve clarity and detail regarding grad buffer shard ranges and ownership.

Akshat8510 · 2026-03-05T05:32:49Z

Hi @jaredcasper , I have restored the full explanation regarding conceptual partitioning, DP ownership, and the specific range mappings. Thank you for catching this!

Removed detailed attribute documentation from OptimizerConfig.

Akshat8510 · 2026-03-19T03:03:00Z

Hi @jaredcasper, that makes perfect sense. I have removed the redundant Attributes block from the class docstrings in optimizer_config.py to avoid duplication and maintenance issues, relying instead on the existing inline docstrings for each attribute. Thank you!

Akshat8510 · 2026-03-25T03:36:51Z

Hi @chtruong814, could you pls also review my PR ?. The PR is ready for a final look.

Phlip79 · 2026-04-03T16:17:09Z

/ok to test dd73cfe

svcnvidia-nemo-ci · 2026-04-03T17:15:48Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/23955018970

Signed-off-by: Akshat Kumar <akshat230405@gmail.com>

Signed-off-by: Akshat Kumar <akshat230405@gmail.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

Signed-off-by: Akshat Kumar <akshat230405@gmail.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>

# Description NVIDIA#4024 NVIDIA#3746 NVIDIA#2799 NVIDIA#3676 NVIDIA#4133 NVIDIA#3295 NVIDIA#3831 NVIDIA#3754 ## Type of change - [ ] New feature (non-breaking change which adds functionality) - [ ] Infra/Build change (changes to CI/CD workflows or build scripts) - [ ] Code refactoring - [ ] Documentation change - [x] Bug fix - [ ] Breaking change ## Changes - Content 1 - Content 2 - Content 3 - Content 4 ## Checklist - [ ] I have read and followed the contributing guidelines - [ ] The functionality is complete - [ ] I have commented my code, particularly in coverage report uploading steps - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added/updated tests that prove my feature works - [ ] New and existing unit tests pass locally --------- Signed-off-by: Cory Ye <cye@nvidia.com> Signed-off-by: yexiaochuan <yexiaochuan@baidu.com> Signed-off-by: Akshat Kumar <akshat230405@gmail.com> Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com> Signed-off-by: jinliangl <jinliangl@nvidia.com> Co-authored-by: Ritesh Patel <ripatel@nvidia.com> Co-authored-by: Cory Ye <cye@nvidia.com> Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com> Co-authored-by: Xiaochuan Ye <tap91624@gmail.com> Co-authored-by: yexiaochuan <yexiaochuan@baidu.com> Co-authored-by: Akshat8510 <akshat230405@gmail.com> Co-authored-by: Kunlun Li <94586211+kunlunl@users.noreply.github.com> Co-authored-by: eternally-z <zzywzj@gmail.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Co-authored-by: Youngeun Kwon <youngeunk@nvidia.com> Co-authored-by: xuwchen <xuwenc@nvidia.com> Co-authored-by: Li Jinliang <jinliangl@nvidia.com>

Signed-off-by: Akshat Kumar <akshat230405@gmail.com>

Akshat8510 requested review from a team as code owners January 4, 2026 04:41

github-actions Bot requested a review from Phlip79 January 4, 2026 04:41

github-actions Bot added the community-request label Jan 4, 2026

Akshat8510 marked this pull request as draft January 4, 2026 09:57

Akshat8510 marked this pull request as ready for review January 4, 2026 09:57

Akshat8510 force-pushed the docstring-optimizer branch 2 times, most recently from 7199c88 to 8e80e8f Compare January 9, 2026 05:02

chtruong814 added the needs-follow-up Issue needs follow-up label Jan 11, 2026

chtruong814 added needs-follow-up Issue needs follow-up and removed needs-follow-up Issue needs follow-up labels Jan 23, 2026

Akshat8510 force-pushed the docstring-optimizer branch 2 times, most recently from c97b8cd to 71a045f Compare January 29, 2026 14:02

chtruong814 added needs-follow-up Issue needs follow-up and removed needs-follow-up Issue needs follow-up labels Jan 29, 2026

chtruong814 added needs-follow-up Issue needs follow-up and removed needs-follow-up Issue needs follow-up labels Feb 9, 2026

chtruong814 added needs-follow-up Issue needs follow-up and removed needs-follow-up Issue needs follow-up labels Feb 14, 2026

jaredcasper reviewed Mar 4, 2026

View reviewed changes

Comment thread megatron/core/optimizer/distrib_optimizer.py

Comment thread megatron/core/optimizer/distrib_optimizer.py

Comment thread megatron/core/optimizer/distrib_optimizer.py Outdated

chtruong814 added the needs-follow-up Issue needs follow-up label Mar 5, 2026

Akshat8510 added 3 commits March 5, 2026 09:42

removing logical changes causes by mergeconflict

5493532

Signed-off-by: Akshat Kumar <akshat230405@gmail.com>

Merge branch 'main' into docstring-optimizer

0a2a3cf

Refine docstrings for grad buffer mapping methods

70e67d6

Updated docstrings to improve clarity and detail regarding grad buffer shard ranges and ownership.

Akshat8510 requested a review from jaredcasper March 14, 2026 12:37

jaredcasper reviewed Mar 18, 2026

View reviewed changes

Comment thread megatron/core/optimizer/optimizer_config.py Outdated

Akshat8510 added 2 commits March 19, 2026 08:29

Remove attribute documentation from OptimizerConfig

cfdc9d5

Removed detailed attribute documentation from OptimizerConfig.

Merge branch 'main' into docstring-optimizer

dd73cfe

Akshat8510 requested a review from jaredcasper March 19, 2026 03:04

jaredcasper approved these changes Mar 19, 2026

View reviewed changes

chtruong814 added needs-follow-up Issue needs follow-up and removed needs-follow-up Issue needs follow-up labels Mar 19, 2026

ericharper approved these changes Apr 3, 2026

View reviewed changes

ericharper added this pull request to the merge queue Apr 3, 2026

Merged via the queue into NVIDIA:main with commit a025a69 Apr 3, 2026
42 checks passed

chtruong814 removed the needs-follow-up Issue needs follow-up label Apr 3, 2026

yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026

docs: update mcore optimizer docstrings to google style (NVIDIA#2799)

e45aec6

Signed-off-by: Akshat Kumar <akshat230405@gmail.com>

yhgalaxy pushed a commit to yhgalaxy/Megatron-LM that referenced this pull request Jun 17, 2026

docs: update mcore optimizer docstrings to google style (NVIDIA#2799)

e792ba1

Signed-off-by: Akshat Kumar <akshat230405@gmail.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

zhaoyinglia mentioned this pull request Jul 16, 2026

cherry-pick fix pr for megatron_fsdp flagos-ai/Megatron-LM-FL#79

Merged

13 tasks

terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026

docs: update mcore optimizer docstrings to google style (NVIDIA#2799)

176aa7d

Signed-off-by: Akshat Kumar <akshat230405@gmail.com>

Uh oh!

Conversation

Akshat8510 commented Jan 4, 2026

Description

Changes

Files Updated:

Uh oh!

copy-pr-bot Bot commented Jan 4, 2026

Uh oh!

Akshat8510 commented Jan 4, 2026

Uh oh!

Phlip79 commented Jan 9, 2026

Uh oh!

Akshat8510 commented Jan 9, 2026

Uh oh!

Akshat8510 commented Jan 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Akshat8510 commented Jan 23, 2026

Uh oh!

Akshat8510 commented Jan 23, 2026

Uh oh!

Akshat8510 commented Jan 29, 2026

Uh oh!

Akshat8510 commented Feb 9, 2026

Uh oh!

Akshat8510 commented Feb 9, 2026

Uh oh!

Akshat8510 commented Feb 14, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Akshat8510 commented Mar 5, 2026

Uh oh!

Akshat8510 commented Mar 5, 2026

Uh oh!

Uh oh!

Akshat8510 commented Mar 19, 2026

Uh oh!

Akshat8510 commented Mar 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Phlip79 commented Apr 3, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Akshat8510 commented Jan 9, 2026 •

edited

Loading

Akshat8510 commented Mar 25, 2026 •

edited

Loading