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.

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.
`_is_qwen3_moe`	Check whether model uses the Qwen3 MoE state-dict adapter.
`_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_original_config_json`	Copy the original pretrained `config.json` with `quantization_config` stripped.

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.

On rank 0, this saves config.json, generation_config.json, and tokenizer artifacts into the provided consolidated directory, 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.

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,
    config_name: str
) -> bool

nemo_automodel.components.checkpoint.addons._extract_target_modules(
    model: torch.nn.Module,
    v4_compatible: bool = False
) -> 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 unless the model is Qwen3 MoE in non-legacy mode (where target_parameters is used instead).

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

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

Extract target_parameters for PEFT v0.18+ ParamWrapper format.

Returns fused expert parameter paths for Qwen3 MoE when not in legacy mode, or an empty list otherwise.

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._is_qwen3_moe(
    model: torch.nn.Module
) -> bool

Check whether model uses the Qwen3 MoE state-dict adapter.

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
) -> 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_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._REMOVED_TRANSFORMERS_SYMBOLS = {'NEED_SETUP_CACHE_CLASSES_MAPPING': ('transformers.generation.utils', '{}')}