`nemo_rl.utils.checkpoint`#

Checkpoint management utilities for the rl algorithm loop.

It handles logic at the algorithm level. Each RL Actor is expected to have its own checkpoint saving function (called by the algorithm loop).

Module Contents#

Classes#

`PretrainedCheckpointConfig`	Configuration for restoring initial weights from a pre-existing Megatron checkpoint.
`CheckpointingConfig`	Configuration for checkpoint management.
`CheckpointManager`	Manages model checkpoints during training.

Functions#

_load_checkpoint_history

Load the history of checkpoints and their metrics.

Data#

PathLike

API#

nemo_rl.utils.checkpoint.PathLike#: None

class nemo_rl.utils.checkpoint.PretrainedCheckpointConfig#

Bases: typing.TypedDict

Configuration for restoring initial weights from a pre-existing Megatron checkpoint.

When set, the policy will restore its initial weights from this checkpoint instead of loading them from model_name. Supported by the Megatron backend only; DTensor backends continue to use HuggingFace weights via model_name.

.. attribute:: path

Filesystem path to the checkpoint to load.

For "megatron_bridge" format: may be either a specific iteration directory that contains a run_config.yaml file (e.g. /checkpoints/iter_0005000/) or a checkpoint root directory that contains iter_* subdirectories. When a root directory is given the latest iter_* subdirectory is used automatically.
For "megatron_lm" format: may be the checkpoint root directory (containing iter_* subdirectories and a latest_checkpointed_iteration.txt tracker file) or a specific iteration directory (e.g. /mlm_checkpoints/iter_0005000/). The checkpoint must use the torch_dist format (i.e. contain a metadata.json file); the legacy torch format is not supported.

.. attribute:: format

Checkpoint format. Use "megatron_bridge" for checkpoints saved by megatron-bridge (e.g. produced by a prior NeMo-RL run) and "megatron_lm" for checkpoints saved by upstream Megatron-LM.

Initialization

Initialize self. See help(type(self)) for accurate signature.

path: str#: None

format: Literal[megatron_bridge, megatron_lm]#: None

class nemo_rl.utils.checkpoint.CheckpointingConfig#

Bases: typing.TypedDict

Configuration for checkpoint management.

Attributes: enabled (bool): Whether checkpointing is enabled. checkpoint_dir (PathLike): Directory where checkpoints will be saved. metric_name (str | None): Name of the metric to use for determining best checkpoints. Must be of the form “val:<metric_name>” or “train:<metric_name>” to indicate whether the metric should be taken from the validation or training metrics. higher_is_better (bool): Whether higher values of the metric indicate better performance. keep_top_k (Optional[int]): Number of best checkpoints to keep. If None, all checkpoints are kept. model_save_format (str | None): Format for saving model (v2 allowed values: “torch_save” or “safetensors”, v1 allowed values: None). save_consolidated (bool): Whether to save consolidated checkpoints (for HF compatibility). model_cache_dir (str): Directory for model cache (for safetensors format). model_repo_id (str): Repository ID for the model (for safetensors format). is_peft (bool): Whether the model uses PEFT. save_optimizer (bool): Whether to save optimizer state with checkpoints.

Initialization

Initialize self. See help(type(self)) for accurate signature.

enabled: bool#: None

checkpoint_dir: nemo_rl.utils.checkpoint.PathLike#: None

metric_name: str | None#: None

higher_is_better: bool#: None

save_period: int#: None

keep_top_k: NotRequired[int]#: None

checkpoint_must_save_by: NotRequired[str | None]#: None

pretrained_checkpoint: NotRequired[nemo_rl.utils.checkpoint.PretrainedCheckpointConfig]#: None

save_optimizer: NotRequired[bool]#: None

model_save_format: NotRequired[str | None]#: None

save_consolidated: NotRequired[bool]#: None

model_cache_dir: NotRequired[str]#: None

model_repo_id: NotRequired[str]#: None

is_peft: NotRequired[bool]#: None

peft_config: NotRequired[Any]#: None

is_async: NotRequired[bool]#: None

class nemo_rl.utils.checkpoint.CheckpointManager( config: nemo_rl.utils.checkpoint.CheckpointingConfig, )#

Manages model checkpoints during training.

This class handles creating checkpoint dirs, saving training info, and configurations. It also provides utilities for keeping just the top-k checkpoints. The checkpointing structure looks like this:

checkpoint_dir/
    step_0/
        training_info.json
        config.yaml
        policy.py (up to the algorithm loop to save here)
        policy_optimizer.py (up to the algorithm loop to save here)
        ...
    step_1/
        ...

Attributes: Derived from the CheckpointingConfig.

Initialization

Initialize the checkpoint manager.

Parameters:: config (CheckpointingConfig)

static get_resume_paths( last_checkpoint_path: Optional[nemo_rl.utils.checkpoint.PathLike], *, model_component: Literal[policy, value] = 'policy', ) → tuple[Optional[pathlib.Path], Optional[pathlib.Path]]#

Get weights and optimizer paths for resuming from a checkpoint.

Parameters:

last_checkpoint_path – Path to the last checkpoint, or None if starting fresh.
model_component – Model subtree to resolve. Policy is the default for backward compatibility with algorithms that only checkpoint a policy.

Returns:

Tuple of (weights_path, optimizer_path). Both are None if no checkpoint. optimizer_path is None if checkpoint exists but optimizer state was not saved.

init_tmp_checkpoint( step: int, training_info: Mapping[str, Any], run_config: Optional[pydantic.BaseModel] = None, ) → nemo_rl.utils.checkpoint.PathLike#

Initialize a temporary checkpoint directory.

Creates a temporary directory for a new checkpoint and saves training info and configuration. The directory is named ‘tmp_step_{step}’ and will be renamed to ‘step_{step}’ when the checkpoint is completed. We do it this way to allow the algorithm loop to save any files it wants to save in a safe, temporary directory.

Parameters:

step (int) – The training step number.
training_info (dict[str, Any]) – Dictionary containing training metrics and info.
run_config (Optional[BaseModel]) – Optional configuration for the training run.

Returns:

Path to the temporary checkpoint directory.

Return type:

PathLike

_rename_checkpoint( checkpoint_path: nemo_rl.utils.checkpoint.PathLike, ) → None#

Rename tmp_step_N to step_N.

If step_N already exists (defensive guard for edge cases, e.g. resuming training), performs a pseudo-atomic swap via an intermediate old_step_N directory.

finalize_checkpoint( checkpoint_path: nemo_rl.utils.checkpoint.PathLike, ) → None#

Complete a checkpoint synchronously (rename + delete old).

This is the original synchronous API, preserved for backward compatibility. For async-aware usage, prefer begin_finalization() + finalize_pending().

Parameters:: checkpoint_path (PathLike) – Path to the temporary checkpoint directory.

begin_finalization( checkpoint_path: nemo_rl.utils.checkpoint.PathLike, wait_fn: Optional[Callable[[], None]] = None, ) → None#

Start background finalization of a checkpoint.

Spawns a daemon thread that calls wait_fn (blocks until async writers finish), renames tmp_step_N to step_N, then queues old-checkpoint deletion. All writes to checkpoint_path must be complete before calling this. If a previous finalization is still active, blocks until it completes.

Parameters:

checkpoint_path – Path to tmp_step_N directory from init_tmp_checkpoint().
wait_fn – Callable that blocks until all async writes are complete. For Megatron async save: policy.finalize_async_save. For sync saves: None (rename immediately).

finalize_pending() → None#

Block until the in-flight rename completes. Does NOT wait for deletion.

Re-raises any exception that occurred in the background thread. No-op if nothing is pending. Safe to call multiple times.

shutdown() → None#

Block until rename + all queued deletions complete. Call at training exit.

Safe to call multiple times.

static _warn_on_delete_failure(future: Future[None]) → None#

Surface background old-checkpoint deletion errors as warnings.

Deletion is best-effort cleanup that runs off the critical path, so a failure should not abort training — but it must not be silent either.

__enter__() → nemo_rl.utils.checkpoint.CheckpointManager#: Enter a context that guarantees shutdown() flushes on exit.

__exit__( exc_type: Any, exc_val: Any, exc_tb: Any, ) → Literal[False]#

Flush pending finalizations when leaving the context.

If an exception is already propagating out of the with block, the flush is best-effort and never masks the original exception. On a normal exit, finalization errors propagate so a failed final checkpoint is not silently dropped. Always returns False so exceptions are re-raised.

property has_pending_finalization: bool#: Whether a background finalization is in-flight.

remove_old_checkpoints(exclude_latest: bool = True) → None#

Remove checkpoints that are not in the top-k or latest based on the (optional) metric.

If keep_top_k is set, this method removes all checkpoints except the top-k best ones. The “best” checkpoints are determined by:

If a metric is provided: the given metric value and the higher_is_better setting. When multiple checkpoints have the same metric value, more recent checkpoints (higher step numbers) are preferred.
If no metric is provided: the step number. The most recent k checkpoints are kept.

Parameters:: exclude_latest (bool) – Whether to exclude the latest checkpoint from deletion. (may result in K+1 checkpoints)

get_best_checkpoint_path() → Optional[str]#

Get the path to the best checkpoint based on the metric.

Returns the path to the checkpoint with the best metric value. If no checkpoints exist, returns None. If some checkpoints are missing the metric, they are filtered out with a warning. If no checkpoints have the metric, returns the latest checkpoint.

Returns:: Path to the best checkpoint, or None if no checkpoints exist.
Return type:: Optional[str]

get_latest_checkpoint_path() → Optional[str]#

Get the path to the latest checkpoint.

Returns the path to the checkpoint with the highest step number.

Returns:: Path to the latest checkpoint, or None if no checkpoints exist.
Return type:: Optional[str]

load_training_info( checkpoint_path: Optional[nemo_rl.utils.checkpoint.PathLike] = None, ) → Optional[dict[str, Any]]#

Load the training info from a checkpoint.

Parameters:: checkpoint_path (Optional[PathLike]) – Path to the checkpoint. If None, returns None.
Returns:: Dictionary containing the training info, or None if checkpoint_path is None.
Return type:: Optional[dict[str, Any]]

nemo_rl.utils.checkpoint._load_checkpoint_history( checkpoint_dir: pathlib.Path, ) → list[tuple[int, nemo_rl.utils.checkpoint.PathLike, dict[str, Any]]]#

Load the history of checkpoints and their metrics.

Parameters:: checkpoint_dir (Path) – Directory containing the checkpoints.
Returns:: List of tuples containing (step_number, checkpoint_path, info) for each checkpoint.
Return type:: list[tuple[int, PathLike, dict[str, Any]]]

nemo_rl.utils.checkpoint#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.utils.checkpoint`#