Skip to content

support loading checkpoints from a given step - #1040

Merged
ananthsub merged 4 commits into
NVIDIA-NeMo:mainfrom
ananthsub:ckpt-step
Oct 23, 2025
Merged

support loading checkpoints from a given step#1040
ananthsub merged 4 commits into
NVIDIA-NeMo:mainfrom
ananthsub:ckpt-step

Conversation

@ananthsub

Copy link
Copy Markdown
Contributor

Add support for ckpt_step in the checkpoint config.

This allows users to specify a specific checkpoint directory to resume from based on the iteration. This is an opt-in feature. The default is still to resume from the latest checkpoint saved.

Special handling is needed for PEFT where there are 2 checkpoint loads:

  1. an initial load from pretrained_checkpoint
  2. potentially a load from the load directory in case of resuming the adapter states within a training run.

ckpt_step is intended to only take effect for load directories as pretrained_checkpoint should be unchanged over the course of training.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub
ananthsub requested a review from dimapihtar October 22, 2025 10:40
@ananthsub ananthsub added the r0.2.0 Cherry-pick label for r0.2.0 release branch label Oct 22, 2025
@copy-pr-bot

copy-pr-bot Bot commented Oct 22, 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.

@ananthsub

Copy link
Copy Markdown
Contributor Author

/ok to test 2d65c06

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub

Copy link
Copy Markdown
Contributor Author

/ok to test deaaa85

@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!

@ananthsub
ananthsub merged commit ed15298 into NVIDIA-NeMo:main Oct 23, 2025
35 checks passed
@ananthsub
ananthsub deleted the ckpt-step branch October 23, 2025 12:10
chtruong814 pushed a commit that referenced this pull request Oct 23, 2025
* support loading checkpoints from a given step

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

* support loading checkpoints from a given step

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

* add config test

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

* fix tests

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

---------

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: NeMo Bot <nemo-bot@nvidia.com>
ananthsub added a commit that referenced this pull request Oct 23, 2025
* support loading checkpoints from a given step



* support loading checkpoints from a given step



* add config test



* fix tests



---------

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: NeMo Bot <nemo-bot@nvidia.com>
Co-authored-by: Ananth Subramaniam <ansubramania@nvidia.com>
paul-gibbons pushed a commit to paul-gibbons/Megatron-Bridge that referenced this pull request Oct 29, 2025
* support loading checkpoints from a given step

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

* support loading checkpoints from a given step

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

* add config test

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

* fix tests

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

---------

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Paul Gibbons <pgibbons@nvidia.com>
nv-mollys pushed a commit that referenced this pull request Oct 31, 2025
* support loading checkpoints from a given step

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

* support loading checkpoints from a given step

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

* add config test

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

* fix tests

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

---------

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: mollys <mollys@mollys.nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

r0.2.0 Cherry-pick label for r0.2.0 release branch

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants