nemo_automodel.components.checkpoint.config

Public config surface for the checkpoint component.

CheckpointingConfig holds the typed parameters that drive checkpointing behaviour and exposes .build() to construct the Checkpointer engine (defined in checkpointing.py). Every field has a sensible default so the recipe layer can construct it directly from the YAML checkpoint: block plus the model-derived model_repo_id / model_cache_dir / is_peft arguments — there is no separate builder/adapter.

Module Contents

Classes

Name	Description
`CheckpointingConfig`	Configuration for checkpointing.
`SaveConsolidatedMode`	Controls when consolidated HF safetensors are exported.

Functions

Name	Description
`_is_geq_torch_2_9`	Check if the current torch version is greater than or equal to 2.9.0.
`_is_leq_torch_2_7_1`	Check if the current torch version is less than or equal to 2.7.1.
`_normalize_save_consolidated`	Normalize legacy bools and string aliases to a consolidated export mode.

Data

_TORCH_2_7_1

_TORCH_2_9

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.

class nemo_automodel.components.checkpoint.config.SaveConsolidatedMode

Bases: enum.Enum

Controls when consolidated HF safetensors are exported.

EVERY

= 'every'

FALSE

= 'false'

FINAL

= 'final'

nemo_automodel.components.checkpoint.config._is_geq_torch_2_9() -> bool

Check if the current torch version is greater than or equal to 2.9.0.

nemo_automodel.components.checkpoint.config._is_leq_torch_2_7_1() -> bool

Check if the current torch version is less than or equal to 2.7.1.

nemo_automodel.components.checkpoint.config._normalize_save_consolidated(
    value: bool | str | nemo_automodel.components.checkpoint.config.SaveConsolidatedMode
) -> nemo_automodel.components.checkpoint.config.SaveConsolidatedMode

Normalize legacy bools and string aliases to a consolidated export mode.

nemo_automodel.components.checkpoint.config._TORCH_2_7_1 = (2, 7, 1)

nemo_automodel.components.checkpoint.config._TORCH_2_9 = (2, 9)