nemo_automodel.components.checkpoint.addons

Module Contents

Classes

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

_get_hf_peft_config	Get the minimal PEFT config in the format expected by Hugging Face.
_get_automodel_peft_metadata	Get the PEFT metadata in the format expected by Automodel.
_is_qwen3_moe	Check whether model uses the Qwen3 MoE state-dict adapter.
_extract_target_parameters	Extract target_parameters for PEFT v0.18+ ParamWrapper format.
_extract_target_modules	Extract the target modules from the model used by LoRA/PEFT layers.
_maybe_strip_quantization_config	Remove quantization_config from the HF config when no parameters are quantized.
_config_exists
_save_original_config_json	Copy the original pretrained config.json with quantization_config stripped.
_maybe_save_custom_model_code	Save the custom model code if it exists. This function preserves the original directory structure.

API

class nemo_automodel.components.checkpoint.addons.CheckpointAddon

Bases: typing.Protocol

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

pre_save(**kwargs) → None

post_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.

pre_save(**kwargs) → None

Pre-save hook to emit consolidated HF artifacts.

Expected kwargs: model_state (ModelState): Wrapper holding the model parts. hf_metadata_dir (str): Target directory for HF metadata artifacts. tokenizer (PreTrainedTokenizerBase | None): Optional tokenizer to save.

post_save(**kwargs) → None

Move the saved HF metadata to the consolidated directory.

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

Expected kwargs: consolidated_path (str): Target directory for consolidated artifacts. hf_metadata_dir (str): Target directory for HF metadata 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.

pre_save(**kwargs) → None

Pre-save hook to emit PEFT artifacts.

Expected kwargs: model_path (str): Directory in which to save PEFT files. tokenizer (PreTrainedTokenizerBase | None): Optional tokenizer to save. model_state (ModelState): Wrapper holding the model parts. peft_config (PeftConfig): PEFT configuration for serialization.

post_save(**kwargs) → None

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 – Source PEFT configuration.

• model_state – Model wrapper used to infer target modules and model task.

• v4_compatible – When True, use legacy per-expert expansion format.

Returns:

A dictionary containing the minimal HF-compatible PEFT configuration (e.g., task type, LoRA rank/alpha, and discovered target modules).

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 – Source PEFT configuration.

Returns:

A dict containing Automodel-specific PEFT metadata fields filtered from the full 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._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._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._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._config_exists(original_model_path: str, config_name: str) → bool

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._maybe_save_custom_model_code(

original_model_path: str | None,

hf_metadata_dir: str,

) → None

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