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

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