nemo_automodel.components.checkpoint.checkpointing

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-aware helpers for saving/loading auxiliary state and utilities to initialize from a base HF checkpoint.

nemo_automodel.components.checkpoint.checkpointing.Checkpointer._do_load(
    state_dict: dict[str, torch.Tensor],
    path: str,
    storage_reader: typing.Optional[nemo_automodel.components.checkpoint._backports.hf_storage._HuggingFaceStorageReader] = None,
    is_init_step: bool = False
) -> dict[str, torch.Tensor]

nemo_automodel.components.checkpoint.checkpointing.Checkpointer._do_save(
    state_dict: dict[str, torch.Tensor],
    path: str,
    storage_writer: typing.Optional[nemo_automodel.components.checkpoint._backports.hf_storage._HuggingFaceStorageWriter] = None
) -> typing.Optional[torch.distributed.checkpoint.state_dict_saver.AsyncSaveResponse]

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

nemo_automodel.components.checkpoint.checkpointing.Checkpointer._get_storage_reader(
    model_path: str,
    key_mapping: typing.Optional[dict[str, str]],
    is_init_step: bool = False,
    is_safetensors: bool | None = None
) -> typing.Optional[nemo_automodel.components.checkpoint._backports.hf_storage._HuggingFaceStorageReader]

nemo_automodel.components.checkpoint.checkpointing.Checkpointer._get_storage_writer(
    consolidated_output_path: typing.Optional[str],
    fqn_to_index_mapping: typing.Optional[dict[str, int]],
    fqn_to_dtype_mapping: typing.Optional[dict[str, str]],
    model_path: str,
    consolidate_on_all_ranks: bool = False
) -> typing.Optional[nemo_automodel.components.checkpoint._backports.hf_storage._HuggingFaceStorageWriter]

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]
) -> typing.Optional[dict[str, int]]

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]
) -> typing.Optional[dict[str, str]]

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

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

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

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

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

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

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

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: typing.Optional[dict[str, str]] = None,
    allow_checkpoint_key_subset: bool = False
) -> None

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

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

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

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

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

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

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

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.

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

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:

Returns: nn.Module

The same module instance after transformation.

Rename state-dict keys using regex-based key_mapping.

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:

Returns: dict[str, torch.Tensor]

A new dict with renamed keys.

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:

Returns: Optional[dict[str, torch.Tensor]]

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

Assign keys to deterministic size-based shards.

Create directories on all ranks and synchronize across ranks.

Parameters:

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

Equally divide the state dict keys into num_shards shards.

Return checkpoint FQNs present in metadata.

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:

Returns: str | None

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

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

Initialize the PEFT adapters with the scaled weights.

Parameters:

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

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.

True if the model was loaded with trust_remote_code (HF dynamic modules).

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

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:

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:

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:

Load a safetensors checkpoint into a state dict.

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

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.

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

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

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

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

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

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

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 + rope_kwargs (e.g. Nemotron-NAS).
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:

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.

Whether to output consolidated HF weights along with sharded weights.

Whether to write HF metadata/artifacts for a checkpoint.

Summarize state-dict key mismatches for checkpoint load diagnostics.

Educate users about the cost of inline HF consolidation.

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

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

Save a config to a weights path.

Parameters:

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:

Returns: nn.Module

The same module instance