Skip to content

fix: model save directly after validation - #3605

Merged
yuki-97 merged 1 commit into
NVIDIA-NeMo:mainfrom
tdene:tde/fix_save_after_validation
Aug 14, 2026
Merged

yuki-97 merged 1 commit into
NVIDIA-NeMo:mainfrom
tdene:tde/fix_save_after_validation

Conversation

@tdene

@tdene tdene commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

save_checkpoint assumes the training model is on GPU. But when it is called during colocated, right after a validation run, the model is not on GPU.

This is probably best fixed with a bigger refactor, but this PR prevents crashes. Optimizations beyond that are not a priority due to how rare this issue is.

Issues

List issues that this PR closes (syntax):

Usage

  • You can potentially add a usage example below
# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • ...

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
tdene requested review from a team as code owners August 12, 2026 19:09
@copy-pr-bot

copy-pr-bot Bot commented Aug 12, 2026

Copy link
Copy Markdown

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@tdene tdene added the CI:L1 Run doctests, unit tests, and functional tests label Aug 12, 2026
@tdene

tdene commented Aug 12, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test b9dad68

@cspades cspades 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.

when it is called during colocated, right after a validation run, the model is not on GPU, because ...

The inference model is used for validation?

@tdene tdene added CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) and removed CI:L1 Run doctests, unit tests, and functional tests labels Aug 14, 2026
@yuki-97
yuki-97 enabled auto-merge (squash) August 14, 2026 12:44
@yuki-97
yuki-97 merged commit 54dc1f0 into NVIDIA-NeMo:main Aug 14, 2026
232 of 248 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants