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.
_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,

) → 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.

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._extract_target_modules(model: torch.nn.Module) → 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 Hugging Face projection names so that the saved adapter_config.json is compatible with vLLM, TensorRT-LLM and the Hugging Face PEFT library.

For MoE expert LoRA (GroupedExpertsLoRA / GroupedExpertsDeepEPLoRA), the grouped 3-D adapter parameters are expanded to per-expert HF projection names (e.g. model.layers.0.mlp.experts.0.gate_proj).

.. note::

When torch.compile is used, module names get prefixed with _orig_mod.. This function strips those prefixes to get the original module names.

Parameters:

model – The model whose named modules are scanned.

Returns:

A sorted list of unique module name prefixes that contain LoRA layers.

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.