Skip to content

[Dev] Support LayerWiseDistributedOptimizer with torch_dist checkpoint format - #1928

Merged
skyw merged 23 commits into
NVIDIA:devfrom
BoxiangW:boxiangw/dist_ckpt_with_dist_muon
Nov 4, 2025
Merged

[Dev] Support LayerWiseDistributedOptimizer with torch_dist checkpoint format#1928
skyw merged 23 commits into
NVIDIA:devfrom
BoxiangW:boxiangw/dist_ckpt_with_dist_muon

Conversation

@BoxiangW

@BoxiangW BoxiangW commented Oct 24, 2025

Copy link
Copy Markdown
Contributor

PR to main

Support variable DP size loading (DP-reshard). See https://wandb.ai/nvidia/muon-test/reports/Dist-ckpt-save-load-loss-verification--VmlldzoxNDkxODkyNw for a simple reference on loss curve report

It is loading at step 100 and the loss curve is overlapped with previous run

Screenshot 2025-10-31 at 11 34 50 AM

@BoxiangW
BoxiangW requested review from a team as code owners October 24, 2025 21:58
@copy-pr-bot

copy-pr-bot Bot commented Oct 24, 2025

Copy link
Copy Markdown

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.

@BoxiangW BoxiangW self-assigned this Oct 24, 2025
@BoxiangW BoxiangW added Run tests Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. labels Oct 31, 2025
@BoxiangW BoxiangW added this to the Core 0.15 milestone Oct 31, 2025
@BoxiangW BoxiangW linked an issue Oct 31, 2025 that may be closed by this pull request
@BoxiangW

Copy link
Copy Markdown
Contributor Author

/ok to test 387793d

@skyw

skyw commented Nov 3, 2025

Copy link
Copy Markdown
Contributor

I remember the license check has been fixed. @BoxiangW could you update the branch and retry?

@skyw
skyw requested a review from FDecaYed November 3, 2025 02:29
@BoxiangW

BoxiangW commented Nov 3, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 073848a

@BoxiangW

BoxiangW commented Nov 3, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 3bf9b7b

@BoxiangW

BoxiangW commented Nov 3, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 617e1d1

BoxiangW and others added 15 commits November 3, 2025 18:00
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
fix init_state_fn not initialized error
remove debug logging
@FDecaYed

FDecaYed commented Nov 3, 2025

Copy link
Copy Markdown
Contributor

/ok to test 6c1aed9

@FDecaYed

FDecaYed commented Nov 3, 2025

Copy link
Copy Markdown
Contributor

/ok to test 1018a43

@dimapihtar dimapihtar left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

LGTM. Thank you!

@BoxiangW

BoxiangW commented Nov 3, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test b6c156d

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW

BoxiangW commented Nov 4, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 12a6250

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW

BoxiangW commented Nov 4, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 31e116f

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. Run tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[ENHANCEMENT] Distributed Muon and local Muon integration

4 participants