> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/checkpoint/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 ## Subpackages * **[`nemo_automodel.components.checkpoint._backports`](/nemo-automodel/nemo_automodel/components/checkpoint/_backports)** ## Submodules * **[`nemo_automodel.components.checkpoint._torch_backports`](/nemo-automodel/nemo_automodel/components/checkpoint/_torch_backports)** * **[`nemo_automodel.components.checkpoint.addons`](/nemo-automodel/nemo_automodel/components/checkpoint/addons)** * **[`nemo_automodel.components.checkpoint.checkpointing`](/nemo-automodel/nemo_automodel/components/checkpoint/checkpointing)** * **[`nemo_automodel.components.checkpoint.config`](/nemo-automodel/nemo_automodel/components/checkpoint/config)** * **[`nemo_automodel.components.checkpoint.conversion_mapping`](/nemo-automodel/nemo_automodel/components/checkpoint/conversion_mapping)** * **[`nemo_automodel.components.checkpoint.lifecycle`](/nemo-automodel/nemo_automodel/components/checkpoint/lifecycle)** * **[`nemo_automodel.components.checkpoint.state_dict_adapter`](/nemo-automodel/nemo_automodel/components/checkpoint/state_dict_adapter)** * **[`nemo_automodel.components.checkpoint.stateful_wrappers`](/nemo-automodel/nemo_automodel/components/checkpoint/stateful_wrappers)** * **[`nemo_automodel.components.checkpoint.utils`](/nemo-automodel/nemo_automodel/components/checkpoint/utils)** ## Package Contents ### Classes | Name | Description | | ----------------------------------------------------------------------------------------- | -------------------------------- | | [`CheckpointingConfig`](#nemo_automodel-components-checkpoint-config-CheckpointingConfig) | Configuration for checkpointing. | ### API ```python class nemo_automodel.components.checkpoint.config.CheckpointingConfig( enabled: bool = True, checkpoint_dir: str | pathlib.Path = 'checkpoints/', model_save_format: str = 'safetensors', model_cache_dir: str | pathlib.Path | None = None, model_repo_id: str | None = None, save_consolidated: bool | str | nemo_automodel.components.checkpoint.config.SaveConsolidatedMode = 'final', is_peft: bool = False, model_state_dict_keys: list[str] | None = None, is_async: bool = False, wait_for_staging: bool = False, cpu_offload: bool = False, load_metadata_with_gloo: bool = False, allow_legacy_pickle_restore: bool = False, dequantize_base_checkpoint: bool | None = None, original_model_root_dir: str | None = None, skip_task_head_prefixes_for_base_model: list[str] | None = None, single_rank_consolidation: bool = False, staging_dir: str | None = None, v4_compatible: bool = False, legacy_paramwrapper_layout: bool = False, diffusers_compatible: bool = False, best_metric_key: str = 'default', consolidation_timeout_minutes: int = 30, max_recent_checkpoints: int | None = None ) ``` Dataclass Configuration for checkpointing. Every field has a default so the recipe layer can construct this directly from the YAML `checkpoint:` block merged with the model-derived `model_repo_id` / `model_cache_dir` / `is_peft` values. When `model_cache_dir` is `None` it falls back to the HF hub cache. **`allow_legacy_pickle_restore`** `bool = False` --- **`best_metric_key`** `str = 'default'` --- **`checkpoint_dir`** `str | Path = 'checkpoints/'` --- **`consolidation_timeout_minutes`** `int = 30` --- **`cpu_offload`** `bool = False` --- **`dequantize_base_checkpoint`** `bool | None = None` --- **`diffusers_compatible`** `bool = False` --- **`enabled`** `bool = True` --- **`is_async`** `bool = False` --- **`is_peft`** `bool = False` --- **`legacy_paramwrapper_layout`** `bool = False` --- **`load_metadata_with_gloo`** `bool = False` --- **`max_recent_checkpoints`** `int | None = None` --- **`model_cache_dir`** `str | Path | None = None` --- **`model_repo_id`** `str | None = None` --- **`model_save_format`** `str = 'safetensors'` --- **`model_state_dict_keys`** `list[str] | None = None` --- **`original_model_root_dir`** `str | None = None` --- **`save_consolidated`** `bool | str | SaveConsolidatedMode = 'final'` --- **`single_rank_consolidation`** `bool = False` --- **`skip_task_head_prefixes_for_base_model`** `list[str] | None = None` --- **`staging_dir`** `str | None = None` --- **`v4_compatible`** `bool = False` --- **`wait_for_staging`** `bool = False` --- ```python nemo_automodel.components.checkpoint.config.CheckpointingConfig.__post_init__() ``` Resolve the cache dir, enforce PEFT constraints, and coerce the save format/mode. ```python nemo_automodel.components.checkpoint.config.CheckpointingConfig.build( 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: torch.distributed.ProcessGroup | None = None ) -> nemo_automodel.components.checkpoint.checkpointing.Checkpointer ``` Build the `Checkpointer` engine for this config. `Checkpointer` is imported lazily to avoid a circular import (`checkpointing.py` imports `CheckpointingConfig` from this module) and to keep the heavy DCP/safetensors deps out of module load. **Parameters:** **`dp_rank`** `int` Data-parallel rank. --- **`tp_rank`** `int` Tensor-parallel rank. --- **`pp_rank`** `int` Pipeline-parallel rank. --- **`moe_mesh`** `DeviceMesh | None` — default: None Optional device mesh for MoE checkpointing. --- **`process_group`** `ProcessGroup | None` — default: None Process group used for distributed checkpoint collectives. --- **`pp_group`** `ProcessGroup | None` — default: None Optional pipeline-parallel process group. Threaded to the PEFT save path so adapter weights are gathered across PP stages (required for complete adapters when `pp_size > 1`). --- **Returns:** `Checkpointer` class:`Checkpointer`. ## Docs - [nemo_automodel.components.checkpoint.addons](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/checkpoint/addons.md) - [nemo_automodel.components.checkpoint.checkpointing](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/checkpoint/checkpointing.md) - [nemo_automodel.components.checkpoint.config](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/checkpoint/config.md) - [nemo_automodel.components.checkpoint.conversion_mapping](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/checkpoint/conversion_mapping.md) - [nemo_automodel.components.checkpoint.lifecycle](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/checkpoint/lifecycle.md) - [nemo_automodel.components.checkpoint.state_dict_adapter](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/checkpoint/state_dict_adapter.md) - [nemo_automodel.components.checkpoint.stateful_wrappers](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/checkpoint/stateful_wrappers.md) - [nemo_automodel.components.checkpoint.utils](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/checkpoint/utils.md)