> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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.checkpointing

## Module Contents

### Classes

| Name                                                                                                 | Description                                                                |
| ---------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`Checkpointer`](#nemo_automodel-components-checkpoint-checkpointing-Checkpointer)                   | High-level checkpoint manager built on torch.distributed.checkpoint (DCP). |
| [`_AsyncSaveContext`](#nemo_automodel-components-checkpoint-checkpointing-_AsyncSaveContext)         | Internal container for async checkpointing state.                          |
| [`_DTensorInitCopyMode`](#nemo_automodel-components-checkpoint-checkpointing-_DTensorInitCopyMode)   | Adapt full-tensor initialization copies to an already-sharded destination. |
| [`_ModelSavePlanner`](#nemo_automodel-components-checkpoint-checkpointing-_ModelSavePlanner)         | Keep model save plans in one checkpointer-scoped cache namespace.          |
| [`_OptimizerSavePlanner`](#nemo_automodel-components-checkpoint-checkpointing-_OptimizerSavePlanner) | Keep optimizer save plans in one checkpointer-scoped cache namespace.      |

### Functions

| Name                                                                                                                                             | Description                                                                                                        |
| ------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------------------------------------ |
| [`_adapter_path`](#nemo_automodel-components-checkpoint-checkpointing-_adapter_path)                                                             | Return the PEFT adapter safetensors path inside a checkpoint dir (local or `msc://`).                              |
| [`_apply`](#nemo_automodel-components-checkpoint-checkpointing-_apply)                                                                           | Apply a transformation function to parameters (and gradients) only.                                                |
| [`_apply_adapter_forced_dtype_mapping`](#nemo_automodel-components-checkpoint-checkpointing-_apply_adapter_forced_dtype_mapping)                 | Let model adapters override original HF dtype metadata for export-only keys.                                       |
| [`_apply_key_mapping`](#nemo_automodel-components-checkpoint-checkpointing-_apply_key_mapping)                                                   | Rename state-dict keys using a regex mapping or a scoped rename callable.                                          |
| [`_collect_global_tensor_sizes`](#nemo_automodel-components-checkpoint-checkpointing-_collect_global_tensor_sizes)                               | Collect logical tensor sizes across pipeline stages without moving tensor data.                                    |
| [`_convert_checkpoint_with_transformers`](#nemo_automodel-components-checkpoint-checkpointing-_convert_checkpoint_with_transformers)             | Convert a checkpoint using transformers' conversion mapping for models that need tensor merging.                   |
| [`_create_dirs`](#nemo_automodel-components-checkpoint-checkpointing-_create_dirs)                                                               | Create local directory paths and ignore cloud paths.                                                               |
| [`_divide_keys_by_size`](#nemo_automodel-components-checkpoint-checkpointing-_divide_keys_by_size)                                               | Assign keys to deterministic size-based shards.                                                                    |
| [`_ensure_dirs`](#nemo_automodel-components-checkpoint-checkpointing-_ensure_dirs)                                                               | Create directories on all ranks and synchronize across ranks.                                                      |
| [`_ensure_msc_available`](#nemo_automodel-components-checkpoint-checkpointing-_ensure_msc_available)                                             | Raise an error if MSC is not installed but a cloud path is used.                                                   |
| [`_ensure_shared_dirs`](#nemo_automodel-components-checkpoint-checkpointing-_ensure_shared_dirs)                                                 | Create shared DCP directories on group rank zero and synchronize the group.                                        |
| [`_equally_divide_layers`](#nemo_automodel-components-checkpoint-checkpointing-_equally_divide_layers)                                           | Equally divide the state dict keys into num\_shards shards.                                                        |
| [`_format_restricted_load_error`](#nemo_automodel-components-checkpoint-checkpointing-_format_restricted_load_error)                             | -                                                                                                                  |
| [`_get_checkpoint_metadata`](#nemo_automodel-components-checkpoint-checkpointing-_get_checkpoint_metadata)                                       | Read checkpoint metadata, including saved tensor sizes and dtypes.                                                 |
| [`_get_hf_safetensors_reference_path`](#nemo_automodel-components-checkpoint-checkpointing-_get_hf_safetensors_reference_path)                   | Return the local HF safetensors reference directory for a model.                                                   |
| [`_get_original_hf_index_total_size`](#nemo_automodel-components-checkpoint-checkpointing-_get_original_hf_index_total_size)                     | Return the original HF safetensors index total size, if available.                                                 |
| [`_get_shared_parameter_names`](#nemo_automodel-components-checkpoint-checkpointing-_get_shared_parameter_names)                                 | Find checkpoint names referring to the same live parameter.                                                        |
| [`_init_peft_adapters`](#nemo_automodel-components-checkpoint-checkpointing-_init_peft_adapters)                                                 | Initialize the PEFT adapters with the scaled weights.                                                              |
| [`_is_bin_checkpoint`](#nemo_automodel-components-checkpoint-checkpointing-_is_bin_checkpoint)                                                   | Return True if path looks like a PyTorch .bin checkpoint.                                                          |
| [`_is_custom_model`](#nemo_automodel-components-checkpoint-checkpointing-_is_custom_model)                                                       | True if the model has a custom implementation in nemo\_automodel/components/models/.                               |
| [`_is_model_checkpoint_path`](#nemo_automodel-components-checkpoint-checkpointing-_is_model_checkpoint_path)                                     | Return whether a checkpoint path identifies model weights.                                                         |
| [`_is_safetensors_checkpoint`](#nemo_automodel-components-checkpoint-checkpointing-_is_safetensors_checkpoint)                                   | Return True if path looks like a safetensors checkpoint (so we can preserve dtype); else DCP or other.             |
| [`_load_full_state_dict_into_model`](#nemo_automodel-components-checkpoint-checkpointing-_load_full_state_dict_into_model)                       | Load a full (non-sharded) state dict into a potentially FSDP-wrapped model.                                        |
| [`_load_hf_bin_checkpoint`](#nemo_automodel-components-checkpoint-checkpointing-_load_hf_bin_checkpoint)                                         | Load a HuggingFace .bin checkpoint into a state dict.                                                              |
| [`_load_hf_checkpoint_preserving_dtype`](#nemo_automodel-components-checkpoint-checkpointing-_load_hf_checkpoint_preserving_dtype)               | Load a HuggingFace checkpoint into a new state dict so tensor dtypes                                               |
| [`_load_hf_safetensors_checkpoint`](#nemo_automodel-components-checkpoint-checkpointing-_load_hf_safetensors_checkpoint)                         | Load a safetensors checkpoint into a state dict.                                                                   |
| [`_load_safetensors`](#nemo_automodel-components-checkpoint-checkpointing-_load_safetensors)                                                     | Read a safetensors file from a local path or an `msc://` cloud path.                                               |
| [`_materialize_to_hf_views_for_save`](#nemo_automodel-components-checkpoint-checkpointing-_materialize_to_hf_views_for_save)                     | Replace non-contiguous tensor values in `state_dict` with contiguous copies in place.                              |
| [`_maybe_adapt_state_dict_from_hf`](#nemo_automodel-components-checkpoint-checkpointing-_maybe_adapt_state_dict_from_hf)                         | Custom models use state dict adapters to convert the state dict from the Hugging Face format to the native format. |
| [`_maybe_adapt_state_dict_to_hf`](#nemo_automodel-components-checkpoint-checkpointing-_maybe_adapt_state_dict_to_hf)                             | Custom models use state dict adapters to convert the state dict to the Hugging Face format.                        |
| [`_maybe_msc_reader`](#nemo_automodel-components-checkpoint-checkpointing-_maybe_msc_reader)                                                     | Return an MSC filesystem reader for `msc://` paths, else the given reader.                                         |
| [`_maybe_msc_writer`](#nemo_automodel-components-checkpoint-checkpointing-_maybe_msc_writer)                                                     | Return an MSC filesystem writer for `msc://` paths, else the given writer.                                         |
| [`_model_has_dtensors`](#nemo_automodel-components-checkpoint-checkpointing-_model_has_dtensors)                                                 | True if any parameter is a DTensor (model is already sharded).                                                     |
| [`_new_gloo_process_group`](#nemo_automodel-components-checkpoint-checkpointing-_new_gloo_process_group)                                         | Create a Gloo group with the same membership as `process_group`.                                                   |
| [`_normalize_dtype_mapping_to_state_dict_keys`](#nemo_automodel-components-checkpoint-checkpointing-_normalize_dtype_mapping_to_state_dict_keys) | Align original HF dtype metadata with the keys that will be exported.                                              |
| [`_read_paramwrapper_layout_metadata`](#nemo_automodel-components-checkpoint-checkpointing-_read_paramwrapper_layout_metadata)                   | Return the fused expert LoRA layout stamped into a PEFT checkpoint, if any.                                        |
| [`_reinit_non_persistent_buffers`](#nemo_automodel-components-checkpoint-checkpointing-_reinit_non_persistent_buffers)                           | Recompute non-persistent buffers that are not saved in checkpoints.                                                |
| [`_save_safetensors`](#nemo_automodel-components-checkpoint-checkpointing-_save_safetensors)                                                     | Write a safetensors file to a local path or an `msc://` cloud path.                                                |
| [`_should_dequantize_base_checkpoint`](#nemo_automodel-components-checkpoint-checkpointing-_should_dequantize_base_checkpoint)                   | Return whether this load requires checkpoint dequantization.                                                       |
| [`_should_write_consolidated_safetensors`](#nemo_automodel-components-checkpoint-checkpointing-_should_write_consolidated_safetensors)           | Whether to output consolidated HF weights along with sharded weights.                                              |
| [`_should_write_hf_metadata`](#nemo_automodel-components-checkpoint-checkpointing-_should_write_hf_metadata)                                     | Whether to write HF metadata/artifacts for a checkpoint.                                                           |
| [`_summarize_state_dict_key_diff`](#nemo_automodel-components-checkpoint-checkpointing-_summarize_state_dict_key_diff)                           | Summarize state-dict key mismatches for checkpoint load diagnostics.                                               |
| [`_unwrap_ddp_model`](#nemo_automodel-components-checkpoint-checkpointing-_unwrap_ddp_model)                                                     | Return the module that owns model metadata hidden by DDP.                                                          |
| [`_warn_if_inline_consolidation_enabled`](#nemo_automodel-components-checkpoint-checkpointing-_warn_if_inline_consolidation_enabled)             | Educate users about the cost of inline HF consolidation.                                                           |
| [`_warn_if_large_inline_consolidation`](#nemo_automodel-components-checkpoint-checkpointing-_warn_if_large_inline_consolidation)                 | Warn when inline consolidated export is large enough to waste GPU allocation time.                                 |
| [`load_torch_ckpt`](#nemo_automodel-components-checkpoint-checkpointing-load_torch_ckpt)                                                         | Load a torch checkpoint with restricted unpickling by default.                                                     |
| [`save_config`](#nemo_automodel-components-checkpoint-checkpointing-save_config)                                                                 | Save a config to a weights path.                                                                                   |
| [`save_losses`](#nemo_automodel-components-checkpoint-checkpointing-save_losses)                                                                 | Write checkpoint loss metadata to `weights_path/losses.json`.                                                      |
| [`to_empty_parameters_only`](#nemo_automodel-components-checkpoint-checkpointing-to_empty_parameters_only)                                       | Move parameters to the specified device without copying storage, skipping buffers.                                 |

### Data

[`MSC_AVAILABLE`](#nemo_automodel-components-checkpoint-checkpointing-MSC_AVAILABLE)

[`_CONSOLIDATED_SIZE_WARNING_THRESHOLD_BYTES`](#nemo_automodel-components-checkpoint-checkpointing-_CONSOLIDATED_SIZE_WARNING_THRESHOLD_BYTES)

[`_DEFAULT_HF_CONSOLIDATED_SHARD_SIZE_BYTES`](#nemo_automodel-components-checkpoint-checkpointing-_DEFAULT_HF_CONSOLIDATED_SHARD_SIZE_BYTES)

[`_MODELS_REQUIRING_BUFFER_REINIT`](#nemo_automodel-components-checkpoint-checkpointing-_MODELS_REQUIRING_BUFFER_REINIT)

[`_load_metadata_groups`](#nemo_automodel-components-checkpoint-checkpointing-_load_metadata_groups)

[`load_hf_safetensors_state_dict`](#nemo_automodel-components-checkpoint-checkpointing-load_hf_safetensors_state_dict)

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

### API

```python
class nemo_automodel.components.checkpoint.checkpointing.Checkpointer(
    config: nemo_automodel.components.checkpoint.config.CheckpointingConfig,
    dp_rank: int,
    tp_rank: int,
    pp_rank: int,
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    process_group: torch.distributed.ProcessGroup | None = None,
    pp_group: typing.Optional[torch.distributed.ProcessGroup] = None
)
```

High-level checkpoint manager built on torch.distributed.checkpoint (DCP).

Supports:

* HF sharded safetensors via custom storage reader/writer
* Optional consolidated export (config, generation config, tokenizer)
* PEFT adapter save/load handling
* Async save for torch >= 2.9.0

Also provides DP- and global-rank-aware helpers for saving/loading
auxiliary state and utilities to initialize from a base HF checkpoint.

**`_addons`** `= []`

---

**`_consolidation_error`** `BaseException | None = None`

---

**`_consolidation_thread`** `Thread | None = None`

---

**`_load_process_group`** `ProcessGroup | None = None`

---

**`_model_ctx`**

---

**`_optim_ctx`**

---

**`_planner_cache_namespace`** `= uuid.uuid4().hex`

---

**`lifecycle`**

---

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._do_load(
    state_dict: dict[str, torch.Tensor],
    path: str,
    storage_reader: torch.distributed.checkpoint.storage.StorageReader | None = None,
    is_init_step: bool = False
) -> dict[str, torch.Tensor]
```

Load a state dictionary from `path` using DCP or PEFT special-case logic.

**Parameters:**

**`state_dict`** `dict[str, torch.Tensor]`

Mutable mapping to destination tensors of arbitrary shape/dtype. DTensors retain their
global shape and mesh placements; each rank loads its local shard on the existing device.

---

**`path`** `str`

Checkpoint directory path.

---

**`storage_reader`** `StorageReader | None` — default: None

Optional HF storage reader for safetensors.

---

**`is_init_step`** `bool` — default: False

True if loading from a base checkpoint during initialization.

---

**Returns:** `dict[str, torch.Tensor]`

The populated mapping, preserving tensor shapes/dtypes/devices and DTensor placements. DCP fills

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._do_save(
    state_dict: dict[str, torch.Tensor],
    path: str,
    storage_writer: torch.distributed.checkpoint.storage.StorageWriter | None = None
) -> typing.Optional[torch.distributed.checkpoint.state_dict_saver.AsyncSaveResponse]
```

Save a state dictionary to `path` using DCP or PEFT special-case logic.

* For PEFT model saves: only rank 0 writes `adapter_model.safetensors`.
* If async mode is enabled, schedule an asynchronous save.

**Parameters:**

**`state_dict`** `dict[str, torch.Tensor]`

State dict to be serialized.

---

**`path`** `str`

Checkpoint directory path.

---

**`storage_writer`** `StorageWriter | None` — default: None

Optional HF storage writer for safetensors sharding.

---

**Returns:** `Optional[AsyncSaveResponse]`

Optional Future object if async mode is enabled.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._get_load_process_group() -> torch.distributed.ProcessGroup | None
```

Reuse a world-owned CPU metadata group with the caller's membership.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._get_original_model_path(
    model_state: nemo_automodel.components.checkpoint.stateful_wrappers.ModelState
) -> str | None
```

Get the path to the original model from the Hugging Face checkpoint.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._get_storage_reader(
    model_path: str,
    key_mapping: dict[str, str] | collections.abc.Callable[[str], str] | None,
    is_init_step: bool = False,
    is_safetensors: bool | None = None
) -> torch.distributed.checkpoint.storage.StorageReader | None
```

Construct a Hugging Face storage reader when loading safetensors or during init.

Prefers the upstream `torch.distributed.checkpoint.hf_storage.HuggingFaceStorageReader`
when no `key_mapping` is needed, since it uses safetensors' native `get_slice()` for
efficient partial reads (only the bytes for the local DTensor shard are read from disk).
Falls back to the backported reader when `key_mapping` is required or when the upstream
reader is not available.

**Parameters:**

**`model_path`** `str`

Path to the model checkpoint directory or HF snapshot.

---

**`key_mapping`** `dict[str, str] | Callable[[str], str] | None`

Optional regex mapping or callable renaming checkpoint keys.

---

**`is_init_step`** `bool` — default: False

If True, always produce a reader for base HF load.

---

**`is_safetensors`** `bool | None` — default: None

Whether `model_path` holds a safetensors checkpoint; computed
from the directory contents when not supplied.

---

**Returns:** `StorageReader | None`

Configured storage reader, or None for the default DCP FileSystemReader.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._get_storage_writer(
    consolidated_output_path: str | None,
    fqn_to_index_mapping: dict[str, int] | None,
    fqn_to_dtype_mapping: dict[str, str] | None,
    model_path: str,
    consolidation_handled_externally: bool = False
) -> torch.distributed.checkpoint.storage.StorageWriter | None
```

Construct a Hugging Face storage writer for sharded safetensors.

**Parameters:**

**`consolidated_output_path`** `str | None`

Optional path for consolidated artifacts.

---

**`fqn_to_index_mapping`** `dict[str, int] | None`

Optional mapping from FQN to shard index.

---

**`fqn_to_dtype_mapping`** `dict[str, str] | None`

Optional mapping from FQN to original HF safetensors dtype string.

---

**`model_path`** `str`

Path where the model checkpoint is saved.

---

**`consolidation_handled_externally`** `bool` — default: False

If True, consolidation happens outside the writer
(inline on all ranks in sync mode, or on a background thread in async mode), so
the writer's own finish() consolidation is disabled.

---

**Returns:** `StorageWriter | None`

Configured storage writer or None for non-safetensors.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._join_deferred_consolidation() -> None
```

Wait for a pending background consolidation and surface its error, if any.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._load_model_in_parts(
    model_state: nemo_automodel.components.checkpoint.stateful_wrappers.ModelState,
    load_parts: collections.abc.Iterator[nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart],
    model_state_dict: dict[str, torch.Tensor],
    model_path: str,
    storage_reader: torch.distributed.checkpoint.storage.StorageReader
) -> None
```

Load, convert, and release one checkpoint part at a time.

**Parameters:**

**`model_state`** `ModelState`

Wrapper for the model whose final parameter storage is populated.

---

**`load_parts`** `Iterator[CheckpointLoadPart]`

Adapter-owned sequence of checkpoint destinations and finish callbacks.

---

**`model_state_dict`** `dict[str, torch.Tensor]`

Native names mapped to final tensors of arbitrary shape/dtype on their destination
devices. DTensors retain their global shape and mesh placements. Every key must be completed exactly
once across `load_parts`.

---

**`model_path`** `str`

Hugging Face safetensors checkpoint directory.

---

**`storage_reader`** `StorageReader`

Reader already bound to `model_path`. Its parsed metadata is reused across parts.

---

**Raises:**

* `RuntimeError`: If the checkpoint is missing a requested tensor or the parts do not cover all model tensors.
* `TypeError`: If the adapter yields an object other than `CheckpointLoadPart`.
* `ValueError`: If a part is empty or repeats checkpoint or model keys.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._maybe_build_consolidated_index(
    model_state: nemo_automodel.components.checkpoint.stateful_wrappers.ModelState,
    state_dict: dict[str, torch.Tensor]
) -> dict[str, int] | None
```

Build FQN to shard index mapping for consolidated HF export.

Uses the base checkpoint index (if present), removes non-persistent keys,
and assigns new keys to the last shard by default.

**Parameters:**

**`model_state`** `ModelState`

Wrapper exposing the primary model part.

---

**`state_dict`** `dict[str, torch.Tensor]`

Current pipeline stage's subset of the exported state dict. Each value is a tensor of
arbitrary shape representing its full logical tensor, including when its per-rank storage is sharded.

---

**Returns:** `dict[str, int] | None`

Mapping from FQN to shard index, or None when not consolidating.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._maybe_build_original_dtype_mapping(
    model_state: nemo_automodel.components.checkpoint.stateful_wrappers.ModelState,
    state_dict: dict[str, torch.Tensor]
) -> dict[str, str] | None
```

Build FQN to target safetensors dtype mapping for consolidated export.

Original HF safetensors headers provide the baseline mapping when available.
Model-owned adapter overrides are applied even for config-only runs so
intrinsically fp32 tensors retain their required export dtype.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._maybe_log_final_offline_consolidation_hint(
    model_dir: str,
    is_final_checkpoint: bool = False
) -> None
```

Log the final-checkpoint helper hint when consolidated export was disabled.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._maybe_write_offline_consolidation_script(
    model_dir: str
) -> None
```

Write a conservative helper script for offline HF safetensors consolidation.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._release_async_stager(
    context: nemo_automodel.components.checkpoint.checkpointing._AsyncSaveContext
) -> None
```

staticmethod

Close a completed async stager so the next save uses a fresh instance.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer._schedule_deferred_consolidation(
    future: torch.distributed.checkpoint.state_dict_saver.AsyncSaveResponse | None,
    model_dir: str,
    consolidated_dir: str,
    fqn_to_index_mapping: dict[str, int] | None,
    fqn_to_dtype_mapping: dict[str, str] | None,
    process_group: torch.distributed.ProcessGroup | None
) -> None
```

Consolidate HF safetensors on a background thread once the async upload completes.

Every rank schedules the same distributed consolidation used in sync mode, so the
shards written by the async save are merged in parallel across ranks without
blocking the training loop. Collectives inside the consolidation run on the
dedicated Gloo group created at init.

**Parameters:**

**`future`** `AsyncSaveResponse | None`

Async save response whose `upload_completion` gates the consolidation.

---

**`model_dir`** `str`

Directory holding the sharded safetensors written by the async save.

---

**`consolidated_dir`** `str`

Output directory for the consolidated HF safetensors.

---

**`fqn_to_index_mapping`** `dict[str, int] | None`

Mapping from tensor FQN to consolidated output file index.

---

**`fqn_to_dtype_mapping`** `dict[str, str] | None`

Optional mapping from tensor FQN to original HF dtype string.

---

**`process_group`** `torch.distributed.ProcessGroup | None`

Group the consolidation collectives run on; the same one the
synchronous path and the save addons use.

---

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.async_wait() -> None
```

Wait for the async save (and any deferred consolidation) to finish.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.close() -> None
```

Close owned save resources and release the shared load-group reference.

Load metadata groups are reused across checkpointers with the same caller
group and destroyed by the final distributed-world teardown.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.finalize() -> None
```

Publish any final async checkpoint and close owned resources.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.initialize_model_weights(
    model: torch.nn.Module,
    device: torch.device,
    peft_init_method: str | None = None
) -> None
```

staticmethod

Materialize meta-device parameters and initialize model weights.

Moves empty parameter shells to the target device, resets HF initialization
flags, calls the model's weight initialization method, and initializes any
PEFT adapters.

**Parameters:**

**`model`** `torch.nn.Module`

Model whose weights should be initialized.

---

**`device`** `torch.device`

Target device for materialized parameters.

---

**`peft_init_method`** `str | None` — default: None

Initialization method for PEFT adapters (e.g. "xavier").

---

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.load_base_model(
    model: torch.nn.Module,
    device: torch.device,
    root_dir: str,
    model_name: str | None,
    load_base_model: bool = True
) -> None
```

Load a model from the base Hugging Face checkpoint in parallel.

**Parameters:**

**`model`** `torch.nn.Module`

Model to load state into

---

**`device`** `torch.device`

Device to load model onto

---

**`root_dir`** `str`

Root directory of the model cache or snapshots

---

**`model_name`** `str | None`

Name of the model or an absolute path to a snapshot

---

**`load_base_model`** `bool` — default: True

If True, restore from HF base checkpoint

---

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.load_distributed_state(
    state: typing.Any,
    state_name: str,
    path: str
) -> None
```

Load a custom stateful object previously saved with DCP.

**Parameters:**

**`state`** `Any`

Stateful object whose state dict maps keys to tensors of arbitrary shape/dtype or auxiliary
values. Tensor devices and DTensor mesh placements define the rank-local load destinations.

---

**`state_name`** `str`

Subdirectory containing the saved state.

---

**`path`** `str`

Checkpoint root directory.

---

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.load_model(
    model: torch.nn.Module,
    model_path: str,
    is_init_step: bool = False,
    use_checkpoint_id: bool = True,
    key_mapping: dict[str, str] | collections.abc.Callable[[str], str] | None = None,
    allow_checkpoint_key_subset: bool = False
) -> None
```

Load model weights from `model_path`.

Behavior:

* For PEFT (non-init): rank 0 reads `adapter_model.safetensors`, then broadcasts.
* Otherwise: use DCP with a Hugging Face or default storage reader to populate the state dict.
* If the model exposes a `state_dict_adapter`, convert to/from HF format as needed.
* For models requiring tensor merging (e.g., Mixtral), uses transformers' conversion mapping.

**Parameters:**

**`model`** `nn.Module`

Model or parallelized model parts to load into.

---

**`model_path`** `str`

Path to the model checkpoint directory or HF snapshot.

---

**`is_init_step`** `bool` — default: False

If True, treat load as initialization from a base checkpoint.

---

**`use_checkpoint_id`** `bool` — default: True

Pass `checkpoint_id` to DCP if True; disable when using direct HF paths.

---

**`key_mapping`** `dict[str, str] | Callable[[str], str] | None` — default: None

Optional regex mapping or callable renaming HF checkpoint keys.
Tensor-merging conversions require a regex mapping.

---

**`allow_checkpoint_key_subset`** `bool` — default: False

If True, keep the model's current initialization for
parameters that are absent from the checkpoint instead of requiring an exact key match.

---

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.load_on_dp_ranks(
    state: typing.Any,
    state_name: str,
    path: str
) -> None
```

Load state shared by all tensor- and pipeline-parallel peers.

This helper is intended for data-parallel-scoped state such as a
stateful dataloader. All TP/PP peers in a DP rank load the same state.

**Parameters:**

**`state`** `Any`

Stateful object to load

---

**`state_name`** `str`

Name of the stateful object

---

**`path`** `str`

Path to load stateful object

---

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.load_on_global_ranks(
    state: typing.Any,
    state_name: str,
    path: str
) -> None
```

Load state unique to this global rank, with legacy DP fallback.

**Parameters:**

**`state`** `Any`

Stateful object to load.

---

**`state_name`** `str`

Name of the stateful object.

---

**`path`** `str`

Path containing the stateful object.

---

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.load_optimizer(
    optimizer: torch.optim.Optimizer | list[torch.optim.Optimizer],
    model: torch.nn.Module | list[torch.nn.Module],
    weights_path: str,
    scheduler: typing.Any | None = None,
    optimizer_part_ids: list[int] | None = None
) -> None
```

Load optimizer (and optional scheduler) state from `weights_path/optim` using DCP.

**Parameters:**

**`optimizer`** `torch.optim.Optimizer | list[torch.optim.Optimizer]`

Optimizer or per-model-part optimizers to populate.

---

**`model`** `nn.Module | list[nn.Module]`

Model or pipeline model parts providing partitioning context.

---

**`weights_path`** `str`

Base directory for checkpoints.

---

**`scheduler`** `Any | None` — default: None

Optional LR scheduler to populate.

---

**`optimizer_part_ids`** `list[int] | None` — default: None

Global pipeline-stage indices corresponding to
per-model-part optimizers.

---

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.maybe_wait_for_staging() -> None
```

Wait for the staging to finish if it is enabled.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.save_distributed_state(
    state: typing.Any,
    state_name: str,
    path: str
) -> None
```

Save a custom stateful object through DCP on all ranks.

This is intended for auxiliary objects whose state dict contains
sharded tensors, for example BAGEL EMA shadows under FSDP2. Rank-0
`torch.save` would only persist rank 0's local shard; DCP sees the
DTensor metadata and writes all shards correctly.

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.save_model(
    model: torch.nn.Module,
    weights_path: str,
    peft_config: typing.Optional[peft.PeftConfig] = None,
    tokenizer: typing.Optional[transformers.tokenization_utils_base.PreTrainedTokenizerBase] = None,
    is_final_checkpoint: bool = False
) -> None
```

Save model weights to `weights_path/model`.

Behavior:

* PEFT: write `adapter_model.safetensors` and metadata on rank 0.
* Safetensors + consolidation: emit HF artifacts under
  `weights_path/model/consolidated` and build a consolidated index.
* Otherwise: use DCP with a Hugging Face or default storage writer to save shards.

**Parameters:**

**`model`** `nn.Module`

Model to checkpoint.

---

**`weights_path`** `str`

Base directory for checkpoints.

---

**`peft_config`** `Optional[PeftConfig]` — default: None

Optional PEFT configuration when saving adapters.

---

**`tokenizer`** `Optional[PreTrainedTokenizerBase]` — default: None

Optional tokenizer to save with consolidated artifacts.

---

**`is_final_checkpoint`** `bool` — default: False

Whether this save is the final scheduled training checkpoint.

---

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.save_on_dp_ranks(
    state: typing.Any,
    state_name: str,
    path: str
) -> None
```

Save state shared by all tensor- and pipeline-parallel peers.

This helper is intended for data-parallel-scoped state such as a
stateful dataloader. Only the TP0/PP0 peer writes each DP rank's state.

**Parameters:**

**`state`** `Any`

Stateful object to save

---

**`state_name`** `str`

Name of the stateful object

---

**`path`** `str`

Path to save stateful object

---

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.save_on_global_ranks(
    state: typing.Any,
    state_name: str,
    path: str
) -> None
```

Save state that is unique to every global process rank.

**Parameters:**

**`state`** `Any`

Stateful object to save.

---

**`state_name`** `str`

Name of the stateful object.

---

**`path`** `str`

Path to save the stateful object.

---

```python
nemo_automodel.components.checkpoint.checkpointing.Checkpointer.save_optimizer(
    optimizer: torch.optim.Optimizer | list[torch.optim.Optimizer],
    model: torch.nn.Module | list[torch.nn.Module],
    weights_path: str,
    scheduler: typing.Any | None = None,
    optimizer_part_ids: list[int] | None = None
) -> None
```

Save optimizer (and optional scheduler) state to `weights_path/optim` using DCP.

**Parameters:**

**`optimizer`** `torch.optim.Optimizer | list[torch.optim.Optimizer]`

Optimizer or per-model-part optimizers whose state will be saved.

---

**`model`** `nn.Module | list[nn.Module]`

Model or pipeline model parts providing partitioning context.

---

**`weights_path`** `str`

Base directory for checkpoints.

---

**`scheduler`** `Any | None` — default: None

Optional LR scheduler to include.

---

**`optimizer_part_ids`** `list[int] | None` — default: None

Global pipeline-stage indices corresponding to
per-model-part optimizers.

---

```python
class nemo_automodel.components.checkpoint.checkpointing._AsyncSaveContext(
    stager: typing.Any | None,
    process_group: typing.Any | None,
    future: typing.Any | None,
    staging_active: bool = False
)
```

Dataclass

Internal container for async checkpointing state.

One instance is maintained for the model save and one for the optimizer save
to keep staging/upload futures and the associated process group and stager
together in a single place.

**`future`** `Any | None`

---

**`process_group`** `Any | None`

---

**`stager`** `Any | None`

---

**`staging_active`** `bool = False`

---

```python
class nemo_automodel.components.checkpoint.checkpointing._DTensorInitCopyMode()
```

**Bases:** `TorchFunctionMode`

Adapt full-tensor initialization copies to an already-sharded destination.

```python
nemo_automodel.components.checkpoint.checkpointing._DTensorInitCopyMode.__torch_function__(
    func: collections.abc.Callable,
    types: tuple[type, ...],
    args: tuple = (),
    kwargs: dict | None = None
) -> typing.Any
```

Preserve `copy_` semantics while matching the destination's DTensor layout.

**Parameters:**

**`func`** `Callable`

Intercepted torch operation.

---

**`types`** `tuple[type, ...]`

Tensor types participating in the operation.

---

**`args`** `tuple` — default: ()

Operation arguments. For `copy_`, the destination has arbitrary
global shape and the source is broadcastable to that shape. A plain
source is privately copied, broadcast from mesh rank zero, and
locally sharded to the destination's placements before the in-place
copy; existing DTensor sources and plain destinations keep their
normal behavior.

---

**`kwargs`** `dict | None` — default: None

Keyword arguments to the operation, including an optional
`other` source tensor with the same contract as above.

---

**Returns:** `Any`

The operation result; `copy_` returns and mutates the original

```python
class nemo_automodel.components.checkpoint.checkpointing._ModelSavePlanner(
    cache_namespace: str
)
```

**Bases:** `DefaultSavePlanner`

Keep model save plans in one checkpointer-scoped cache namespace.

**`_cached_plans_key`** `= f'{cache_namespace}:model'`

---

```python
class nemo_automodel.components.checkpoint.checkpointing._OptimizerSavePlanner(
    cache_namespace: str
)
```

**Bases:** `DefaultSavePlanner`

Keep optimizer save plans in one checkpointer-scoped cache namespace.

**`_cached_plans_key`** `= f'{cache_namespace}:optimizer'`

---

```python
nemo_automodel.components.checkpoint.checkpointing._adapter_path(
    checkpoint_dir: str
) -> str
```

Return the PEFT adapter safetensors path inside a checkpoint dir (local or `msc://`).

```python
nemo_automodel.components.checkpoint.checkpointing._apply(
    module,
    fn,
    recurse = True
) -> torch.nn.Module
```

Apply a transformation function to parameters (and gradients) only.

Mirrors `nn.Module.to_empty` for parameters while skipping buffers. Respects
future flags controlling in-place vs swap behavior and safely handles
wrapper subclasses.

**Parameters:**

**`module`**

Module whose parameters are to be transformed.

---

**`fn`**

Callable applied to each parameter (and its gradient).

---

**`recurse`** — default: True

Whether to recurse into child modules.

---

**Returns:** `nn.Module`

The same module instance after transformation.

```python
nemo_automodel.components.checkpoint.checkpointing._apply_adapter_forced_dtype_mapping(
    model: torch.nn.Module,
    state_dict: dict[str, torch.Tensor],
    fqn_to_dtype_mapping: dict[str, str]
) -> dict[str, str]
```

Let model adapters override original HF dtype metadata for export-only keys.

```python
nemo_automodel.components.checkpoint.checkpointing._apply_key_mapping(
    state_dict: dict[str, torch.Tensor],
    key_mapping: dict[str, str] | collections.abc.Callable[[str], str]
) -> dict[str, torch.Tensor]
```

Rename state-dict keys using a regex mapping or a scoped rename callable.

This mirrors the renaming logic used by the DCP / HuggingFace storage
reader but operates directly on an in-memory state dict.  It is needed
when loading safetensors checkpoints outside of DCP so that HF checkpoint
keys (e.g. `language_model.model.X`) are translated to the model's
parameter FQNs (e.g. `model.language_model.X`).

**Parameters:**

**`state_dict`** `dict[str, torch.Tensor]`

Original state dict mapping keys to tensors of arbitrary shape.
Tensor layout, dtype, device, and storage are preserved.

---

**`key_mapping`** `dict[str, str] | Callable[[str], str]`

First-match `{regex_pattern: replacement}` pairs or a
callable applying the complete key conversion.

---

**Returns:** `dict[str, torch.Tensor]`

A new dict with renamed keys and the same tensor objects as the input.

```python
nemo_automodel.components.checkpoint.checkpointing._collect_global_tensor_sizes(
    state_dict: dict[str, torch.Tensor],
    process_group: torch.distributed.ProcessGroup | None
) -> dict[str, int]
```

Collect logical tensor sizes across pipeline stages without moving tensor data.

**Parameters:**

**`state_dict`** `dict[str, torch.Tensor]`

Current pipeline stage's key subset, mapping names to tensors of arbitrary shape. Each value must
report the full logical element count through `numel()`; for example, a DTensor reports its global
logical size rather than its per-rank TP/FSDP shard size. Tensors remain on their existing devices and
their placements are not changed.

---

**`process_group`** `torch.distributed.ProcessGroup | None`

Pipeline-parallel process group whose ranks collectively own the logical state dict, or
`None` for a local-only size mapping.

---

**Returns:** `dict[str, int]`

Mapping from tensor names to logical byte sizes, merged across all ranks in `process_group`.

**Raises:**

* `RuntimeError`: If a participating pipeline rank does not provide its size mapping.
* `ValueError`: If pipeline ranks report different logical sizes for the same tensor name.

```python
nemo_automodel.components.checkpoint.checkpointing._convert_checkpoint_with_transformers(
    model: torch.nn.Module,
    model_path: str,
    key_mapping: dict[str, str] | None = None
) -> dict[str, torch.Tensor] | None
```

Convert a checkpoint using transformers' conversion mapping for models that need tensor merging.

This handles MoE models like Mixtral where the checkpoint has individual expert weights
but the model uses grouped expert tensors. The transformers library's WeightConverter
operations handle the tensor merging (MergeModulelist, Concatenate).

This function converts the state dict WITHOUT loading it into the model, so it can be
used with FSDP-aware loading mechanisms.

**Parameters:**

**`model`** `nn.Module`

The model (used to get conversion mapping and target keys).

---

**`model_path`** `str`

Path to the HuggingFace checkpoint directory.

---

**`key_mapping`** `dict[str, str] | None` — default: None

Optional additional key mapping.

---

**Returns:** `dict[str, torch.Tensor] | None`

Converted state dict ready for loading, or None if conversion failed.

```python
nemo_automodel.components.checkpoint.checkpointing._create_dirs(
    dirs: str | None = ()
) -> None
```

Create local directory paths and ignore cloud paths.

```python
nemo_automodel.components.checkpoint.checkpointing._divide_keys_by_size(
    keys: list[str],
    state_dict: dict[str, torch.Tensor],
    target_shard_bytes: int,
    key_size_mapping: dict[str, int] | None = None
) -> dict[str, int]
```

Assign keys to deterministic size-based shards.

**Parameters:**

**`keys`** `list[str]`

Ordered tensor names to assign.

---

**`state_dict`** `dict[str, torch.Tensor]`

Mapping of tensor names to tensors of arbitrary shape. Each value represents its full logical
tensor, including when its per-rank storage is sharded, and is read only for its logical byte size when
`key_size_mapping` is not provided.

---

**`target_shard_bytes`** `int`

Positive target size for each shard in bytes.

---

**`key_size_mapping`** `dict[str, int] | None` — default: None

Optional mapping of tensor names to logical byte sizes, including tensors not present in
the rank-local `state_dict`.

---

**Returns:** `dict[str, int]`

Mapping from every input key to a positive, one-based shard index.

**Raises:**

* `ValueError`: If `target_shard_bytes` is not positive.

```python
nemo_automodel.components.checkpoint.checkpointing._ensure_dirs(
    dirs: str | None = (),
    process_group: torch.distributed.ProcessGroup | None = None
) -> None
```

Create directories on all ranks and synchronize across ranks.

**Parameters:**

**`*dirs`** `str | None` — default: ()

One or more directory paths that should exist.

---

**`process_group`** `torch.distributed.ProcessGroup | None` — default: None

Ranks that must observe the directories before continuing.

---

```python
nemo_automodel.components.checkpoint.checkpointing._ensure_msc_available() -> None
```

Raise an error if MSC is not installed but a cloud path is used.

```python
nemo_automodel.components.checkpoint.checkpointing._ensure_shared_dirs(
    dirs: str | None = (),
    process_group: torch.distributed.ProcessGroup | None = None
) -> None
```

Create shared DCP directories on group rank zero and synchronize the group.

Unlike auxiliary per-rank state, DCP checkpoint directories must be visible
to every rank through the same filesystem.

**Parameters:**

**`*dirs`** `str | None` — default: ()

One or more shared directory paths that should exist.

---

**`process_group`** `torch.distributed.ProcessGroup | None` — default: None

Ranks that must observe the directories before continuing.

---

```python
nemo_automodel.components.checkpoint.checkpointing._equally_divide_layers(
    num_shards: int,
    keys: list[str]
) -> dict[str, int]
```

Equally divide the state dict keys into num\_shards shards.

```python
nemo_automodel.components.checkpoint.checkpointing._format_restricted_load_error(
    f: torch.serialization.FileLike
) -> str
```

```python
nemo_automodel.components.checkpoint.checkpointing._get_checkpoint_metadata(
    path: str,
    storage_reader: torch.distributed.checkpoint.storage.StorageReader | None = None
) -> torch.distributed.checkpoint.metadata.Metadata
```

Read checkpoint metadata, including saved tensor sizes and dtypes.

```python
nemo_automodel.components.checkpoint.checkpointing._get_hf_safetensors_reference_path(
    cache_dir: str | pathlib.Path | None,
    repo_id: str | None
) -> str | None
```

Return the local HF safetensors reference directory for a model.

Prefer the snapshot directory containing `model.safetensors.index.json` for
sharded checkpoints. If no index exists but a snapshot directory is present,
return that directory as the single-file safetensors reference path. Return
None when `repo_id` is None or the repo has no cached snapshot directory.

For example, if the located file is

/opt/models/models--meta-llama--Llama-3.2-3B/snapshots/13afe.../model.safetensors.index.json

this function will return the directory path

/opt/models/models--meta-llama--Llama-3.2-3B/snapshots/13afe...

This will error if the model hasn't been downloaded or if the cache directory is incorrect.

**Parameters:**

**`cache_dir`** `str | Path | None`

Path to cache directory

---

**`repo_id`** `str | None`

Hugging Face repository ID

---

**Returns:** `str | None`

Path to the snapshot/model directory containing safetensors weights, or

```python
nemo_automodel.components.checkpoint.checkpointing._get_original_hf_index_total_size(
    config: nemo_automodel.components.checkpoint.config.CheckpointingConfig
) -> int | None
```

Return the original HF safetensors index total size, if available.

```python
nemo_automodel.components.checkpoint.checkpointing._get_shared_parameter_names(
    model_parts: list[torch.nn.Module]
) -> list[list[str]]
```

Find checkpoint names referring to the same live parameter.

This primarily safeguards HF initialization for encoder-decoder models
sharing embeddings across both stacks. ModelState already handles ordinary
input-embedding/LM-head tying; safetensors may omit additional aliases that
the loader must restore without accepting genuinely missing parameters.

**Parameters:**

**`model_parts`** `list[nn.Module]`

Model or pipeline parts after state-dict normalization.

---

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

Groups of canonical names sharing one parameter object. Equal-valued

```python
nemo_automodel.components.checkpoint.checkpointing._init_peft_adapters(
    model: torch.nn.Module,
    peft_init_method: str
) -> None
```

Initialize the PEFT adapters with the scaled weights.

**Parameters:**

**`model`** `nn.Module`

Model to initialize PEFT adapters for

---

**`peft_init_method`** `str`

Method to initialize PEFT adapters e.g. "xavier". See `LinearLoRA` for more details.

---

```python
nemo_automodel.components.checkpoint.checkpointing._is_bin_checkpoint(
    path: str
) -> bool
```

Return True if path looks like a PyTorch .bin checkpoint.

```python
nemo_automodel.components.checkpoint.checkpointing._is_custom_model(
    module: torch.nn.Module
) -> bool
```

True if the model has a custom implementation in nemo\_automodel/components/models/.

The generic HFCheckpointingMixin (in .common.hf\_checkpointing\_mixin) is
injected into every model by \_get\_mixin\_wrapped\_class and does NOT count
as a "custom model".  Only actual model implementations (e.g. llama,
deepseek\_v3) that live under nemo\_automodel.components.models qualify.

```python
nemo_automodel.components.checkpoint.checkpointing._is_model_checkpoint_path(
    path: str
) -> bool
```

Return whether a checkpoint path identifies model weights.

```python
nemo_automodel.components.checkpoint.checkpointing._is_safetensors_checkpoint(
    path: str
) -> bool
```

Return True if path looks like a safetensors checkpoint (so we can preserve dtype); else DCP or other.

```python
nemo_automodel.components.checkpoint.checkpointing._load_full_state_dict_into_model(
    model_parts: list[torch.nn.Module],
    state_dict: dict[str, torch.Tensor]
) -> None
```

Load a full (non-sharded) state dict into a potentially FSDP-wrapped model.

Every rank must supply the **full** state dict.  PyTorch's
`set_model_state_dict` with `full_state_dict=True` (but **not**
`broadcast_from_rank0`) calls `_distribute_state_dict` which lets
each rank independently slice its local DTensor shard from the full
tensor -- no NCCL collectives are needed.

We intentionally avoid `broadcast_from_rank0=True` because it
introduces an asymmetric workload: rank 0 does a synchronous CPU→GPU
copy (`.to(device)`) per tensor while other ranks only do
`torch.empty` (async allocation).  The non-src ranks race ahead
enqueuing hundreds of NCCL broadcasts that rank 0 cannot keep up with,
leading to a 60 s NCCL watchdog timeout.

After loading, floating-point parameters are converted to match the
checkpoint dtype.  PyTorch's `set_model_state_dict` uses *copy*
semantics (`assign=False`) for non-meta parameters, which preserves
the model's initialisation dtype instead of the checkpoint dtype.
The post-load fixup ensures the safetensors dtype (e.g. bf16) is
honoured.

**Parameters:**

**`model_parts`** `list[nn.Module]`

List of model parts (for pipeline parallelism)

---

**`state_dict`** `dict[str, torch.Tensor]`

Full state dict with regular tensors.  Must be
populated on **every** rank (not just rank 0).

---

```python
nemo_automodel.components.checkpoint.checkpointing._load_hf_bin_checkpoint(
    model_path: str
) -> dict[str, torch.Tensor] | None
```

Load a HuggingFace .bin checkpoint into a state dict.

Handles single-file (pytorch\_model.bin), sharded (pytorch\_model.bin.index.json),
and glob fallback (\*.bin) layouts.
Returns None if no .bin files are found.

**Parameters:**

**`model_path`** `str`

Path to checkpoint file or directory.

---

```python
nemo_automodel.components.checkpoint.checkpointing._load_hf_checkpoint_preserving_dtype(
    model_path: str,
    prefault_safetensors: bool = False
) -> dict[str, torch.Tensor] | None
```

Load a HuggingFace checkpoint into a new state dict so tensor dtypes
match the checkpoint (e.g. bf16). Used when loading the base model so FSDP sees
uniform dtype instead of the model's init dtypes (e.g. float32).
Prefers safetensors but falls back to .bin files.
Returns None if no loadable checkpoint is found.

**Parameters:**

**`model_path`** `str`

Path to checkpoint file or directory.

---

```python
nemo_automodel.components.checkpoint.checkpointing._load_hf_safetensors_checkpoint(
    model_path: str,
    prefault_mmap: bool = False
) -> dict[str, torch.Tensor] | None
```

Load a safetensors checkpoint into a state dict.

On integrated CUDA systems, `prefault_mmap` reads every file-backed tensor
once before installation. This keeps the returned tensors mmap-backed and
reclaimable while avoiding page-by-page migration during the later CUDA copy.

```python
nemo_automodel.components.checkpoint.checkpointing._load_safetensors(
    path: str
) -> dict[str, torch.Tensor]
```

Read a safetensors file from a local path or an `msc://` cloud path.

```python
nemo_automodel.components.checkpoint.checkpointing._materialize_to_hf_views_for_save(
    state_dict: dict[str, torch.Tensor]
) -> None
```

Replace non-contiguous tensor values in `state_dict` with contiguous copies in place.

MoE adapters return non-contiguous strided views into the model's grouped
expert storage for the optimized load path; `safetensors.torch.save`
(which the DCP HF storage writer calls) rejects non-contiguous tensors,
so we materialize one tensor at a time here with `empty_cache` between
iterations. Per-tensor transient is bounded to a single expert weight
instead of allocating the full grouped set up front.

```python
nemo_automodel.components.checkpoint.checkpointing._maybe_adapt_state_dict_from_hf(
    model_part: torch.nn.Module,
    state_dict: dict[str, torch.Tensor],
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    paramwrapper_layout_hint: str | None = None
) -> dict[str, torch.Tensor]
```

Custom models use state dict adapters to convert the state dict from the Hugging Face format to the native format.

`paramwrapper_layout_hint` carries the fused expert LoRA layout recorded in
the checkpoint's automodel\_peft\_config.json (see \_read\_paramwrapper\_layout\_metadata),
so the adapter resolves the peft ParamWrapper layout from metadata instead of shapes.

```python
nemo_automodel.components.checkpoint.checkpointing._maybe_adapt_state_dict_to_hf(
    model_part: torch.nn.Module,
    state_dict: dict[str, torch.Tensor],
    quantization: bool = False,
    kwargs = {}
) -> dict[str, torch.Tensor]
```

Custom models use state dict adapters to convert the state dict to the Hugging Face format.

```python
nemo_automodel.components.checkpoint.checkpointing._maybe_msc_reader(
    path: str,
    storage_reader: torch.distributed.checkpoint.storage.StorageReader | None
) -> torch.distributed.checkpoint.storage.StorageReader | None
```

Return an MSC filesystem reader for `msc://` paths, else the given reader.

```python
nemo_automodel.components.checkpoint.checkpointing._maybe_msc_writer(
    path: str,
    storage_writer: torch.distributed.checkpoint.storage.StorageWriter | None
) -> torch.distributed.checkpoint.storage.StorageWriter | None
```

Return an MSC filesystem writer for `msc://` paths, else the given writer.

```python
nemo_automodel.components.checkpoint.checkpointing._model_has_dtensors(
    module: torch.nn.Module
) -> bool
```

True if any parameter is a DTensor (model is already sharded).

```python
nemo_automodel.components.checkpoint.checkpointing._new_gloo_process_group(
    process_group: torch.distributed.ProcessGroup | None,
    timeout: datetime.timedelta | None = None
) -> torch.distributed.ProcessGroup
```

Create a Gloo group with the same membership as `process_group`.

**Parameters:**

**`process_group`** `torch.distributed.ProcessGroup | None`

Source process group whose membership should be preserved.

---

**`timeout`** `timedelta | None` — default: None

Optional timeout for operations executed on the new group.

---

**Returns:** `torch.distributed.ProcessGroup`

The newly created Gloo process group.

```python
nemo_automodel.components.checkpoint.checkpointing._normalize_dtype_mapping_to_state_dict_keys(
    fqn_to_dtype_mapping: dict[str, str],
    state_dict_keys: list[str],
    base_model_prefix: str | None = None
) -> dict[str, str]
```

Align original HF dtype metadata with the keys that will be exported.

```python
nemo_automodel.components.checkpoint.checkpointing._read_paramwrapper_layout_metadata(
    model_path: str | os.PathLike
) -> str | None
```

Return the fused expert LoRA layout stamped into a PEFT checkpoint, if any.

The PEFT save path records which peft ParamWrapper layout the adapter was
exported in (peft flipped it in 0.19.1, huggingface/peft#3165) inside
automodel\_peft\_config.json. Absent or unstamped checkpoints return None and
the adapter falls back to shape detection.

```python
nemo_automodel.components.checkpoint.checkpointing._reinit_non_persistent_buffers(
    model: torch.nn.Module,
    device: torch.device,
    model_type: str | None = None
) -> None
```

Recompute non-persistent buffers that are not saved in checkpoints.

Non-persistent buffers are not saved in checkpoints, so after meta-device
materialization they contain uninitialized CUDA memory.  When
`initialize_weights()` is skipped (e.g. for Gemma3 to avoid DTensor
issues), these buffers must be recomputed explicitly.

Only runs for models listed in `_MODELS_REQUIRING_BUFFER_REINIT` to
avoid unexpected side-effects on arbitrary HF Hub models.

Handles four patterns:

1. **Standard RoPE** — single `inv_freq` buffer with `rope_init_fn` and
   optional legacy `rope_kwargs` (e.g. Nemotron-NAS, Ling).
2. **Per-layer-type RoPE** — `{layer_type}_inv_freq` buffers via
   `compute_default_rope_parameters` (e.g. Gemma3RotaryEmbedding).
3. **Scaled embedding** — `embed_scale` buffer on `ScaledWordEmbedding`
   modules (Gemma family), recomputed from `scalar_embed_scale`.
4. **Vision position IDs** — `position_ids` buffer on vision embedding
   modules (SigLIP), recomputed from `num_positions`.

**Parameters:**

**`model`** `nn.Module`

Model to reinitialize non-persistent buffers for.

---

**`device`** `torch.device`

Device to create the new buffers on.

---

**`model_type`** `str | None` — default: None

The `config.model_type` string.  If not in
`_MODELS_REQUIRING_BUFFER_REINIT` the function is a no-op.

---

```python
nemo_automodel.components.checkpoint.checkpointing._save_safetensors(
    state_dict: dict[str, torch.Tensor],
    path: str
) -> None
```

Write a safetensors file to a local path or an `msc://` cloud path.

For cloud paths the tensors are serialized to bytes and streamed to the MSC
file handle, since `save_file` only accepts a local filesystem path.

```python
nemo_automodel.components.checkpoint.checkpointing._should_dequantize_base_checkpoint(
    model: torch.nn.Module,
    requested: bool | None
) -> bool
```

Return whether this load requires checkpoint dequantization.

`requested` permits dequantization unless it is explicitly `False`.
The conversion is needed only when the source model config declares a
quantization method; a stale `True` setting must not route BF16 weights
through the quantized full-CPU loading path.

**Parameters:**

**`model`** `nn.Module`

Model whose source checkpoint metadata is being loaded.

---

**`requested`** `bool | None`

Configured dequantization preference.

---

**Returns:** `bool`

Whether the source checkpoint declares quantized weights that should

```python
nemo_automodel.components.checkpoint.checkpointing._should_write_consolidated_safetensors(
    config: nemo_automodel.components.checkpoint.config.CheckpointingConfig,
    is_final_checkpoint: bool = False
) -> bool
```

Whether to output consolidated HF weights along with sharded weights.

```python
nemo_automodel.components.checkpoint.checkpointing._should_write_hf_metadata(
    config: nemo_automodel.components.checkpoint.config.CheckpointingConfig
) -> bool
```

Whether to write HF metadata/artifacts for a checkpoint.

```python
nemo_automodel.components.checkpoint.checkpointing._summarize_state_dict_key_diff(
    expected_keys: set[str],
    loaded_keys: set[str],
    limit: int = 10
) -> dict[str, typing.Any]
```

Summarize state-dict key mismatches for checkpoint load diagnostics.

```python
nemo_automodel.components.checkpoint.checkpointing._unwrap_ddp_model(
    model: torch.nn.Module
) -> torch.nn.Module
```

Return the module that owns model metadata hidden by DDP.

```python
nemo_automodel.components.checkpoint.checkpointing._warn_if_inline_consolidation_enabled(
    config: nemo_automodel.components.checkpoint.config.CheckpointingConfig
) -> None
```

Educate users about the cost of inline HF consolidation.

```python
nemo_automodel.components.checkpoint.checkpointing._warn_if_large_inline_consolidation(
    config: nemo_automodel.components.checkpoint.config.CheckpointingConfig,
    state_dict: dict[str, torch.Tensor],
    fqn_to_index_mapping: dict[str, int] | None,
    is_final_checkpoint: bool = False
) -> None
```

Warn when inline consolidated export is large enough to waste GPU allocation time.

```python
nemo_automodel.components.checkpoint.checkpointing.load_torch_ckpt(
    f: torch.serialization.FileLike,
    map_location: torch.serialization.MAP_LOCATION = None,
    pickle_module: typing.Any = None,
    weights_only: bool | None = None,
    mmap: bool | None = None,
    pickle_load_args: typing.Any = {}
) -> typing.Any
```

Load a torch checkpoint with restricted unpickling by default.

**Parameters:**

**`f`** `FileLike`

File path or binary file object accepted by `torch.load`.

---

**`map_location`** `MAP_LOCATION` — default: None

Device remapping accepted by `torch.load`.

---

**`pickle_module`** `Any` — default: None

Module used to unpickle metadata and objects.

---

**`weights_only`** `bool | None` — default: None

When `False`, explicitly opt into unrestricted pickle loading.
`None` and `True` use restricted loading.

---

**`mmap`** `bool | None` — default: None

Whether to memory-map tensor storages from a file path.

---

**`**pickle_load_args`** `Any` — default: \{}

Additional arguments forwarded to the unpickler.

---

**Returns:** `Any`

The deserialized checkpoint.

**Raises:**

* `RuntimeError`: If restricted loading rejects the artifact.

```python
nemo_automodel.components.checkpoint.checkpointing.save_config(
    config: dict[str, typing.Any],
    weights_path: str
) -> None
```

Save a config to a weights path.

**Parameters:**

**`config`** `dict[str, Any]`

Config to save

---

**`weights_path`** `str`

Path to save config

---

```python
nemo_automodel.components.checkpoint.checkpointing.save_losses(
    losses: dict[str, typing.Any],
    weights_path: str
) -> None
```

Write checkpoint loss metadata to `weights_path/losses.json`.

Mirrors `save_config` so the file lands in the checkpoint directory for
both local and `msc://` roots. Every failure is logged rather than raised,
including a missing `multistorageclient`: this is metadata written on rank 0
only, so raising would strand the other ranks in the next collective.

**Parameters:**

**`losses`** `dict[str, Any]`

Loss values to record. Values must be JSON-serializable.

---

**`weights_path`** `str`

Checkpoint directory.

---

```python
nemo_automodel.components.checkpoint.checkpointing.to_empty_parameters_only(
    model: torch.nn.Module,
    device: torch.device,
    recurse: bool = True,
    dtype: torch.dtype | None = None
) -> torch.nn.Module
```

Move parameters to the specified device without copying storage, skipping buffers.

Mirrors torch.nn.Module.to\_empty but applies only to parameters, not buffers.

**Parameters:**

**`model`** `nn.Module`

The module to transform

---

**`device`** `torch.device`

Target device

---

**`recurse`** `bool` — default: True

Whether to recurse into child modules

---

**Returns:** `nn.Module`

The same module instance

```python
nemo_automodel.components.checkpoint.checkpointing.MSC_AVAILABLE = True
```

```python
nemo_automodel.components.checkpoint.checkpointing._CONSOLIDATED_SIZE_WARNING_THRESHOLD_BYTES = 50 * 1024 ** 3
```

```python
nemo_automodel.components.checkpoint.checkpointing._DEFAULT_HF_CONSOLIDATED_SHARD_SIZE_BYTES = 5 * 1024 ** 3
```

```python
nemo_automodel.components.checkpoint.checkpointing._MODELS_REQUIRING_BUFFER_REINIT: frozenset[str] = frozenset({'bailing_moe', 'gemma3', 'nemotron-nas'})
```

```python
nemo_automodel.components.checkpoint.checkpointing._load_metadata_groups: WeakKeyDictionary[ProcessGroup, ProcessGroup] = WeakKeyDictionary()
```

```python
nemo_automodel.components.checkpoint.checkpointing.load_hf_safetensors_state_dict = _load_hf_safetensors_checkpoint
```

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