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

### 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.                                |
| [`_is_qwen3_moe`](#nemo_automodel-components-checkpoint-addons-_is_qwen3_moe)                                         | Check whether *model* uses the Qwen3 MoE state-dict adapter.                                       |
| [`_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_original_config_json`](#nemo_automodel-components-checkpoint-addons-_save_original_config_json)               | Copy the original pretrained `config.json` with `quantization_config` stripped.                    |

### Data

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

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

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

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

```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:**

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:**

Source PEFT configuration.

Model wrapper used to infer target modules and model task.

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

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

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