feat(adapter/nemo): add keep_mlf_checkpoint_on_train_end flag - #103
Merged
Conversation
This change will add a `keep_mlf_checkpoint_on_train_end` flag to MLFlashpointCheckpointCallback. It defaults to False, but when set to True, it skips the deletion of the final checkpoint at the end of training. This ensures the last checkpoint can be preserved for use in E2E testing. Change-Id: Ic238a9dc94fbd9d3c556f0ad1484e277a21fc71e
g-husam
approved these changes
Apr 17, 2026
| trainer.strategy.checkpoint_io.remove_checkpoint(self.base_container.data) | ||
| else: | ||
| _LOGGER.info( | ||
| "Local rank 0: Skipping final checkpoint cleanup due to keep_mlf_checkpoint_on_train_end=True." |
Collaborator
There was a problem hiding this comment.
nit:
normally we should not use f-strings in log statements, but this variable is just a boolean, so stringifying it always is cheap, and this allows us to auto log the variable name with its value
Suggested change
| "Local rank 0: Skipping final checkpoint cleanup due to keep_mlf_checkpoint_on_train_end=True." | |
| "Local rank 0: Skipping final checkpoint cleanup because keep_mlf_checkpoint_on_train_end=%s.", self.keep_mlf_checkpoint_on_train_end |
| self.skip_every_n_steps = skip_every_n_steps if skip_every_n_steps is not None else 0 | ||
| self._enabled = enabled | ||
| self._replication_manager = None | ||
| self.keep_mlf_checkpoint_on_train_end = keep_mlf_checkpoint_on_train_end |
Collaborator
There was a problem hiding this comment.
nit: let's make this "private" as well:
Suggested change
| self.keep_mlf_checkpoint_on_train_end = keep_mlf_checkpoint_on_train_end | |
| self._keep_mlf_checkpoint_on_train_end = keep_mlf_checkpoint_on_train_end |
Change-Id: I5ca3c29cf84ede858c39b280ec263210039d83ea
Python Code Coverage Summary
Minimum allowed line rate is |
C++ Code Coverage Summary
Minimum allowed line rate is |
This file contains hidden or 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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
This change will add a
keep_mlf_checkpoint_on_train_endflag to MLFlashpointCheckpointCallback. It defaults to False, but when set to True, it skips the deletion of the final checkpoint at the end of training. This ensures thelast checkpoint can be preserved for use in E2E testing.