nemo_automodel.components.checkpoint.utils

Module Contents

Functions

Name	Description
`_checkpoint_contains_target`	Return whether target points at or inside checkpoint.
`_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

_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_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, object] | 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>

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._INCOMPLETE_CHECKPOINT_MARKER = '.incomplete'

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