[checkpoint] Log timings for checkpoint IO save and load #11972
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
ashbhandare
reviewed
Jan 28, 2025
ashbhandare
reviewed
Jan 28, 2025
ashbhandare
reviewed
Jan 28, 2025
ashbhandare
reviewed
Jan 28, 2025
ananthsub
force-pushed
the
ckpt-log-times
branch
from
e15c254 to
c2f960f
Compare
|
beep boop 🤖: 🙏 The following files have warnings. In case you are familiar with these, please try helping us to improve the code base. Your code was analyzed with PyLint. The following annotations have been identified: Mitigation guide:
By applying these rules, we reduce the occurance of this message in future. Thank you for improving NeMo's documentation! |
|
Overall, changes LGTM. Just one question -- what is the intended behavior when async checkpointing is enabled? Do you think these logs will be misleading in that case? |
In this case we'd report the time spent blocked while staging the checkpoint. In follow ups I'd like to make the logging more granular so that we can report time spent on staging vs finalization for the async case |
|
Sorry, one other comment. Do you think we should consider making this logging optional? Or provide the option to log only a single rank? If we're running a job on many GPUs, I wonder if this would clutter the logs. |
maanug-nv
reviewed
Jan 30, 2025
|
Just curious if the calls in these locations should be timed as well: |
Or use |
|
Could you also include the iteration number? Filename is optional. Here is the doc for requirements from Heimdall: https://docs.google.com/document/d/1J5TVmg4msY1HA0kcR6etb4Ha72gBRC2aPen1Z0rinQI/edit?usp=sharing As for cluttering logs across many ranks, I agree. Currently MegatronLM only prints the (min, max) across all ranks for each checkpoint. E.g.: Time above is in milliseconds, |
ananthsub
force-pushed
the
ckpt-log-times
branch
2 times, most recently
from
9c20142 to
fff6310
Compare
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: ananthsub <ananthsub@users.noreply.github.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: ananthsub <ananthsub@users.noreply.github.com>
Signed-off-by: ananthsub <ananthsub@users.noreply.github.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
ananthsub
force-pushed
the
ckpt-log-times
branch
from
2668c5d to
452b664
Compare
Signed-off-by: ananthsub <ananthsub@users.noreply.github.com>
ko3n1g
pushed a commit
that referenced
this pull request
Mar 6, 2025
* [checkpoint] Log timings for checkpoint IO save and load Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com> * add rank to logline for megatron strategy Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com> --------- Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
chtruong814
pushed a commit
that referenced
this pull request
Mar 7, 2025
…2520) * [checkpoint] Log timings for checkpoint IO save and load * add rank to logline for megatron strategy --------- Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com> Co-authored-by: Ananth Subramaniam <ansubramania@nvidia.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Adds timings around checkpoint save/load IO calls within Nemo1 and Nemo2 Lightning strategies.
To test, sanity checked logs from existing checkpointing tests:
Collection: nlp, llm
Changelog
Adds timings around checkpoint save/load IO calls within Nemo1 and Nemo2 Lightning strategies
Usage
# Add a code snippet demonstrating how to use thisGitHub Actions CI
The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.
The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".
Before your PR is "Ready for review"
Pre checks:
PR Type:
If you haven't finished some of the above items you can still open "Draft" PR.
Who can review?
Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.
Additional Information