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.

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

Support writing optimizer checkpoint only on rank0 and make UT pass on A100 #142

Merged
merged 12 commits into from
Mar 1, 2024

Conversation

tocean
Copy link
Contributor

@tocean tocean commented Dec 11, 2023

Description
Support writing optimizer checkpoint only on rank0 and make UT pass on A100.

  • support write checkpoint on rank0.
    With this PR, we don't need to change checkpointing in MS-AMP-Examples.
  • Fix some bugs of TransformerEngine integration and make UT pass on A100
  • Improve document

@tocean tocean force-pushed the yuxiang/megatron_checkpoint branch from 5b2d259 to a946cdc Compare February 28, 2024 02:43
@tocean tocean changed the title support write checkpoint on rank0 Support only writing optimizer checkpoint on rank0 and fix bugs on A100 Feb 28, 2024
@tocean tocean changed the title Support only writing optimizer checkpoint on rank0 and fix bugs on A100 Support writing optimizer checkpoint only on rank0 and fix bugs on A100 Feb 28, 2024
@tocean tocean changed the title Support writing optimizer checkpoint only on rank0 and fix bugs on A100 Support writing optimizer checkpoint only on rank0 and make UT pass on A100 Feb 28, 2024
@tocean tocean merged commit 0a2cd72 into main Mar 1, 2024
9 checks passed
@tocean tocean deleted the yuxiang/megatron_checkpoint branch March 1, 2024 08:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants