nemo_automodel.components.checkpoint.utils

Module Contents

Functions

Name	Description
`_checkpoint_contains_target`	Return whether target points at or inside checkpoint.
`_checkpoint_sort_key`	Order optimizer steps first, then the completed epoch frontier at a tie.
`_checkpoint_step_num`	Return the trailing checkpoint step number, or -1 when the name is not a checkpoint.
`_get_checkpoint_tensor_dtypes`	Inspect checkpoint tensors and return their exact dtypes by key.
`_is_checkpoint_pointer_text_file`	Return whether path looks like a symlink fallback checkpoint pointer.
`_list_existing_checkpoints`	Return existing checkpoint directories whose names end in `step_<N>`.
`_resolve_checkpoint_pointer_target`	Resolve a checkpoint pointer target relative to ckpt_root.
`clear_checkpoint_incomplete`	Clear the marker written by `mark_checkpoint_incomplete`.
`estimate_state_dict_bytes`	Estimate logical bytes in a state dict without materializing tensors.
`estimate_tensor_bytes`	Estimate logical bytes in a tensor without materializing it.
`find_latest_checkpoint`	Resolve the most recent checkpoint directory.
`find_pointer_protected_checkpoints`	Return checkpoints targeted by top-level symlinks or symlink fallback text files.
`format_bytes`	Format bytes as a human-readable GiB value.
`format_missing_checkpoint_dir_error`	Format a helpful error message for a missing checkpoint directory.
`format_output_file_count`	Format the output shard count for user-facing log messages.
`get_rank_safe`	Return the current distributed rank, defaulting to 0 when not initialized.
`get_safetensors_index_total_size`	Return the total checkpoint size recorded in a Hugging Face safetensors index.
`get_tied_lm_head_source_names`	Return candidate checkpoint keys that can source a tied LM head.
`get_world_size_safe`	Return the current distributed world size, defaulting to 1 when not initialized.
`is_checkpoint_incomplete`	Return whether a checkpoint directory was left behind by an interrupted save.
`is_cloud_path`	Check if path is a cloud storage path (MSC).
`is_rank_0`	Return True on the main rank.
`list_automodel_checkpoints`	Return canonical AutoModel `epoch_<E>_step_<S>` checkpoint directories.
`mark_checkpoint_incomplete`	Mark a checkpoint directory as still being written.
`materialize_missing_tied_lm_head`	Populate a missing tied `lm_head.weight` from its embedding source.
`read_checkpoint_metric`	Read a validation metric from checkpoint loss metadata.
`read_checkpoint_pointer`	Resolve a checkpoint pointer symlink or fallback text file.
`resolve_restore_from_to_checkpoint_dir`	Resolve restore_from to a checkpoint directory.
`resolve_trust_remote_code`	Whitelist NVIDIA models to allow remote code execution.

Data

_AUTOMODEL_CHECKPOINT_RE

_CHECKPOINT_STEP_RE

_CheckpointLoadValue

_INCOMPLETE_CHECKPOINT_MARKER

logger

API

nemo_automodel.components.checkpoint.utils._checkpoint_contains_target(
    checkpoint: pathlib.Path,
    target: pathlib.Path
) -> bool

Return whether target points at or inside checkpoint.

nemo_automodel.components.checkpoint.utils._checkpoint_sort_key(
    path: pathlib.Path
) -> tuple[int, int]

Order optimizer steps first, then the completed epoch frontier at a tie.

nemo_automodel.components.checkpoint.utils._checkpoint_step_num(
    path: pathlib.Path
) -> int

Return the trailing checkpoint step number, or -1 when the name is not a checkpoint.

nemo_automodel.components.checkpoint.utils._get_checkpoint_tensor_dtypes(
    pretrained_model_name_or_path: str,
    hf_config: typing.Any,
    load_kwargs: collections.abc.Mapping[str, nemo_automodel.components.checkpoint.utils._CheckpointLoadValue] | None = None
) -> dict[str, torch.dtype]

Inspect checkpoint tensors and return their exact dtypes by key.

This reads checkpoint metadata only by loading tensors on the meta device, so it preserves the per-tensor dtype information without materializing full checkpoint weights in memory.

nemo_automodel.components.checkpoint.utils._is_checkpoint_pointer_text_file(
    path: pathlib.Path,
    mode: int
) -> bool

Return whether path looks like a symlink fallback checkpoint pointer.

nemo_automodel.components.checkpoint.utils._list_existing_checkpoints(
    ckpt_root: pathlib.Path
) -> list[pathlib.Path]

Return existing checkpoint directories whose names end in step_<N>.

nemo_automodel.components.checkpoint.utils._resolve_checkpoint_pointer_target(
    ckpt_root: pathlib.Path,
    raw_target: str
) -> pathlib.Path | None

Resolve a checkpoint pointer target relative to ckpt_root.

nemo_automodel.components.checkpoint.utils.clear_checkpoint_incomplete(
    checkpoint_dir: str | pathlib.Path
) -> None

Clear the marker written by mark_checkpoint_incomplete.

Parameters:

checkpoint_dir

str | Path

Directory of the checkpoint that finished writing.

nemo_automodel.components.checkpoint.utils.estimate_state_dict_bytes(
    state_dict: dict[str, torch.Tensor]
) -> int | None

Estimate logical bytes in a state dict without materializing tensors.

nemo_automodel.components.checkpoint.utils.estimate_tensor_bytes(
    tensor: torch.Tensor
) -> int

Estimate logical bytes in a tensor without materializing it.

nemo_automodel.components.checkpoint.utils.find_latest_checkpoint(
    checkpoint_dir: str | pathlib.Path
) -> str | pathlib.Path | None

Resolve the most recent checkpoint directory.

Preference order

Valid LATEST symlink or txt file under checkpoint_dir
Highest step directory under checkpoint_dir whose name ends in step_<N>, preferring the later epoch frontier when canonical checkpoint steps tie

Directories left behind by an interrupted save are skipped at both steps. The LATEST target needs the same filter as the step scan: re-saving a step whose directory LATEST already points at leaves the pointer aimed at a directory that is being rewritten, so an interruption there would otherwise resume from a partially written checkpoint.

Returns: str | Path | None

Path (or str) of the latest checkpoint directory, or None.

nemo_automodel.components.checkpoint.utils.find_pointer_protected_checkpoints(
    ckpt_root: pathlib.Path,
    checkpoints: list[pathlib.Path]
) -> set[pathlib.Path]

Return checkpoints targeted by top-level symlinks or symlink fallback text files.

nemo_automodel.components.checkpoint.utils.format_bytes(
    num_bytes: int
) -> str

Format bytes as a human-readable GiB value.

nemo_automodel.components.checkpoint.utils.format_missing_checkpoint_dir_error(
    checkpoint_dir: str,
    restore_from: str,
    resolved_ckpt_dir: str
) -> str

Format a helpful error message for a missing checkpoint directory.

nemo_automodel.components.checkpoint.utils.format_output_file_count(
    count: int
) -> str

Format the output shard count for user-facing log messages.

nemo_automodel.components.checkpoint.utils.get_rank_safe() -> int

Return the current distributed rank, defaulting to 0 when not initialized.

nemo_automodel.components.checkpoint.utils.get_safetensors_index_total_size(
    index_path: str | None
) -> int | None

Return the total checkpoint size recorded in a Hugging Face safetensors index.

nemo_automodel.components.checkpoint.utils.get_tied_lm_head_source_names(
    model: torch.nn.Module,
    lm_head_param_name: str | None = None
) -> list[str]

Return candidate checkpoint keys that can source a tied LM head.

Parameters:

model

nn.Module

Model or pipeline stage to inspect.

lm_head_param_name

str | NoneDefaults to None

Optional normalized LM head FQN.

Returns: list[str]

Ordered list of possible source FQNs.

nemo_automodel.components.checkpoint.utils.get_world_size_safe() -> int

Return the current distributed world size, defaulting to 1 when not initialized.

nemo_automodel.components.checkpoint.utils.is_checkpoint_incomplete(
    checkpoint_dir: str | pathlib.Path
) -> bool

Return whether a checkpoint directory was left behind by an interrupted save.

Parameters:

checkpoint_dir

str | Path

Directory to inspect.

Returns: bool

True when the in-progress marker is still present.

nemo_automodel.components.checkpoint.utils.is_cloud_path(
    path: str | pathlib.Path
) -> bool

Check if path is a cloud storage path (MSC).

nemo_automodel.components.checkpoint.utils.is_rank_0() -> bool

Return True on the main rank.

nemo_automodel.components.checkpoint.utils.list_automodel_checkpoints(
    ckpt_root: pathlib.Path
) -> list[pathlib.Path]

Return canonical AutoModel epoch_<E>_step_<S> checkpoint directories.

nemo_automodel.components.checkpoint.utils.mark_checkpoint_incomplete(
    checkpoint_dir: str | pathlib.Path
) -> None

Mark a checkpoint directory as still being written.

A save that is interrupted (wall-clock limit, preemption, OOM) leaves the directory behind with only some of its components written. The marker makes that state observable so a resumed run neither loads nor retains it; it is removed by clear_checkpoint_incomplete once the checkpoint is published.

Parameters:

checkpoint_dir

str | Path

Directory of the checkpoint being written.

nemo_automodel.components.checkpoint.utils.materialize_missing_tied_lm_head(
    state_dict: dict[str, typing.Any],
    model: torch.nn.Module,
    allow_current_lm_head_fallback: bool = False
) -> bool

Populate a missing tied lm_head.weight from its embedding source.

Hugging Face checkpoints for tied-embedding models often omit lm_head.weight entirely. That is fine for unsplit models where tie_weights() can restore the alias, but it breaks pipeline-parallel last stages which own lm_head but not embed_tokens.

Parameters:

state_dict

dict[str, Any]

Checkpoint state dict to mutate in place.

model

nn.Module

Target model or pipeline stage.

allow_current_lm_head_fallback

boolDefaults to False

If True, fall back to the current lm_head tensor when the tied source cannot be found in state_dict. This preserves legacy resume behavior for older checkpoints that were saved without a local lm_head.weight.

Returns: bool

True if a missing lm_head.weight was materialized, else False.

nemo_automodel.components.checkpoint.utils.read_checkpoint_metric(
    checkpoint: pathlib.Path,
    metric_key: str | None
) -> float | None

Read a validation metric from checkpoint loss metadata.

nemo_automodel.components.checkpoint.utils.read_checkpoint_pointer(
    ckpt_root: str | pathlib.Path,
    link_name: str
) -> pathlib.Path | None

Resolve a checkpoint pointer symlink or fallback text file.

nemo_automodel.components.checkpoint.utils.resolve_restore_from_to_checkpoint_dir(
    checkpoint_dir: str | pathlib.Path,
    restore_from: str
) -> str | None

Resolve restore_from to a checkpoint directory.

Returns: str | None

str: resolved checkpoint directory

Raises:

RuntimeError: If restore_from selects a checkpoint that an interrupted save left incomplete.

nemo_automodel.components.checkpoint.utils.resolve_trust_remote_code(
    pretrained_model_name_or_path
)

Whitelist NVIDIA models to allow remote code execution.

Parameters:

pretrained_model_name_or_path

str

The name or path of the pretrained model.

Returns:

True if the model should be loaded with trust_remote_code, False otherwise.

nemo_automodel.components.checkpoint.utils._AUTOMODEL_CHECKPOINT_RE = re.compile('^epoch_(\\d+)_step_(\\d+)$')

nemo_automodel.components.checkpoint.utils._CHECKPOINT_STEP_RE = re.compile('step_(\\d+)$')

nemo_automodel.components.checkpoint.utils._CheckpointLoadValue = str | bool | os.PathLike[str] | Mapping[str, str] | Mapping[str, torch.Tensor] |...

nemo_automodel.components.checkpoint.utils._INCOMPLETE_CHECKPOINT_MARKER = '.incomplete'

nemo_automodel.components.checkpoint.utils.logger = logging.getLogger(__name__)