> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.checkpoint.addons

## Module Contents

### Classes

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

### Functions

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

### Data

[`_MOE_LORA_PARAMETER_SUFFIXES`](#nemo_automodel-components-checkpoint-addons-_MOE_LORA_PARAMETER_SUFFIXES)

[`_REMOVED_TRANSFORMERS_SYMBOLS`](#nemo_automodel-components-checkpoint-addons-_REMOVED_TRANSFORMERS_SYMBOLS)

### API

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

Protocol

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

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

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

```python
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.

```python
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.

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

Pre-save hook to emit consolidated HF artifacts.

```python
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.

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

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

Pre-save hook to emit PEFT artifacts.

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

Protocol

Model-provided writer for consolidated Hugging Face metadata.

```python
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.

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

Validate exporter inputs before any distributed rank writes files.

```python
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.

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

```python
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.

```python
nemo_automodel.components.checkpoint.addons._extract_target_parameters(
    model: torch.nn.Module | list[torch.nn.Module],
    v4_compatible: bool = False,
    pp_group: torch.distributed.ProcessGroup | None = None
) -> 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 and have trainable grouped-expert LoRA
parameters, or an empty list otherwise.

`model` may be a single module or a list of local PP parts. When a PP
group is provided, expert-LoRA presence is unioned across its ranks so all
stages emit identical metadata.

```python
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

```python
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`** `bool` — default: False

When True, use legacy per-expert expansion format.

---

**Returns:** `dict`

A dictionary containing the minimal HF-compatible PEFT configuration

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

```python
nemo_automodel.components.checkpoint.addons._has_trainable_moe_lora_parameters(
    model: torch.nn.Module | list[torch.nn.Module],
    pp_group: torch.distributed.ProcessGroup | None = None
) -> bool
```

Return whether any pipeline stage owns trainable grouped-expert LoRA parameters.

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

```python
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.

```python
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`.

```python
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`.

```python
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.

```python
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).

```python
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.

```python
nemo_automodel.components.checkpoint.addons._MOE_LORA_PARAMETER_SUFFIXES = ('lora_gate_and_up_A', 'lora_gate_and_up_B', 'lora_down_A', 'lora_down_B')
```

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