Skip to content

Fixed the problem of loading universal checkpoint error in multi-machine mode. #7601

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
Merged
tohtana merged 2 commits into from
Sep 28, 2025
Merged

Fixed the problem of loading universal checkpoint error in multi-machine mode. #7601

tohtana merged 2 commits into from
Sep 28, 2025

Conversation

@zhengchenyu
Copy link
Contributor

@zhengchenyu zhengchenyu commented Sep 28, 2025

In a multi-machine environment, loading the stage3 universal checkpoint will produce incorrect results, causing the loss to increase abnormally.

@zhengchenyu
Copy link
Contributor Author

zhengchenyu commented Sep 28, 2025

It seems that unit testing currently only works in a single-machine environment. Therefore, I'll use my own examples in for verification. Env: 2 node, 8 gpu per node. There are three experiments:

  • base The base example is a non-interrupt, non-checkpoint example.
  • bug The bug is based on the current master version of the code, with checkpoint at step 20, then convert to universal checkpoint, and then load the universal checkpoint. Abnormal increase in loss.
  • fix The fix is ​​based on this PR, also with checkpoint at step 20, then convert to universal checkpoint, and then reload the universal checkpoint. It can be seen that the loss is normal.
截屏2025-09-28 19 37 40

tohtana
tohtana approved these changes Sep 28, 2025
View reviewed changes
Copy link
Collaborator

@tohtana tohtana left a comment

Choose a reason for hiding this comment

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

Great catch! Thank you for the fix!

@tohtana tohtana enabled auto-merge (squash) September 28, 2025 17:17
@tohtana tohtana merged commit 47b3fb5 into deepspeedai:master Sep 28, 2025
13 of 14 checks passed
@zhengchenyu zhengchenyu deleted the fix.load.universal.multi.nodes branch September 29, 2025 01:39
mauryaavinash95 pushed a commit to DataStates/DeepSpeed that referenced this pull request Oct 4, 2025
@zhengchenyu
Fixed the problem of loading universal checkpoint error in multi-mach…
b460538 
…ine mode. (deepspeedai#7601)

In a multi-machine environment, loading the stage3 universal checkpoint
will produce incorrect results, causing the loss to increase abnormally.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@tohtana tohtana tohtana approved these changes

@tjruwase tjruwase Awaiting requested review from tjruwase tjruwase is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@zhengchenyu @tohtana