> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.checkpoint.utils

## Module Contents

### Functions

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

### Data

[`_AUTOMODEL_CHECKPOINT_RE`](#nemo_automodel-components-checkpoint-utils-_AUTOMODEL_CHECKPOINT_RE)

[`_CHECKPOINT_STEP_RE`](#nemo_automodel-components-checkpoint-utils-_CHECKPOINT_STEP_RE)

[`_INCOMPLETE_CHECKPOINT_MARKER`](#nemo_automodel-components-checkpoint-utils-_INCOMPLETE_CHECKPOINT_MARKER)

[`logger`](#nemo_automodel-components-checkpoint-utils-logger)

### API

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

Return whether target points at or inside checkpoint.

```python
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.

```python
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.

```python
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.

```python
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>`.

```python
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.

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

Clear the marker written by :func:`mark_checkpoint_incomplete`.

**Parameters:**

**`checkpoint_dir`** `str | Path`

Directory of the checkpoint that finished writing.

---

```python
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.

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

Estimate logical bytes in a tensor without materializing it.

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

Resolve the most recent checkpoint directory.

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.

```python
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.

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

Format bytes as a human-readable GiB value.

```python
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.

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

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

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

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

```python
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.

```python
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 | None` — default: None

Optional normalized LM head FQN.

---

**Returns:** `list[str]`

Ordered list of possible source FQNs.

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

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

```python
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.

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

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

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

Return True on the main rank.

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

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

```python
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 :func:`clear_checkpoint_incomplete` once the checkpoint is
published.

**Parameters:**

**`checkpoint_dir`** `str | Path`

Directory of the checkpoint being written.

---

```python
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`** `bool` — default: 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`.

```python
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.

```python
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.

```python
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.

```python
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.

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

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

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

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