nemo_automodel.components.checkpoint

Subpackages

nemo_automodel.components.checkpoint._backports

Submodules

Package Contents

Classes

Name	Description
`CheckpointingConfig`	Configuration for checkpointing.

API

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

nemo_automodel.components.checkpoint.config.CheckpointingConfig.__post_init__()

Resolve the cache dir, enforce PEFT constraints, and coerce the save format/mode.

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 | NoneDefaults to None

Optional device mesh for MoE checkpointing.

process_group

ProcessGroup | NoneDefaults to None

Process group used for distributed checkpoint collectives.

pp_group

ProcessGroup | NoneDefaults to 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.