nemo_automodel.components.checkpoint.addons

Module Contents

Classes

Name	Description
`CheckpointAddon`	Optional hooks that run around backend IO (used for PEFT and consolidated HF metadata).
`ConsolidatedHFAddon`	Addon that writes consolidated Hugging Face metadata alongside sharded weights.
`PeftAddon`	Addon that writes PEFT-specific metadata and tokenizer alongside adapter weights.
`_ConsolidatedHFMetadataExporter`	Model-provided writer for consolidated Hugging Face metadata.

Functions

Name	Description
`_apply_transformers_compat_guards`	Guard imports of transformers symbols removed in newer versions.
`_config_exists`	-
`_extract_target_modules`	Extract the target modules from the model used by LoRA/PEFT layers.
`_extract_target_parameters`	Extract `target_parameters` for PEFT v0.18+ ParamWrapper format.
`_get_automodel_peft_metadata`	Get the PEFT metadata in the format expected by Automodel.
`_get_hf_peft_config`	Get the minimal PEFT config in the format expected by Hugging Face.
`_group_barrier`	-
`_is_group_rank_0`	-
`_iter_custom_code_classes`	Yield classes referenced by `config.auto_map` (and the model’s own class).
`_maybe_save_custom_model_code`	Save the custom model code if it exists. This function preserves the original directory structure.
`_maybe_strip_quantization_config`	Remove `quantization_config` from the HF config when no parameters are quantized.
`_save_generated_hf_assets`	Run the existing generated Hugging Face metadata export path.
`_save_original_config_json`	Copy the original pretrained `config.json` with `quantization_config` stripped.
`_unwrap_ddp_model`	Return the module that owns export metadata hidden by DDP.

Data

_REMOVED_TRANSFORMERS_SYMBOLS

API

class nemo_automodel.components.checkpoint.addons.CheckpointAddon()

Protocol

Optional hooks that run around backend IO (used for PEFT and consolidated HF metadata).

nemo_automodel.components.checkpoint.addons.CheckpointAddon.post_save(
    kwargs = {}
) -> None

nemo_automodel.components.checkpoint.addons.CheckpointAddon.pre_save(
    kwargs = {}
) -> None

class nemo_automodel.components.checkpoint.addons.ConsolidatedHFAddon()

Addon that writes consolidated Hugging Face metadata alongside sharded weights.

Models can provide a custom consolidated metadata exporter; all other models retain the generated config, custom-code, and tokenizer path. Rank 0 writes the artifacts, then synchronizes ranks.

nemo_automodel.components.checkpoint.addons.ConsolidatedHFAddon.post_save(
    kwargs = {}
) -> None

Copy the saved HF metadata to the consolidated directory.

The reason we keep it this way is because the HF metadata needs to stay available for offline consolidation and re-export, otherwise any changes made to the config during training will be lost.

nemo_automodel.components.checkpoint.addons.ConsolidatedHFAddon.pre_save(
    kwargs = {}
) -> None

Pre-save hook to emit consolidated HF artifacts.

class nemo_automodel.components.checkpoint.addons.PeftAddon()

Addon that writes PEFT-specific metadata and tokenizer alongside adapter weights.

On rank 0, this saves adapter_config.json, automodel_peft_config.json, the tokenizer (if provided), and synchronizes all ranks afterward.

nemo_automodel.components.checkpoint.addons.PeftAddon.post_save(
    kwargs = {}
) -> None

nemo_automodel.components.checkpoint.addons.PeftAddon.pre_save(
    kwargs = {}
) -> None

Pre-save hook to emit PEFT artifacts.

class nemo_automodel.components.checkpoint.addons._ConsolidatedHFMetadataExporter()

Protocol

Model-provided writer for consolidated Hugging Face metadata.

nemo_automodel.components.checkpoint.addons._ConsolidatedHFMetadataExporter.save(
    hf_metadata_dir: str,
    tokenizer: object,
    original_model_path: str | None
) -> None

Write model-specific metadata into the shared Hugging Face metadata directory.

nemo_automodel.components.checkpoint.addons._ConsolidatedHFMetadataExporter.validate(
    tokenizer: object,
    original_model_path: str | None
) -> None

Validate exporter inputs before any distributed rank writes files.

nemo_automodel.components.checkpoint.addons._apply_transformers_compat_guards(
    py_path: str
) -> None

Guard imports of transformers symbols removed in newer versions.

For each copied .py that does from <module> import ... <symbol> ... where <symbol> was removed upstream, insert a preamble defining the symbol on <module> if absent, so the subsequent import resolves. Files that don’t reference such symbols are left byte-for-byte unchanged.

nemo_automodel.components.checkpoint.addons._config_exists(
    original_model_path: str | None,
    config_name: str
) -> bool

nemo_automodel.components.checkpoint.addons._extract_target_modules(
    model: torch.nn.Module | list[torch.nn.Module],
    v4_compatible: bool = False,
    pp_group: torch.distributed.ProcessGroup | None = None
) -> list[str]

Extract the target modules from the model used by LoRA/PEFT layers.

Combined-projection module names (e.g. qkv_proj, gate_up_proj) are expanded to the individual HF projection names for adapter_config.json compatibility with vLLM, TensorRT-LLM, and HF PEFT.

For MoE expert LoRA, grouped 3-D adapter parameters are expanded to per-expert HF projection names when in v4-compatible mode (where per-expert target_modules are used). In v5 mode (v4_compatible=False) the expansion is skipped because target_parameters provides the fused ParamWrapper paths instead.

Strips _orig_mod. (torch.compile) and _checkpoint_wrapped_module. (activation checkpointing) prefixes from module names.

model may be a single module or a list of modules. Under pipeline parallelism the caller passes the full list of this rank’s virtual stages, so names are unioned across all local parts (using only the first part would miss the other stages’ layers). When pp_group is also provided, the discovered names are additionally unioned across the PP group so the result covers every rank’s layers too.

nemo_automodel.components.checkpoint.addons._extract_target_parameters(
    model: torch.nn.Module | list[torch.nn.Module],
    v4_compatible: bool = False
) -> list[str]

Extract target_parameters for PEFT v0.18+ ParamWrapper format.

Returns fused expert parameter paths for adapters that explicitly opt in to PEFT v5 ParamWrapper export, or an empty list otherwise.

model may be a single module or a list of PP parts; the check is a per-model-class property, so the first part is representative.

nemo_automodel.components.checkpoint.addons._get_automodel_peft_metadata(
    peft_config: peft.PeftConfig
) -> dict

Get the PEFT metadata in the format expected by Automodel.

Parameters:

peft_config

PeftConfig

Source PEFT configuration.

Returns: dict

A dict containing Automodel-specific PEFT metadata fields filtered from

nemo_automodel.components.checkpoint.addons._get_hf_peft_config(
    peft_config: peft.PeftConfig,
    model_state: nemo_automodel.components.checkpoint.stateful_wrappers.ModelState,
    v4_compatible: bool = False
) -> dict

Get the minimal PEFT config in the format expected by Hugging Face.

Parameters:

peft_config

PeftConfig

Source PEFT configuration.

model_state

ModelState

Model wrapper used to infer target modules and model task.

v4_compatible

boolDefaults to False

When True, use legacy per-expert expansion format.

Returns: dict

A dictionary containing the minimal HF-compatible PEFT configuration

nemo_automodel.components.checkpoint.addons._group_barrier(
    process_group: torch.distributed.ProcessGroup | None
) -> None

nemo_automodel.components.checkpoint.addons._is_group_rank_0(
    process_group: torch.distributed.ProcessGroup | None
) -> bool

nemo_automodel.components.checkpoint.addons._iter_custom_code_classes(
    model_part: torch.nn.Module
)

Yield classes referenced by config.auto_map (and the model’s own class).

Walks the full MRO so wrappers like FSDP2 (which add mixins / rename the top-level class) don’t hide the original transformers_modules.* class.

nemo_automodel.components.checkpoint.addons._maybe_save_custom_model_code(
    original_model_path: str | None,
    hf_metadata_dir: str,
    model_part: torch.nn.Module | None = None
) -> None

Save the custom model code if it exists. This function preserves the original directory structure.

When original_model_path is a local dir, copy its .py files. When it is an HF hub id (e.g. nvidia/Nemotron-Flash-1B) and the loaded model has auto_map custom code, copy the .py files from the cached transformers_modules directory so the consolidated checkpoint carries modeling_*.py locally and reloads without needing trust_remote_code=True.

nemo_automodel.components.checkpoint.addons._maybe_strip_quantization_config(
    model_part: torch.nn.Module,
    config = None
) -> None

Remove quantization_config from the HF config when no parameters are quantized.

Models loaded from quantized checkpoints (e.g. mxfp4 GPT-OSS) carry a quantization_config on their config object. After dequantization all parameters are standard floating-point, but the stale config entry would still be written to the saved config.json. This strips it so the output checkpoint is a clean bf16 checkpoint, consistent with e.g. unsloth/gpt-oss-20b-BF16.

nemo_automodel.components.checkpoint.addons._save_generated_hf_assets(
    model_part: torch.nn.Module,
    metadata_reference_path: str | None,
    hf_metadata_dir: str,
    tokenizer,
    v4_compatible: bool,
    model_config = None,
    save_custom_model_code: bool = True
) -> None

Run the existing generated Hugging Face metadata export path.

nemo_automodel.components.checkpoint.addons._save_original_config_json(
    original_model_path: str,
    hf_metadata_dir: str,
    config_name: str
) -> None

Copy the original pretrained config.json with quantization_config stripped.

This is used in v4-compatible mode so that downstream consumers (e.g. vLLM) that expect a transformers-v4-style config receive the file verbatim from the original checkpoint, minus any quantization metadata (since saved weights are always bf16).

nemo_automodel.components.checkpoint.addons._unwrap_ddp_model(
    model: torch.nn.Module
) -> torch.nn.Module

Return the module that owns export metadata hidden by DDP.

nemo_automodel.components.checkpoint.addons._REMOVED_TRANSFORMERS_SYMBOLS = {'NEED_SETUP_CACHE_CLASSES_MAPPING': ('transformers.generation.utils', '{}')}