> This page is for version 0.4.0 · 26.04.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt
> 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.
# nemo_automodel.components.models.praxis.lora_export
Fold the adapters of a LoRA run into the weights and write a Praxis checkpoint.
The checkpointer saves a LoRA run as adapters (`adapter_model.safetensors` next to the PEFT configs);
the reference runtime loads full checkpoints. This module merges the adapters into the linears they adapt
and writes the result in the checkpoint layout (`config.json`, `model-XXXXX-of-NNNNN.safetensors`
shards, `model.safetensors.index.json`, `SHA256SUMS`)::
python -m nemo_automodel.components.models.praxis.lora_export --checkpoint runwayml/praxis-1-droid \
\--adapters runs/libero_spatial_lora/epoch_0_step_599/model --out runs/libero_spatial_lora/merged_600
## Module Contents
### Functions
| Name | Description |
| ----------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------- |
| [`_natural_key`](#nemo_automodel-components-models-praxis-lora_export-_natural_key) | - |
| [`adapter_scales`](#nemo_automodel-components-models-praxis-lora_export-adapter_scales) | The scale `alpha / rank` of every adapter module of `model`, by module name. |
| [`export_merged`](#nemo_automodel-components-models-praxis-lora_export-export_merged) | Merge the adapters of a LoRA run into the base checkpoint and write the result to `out`. |
| [`load_adapters`](#nemo_automodel-components-models-praxis-lora_export-load_adapters) | The adapter tensors of a checkpointer's PEFT save, under the parameter names of `policy`. |
| [`main`](#nemo_automodel-components-models-praxis-lora_export-main) | Command line: merge the adapters of a LoRA run into the base checkpoint and write the result. |
| [`merge_lora`](#nemo_automodel-components-models-praxis-lora_export-merge_lora) | Fold the adapters into the weights they adapt and drop the adapter tensors. |
| [`save_checkpoint`](#nemo_automodel-components-models-praxis-lora_export-save_checkpoint) | Write `native_state` as a Praxis checkpoint directory. |
### Data
[`ADAPTER_FILE`](#nemo_automodel-components-models-praxis-lora_export-ADAPTER_FILE)
[`MAX_SHARD_BYTES`](#nemo_automodel-components-models-praxis-lora_export-MAX_SHARD_BYTES)
[`_EXPERT_LORA`](#nemo_automodel-components-models-praxis-lora_export-_EXPERT_LORA)
[`_EXPERT_LORA_A`](#nemo_automodel-components-models-praxis-lora_export-_EXPERT_LORA_A)
[`_LORA_A`](#nemo_automodel-components-models-praxis-lora_export-_LORA_A)
### API
```python
nemo_automodel.components.models.praxis.lora_export._natural_key(
name: str
) -> list[int | str]
```
```python
nemo_automodel.components.models.praxis.lora_export.adapter_scales(
model: torch.nn.Module
) -> dict[str, float]
```
The scale `alpha / rank` of every adapter module of `model`, by module name.
Expert adapters may run at a smaller rank than the linear adapters (`moe_rank_scaling`), so the
scale is read from each module rather than derived from the PEFT config.
```python
nemo_automodel.components.models.praxis.lora_export.export_merged(
checkpoint: str | pathlib.Path,
adapters: str | pathlib.Path,
out: str | pathlib.Path
) -> None
```
Merge the adapters of a LoRA run into the base checkpoint and write the result to `out`.
**Parameters:**
**`checkpoint`** `str | Path`
The base checkpoint directory or its Hugging Face repo id.
---
**`adapters`** `str | Path`
The `model` directory of the LoRA run's checkpoint step.
---
**`out`** `str | Path`
Output directory of the merged checkpoint.
---
```python
nemo_automodel.components.models.praxis.lora_export.load_adapters(
policy: nemo_automodel.components.models.praxis.model.PraxisPolicy,
model_dir: str | pathlib.Path
) -> dict[str, torch.Tensor]
```
The adapter tensors of a checkpointer's PEFT save, under the parameter names of `policy`.
**Parameters:**
**`policy`** `PraxisPolicy`
The model with its `LinearLoRA` layers applied.
---
**`model_dir`** `str | Path`
The `model` directory of the checkpoint step.
---
**Returns:** `dict[str, Tensor]`
Every `lora_A` / `lora_B` parameter of `policy` with its saved value.
```python
nemo_automodel.components.models.praxis.lora_export.main() -> int
```
Command line: merge the adapters of a LoRA run into the base checkpoint and write the result.
```python
nemo_automodel.components.models.praxis.lora_export.merge_lora(
state: dict[str, torch.Tensor],
scales: collections.abc.Mapping[str, float]
) -> dict[str, torch.Tensor]
```
Fold the adapters into the weights they adapt and drop the adapter tensors.
Linear adapters are `<m>.lora_A.weight` of shape \[rank, in] and `<m>.lora_B.weight` of shape
\[out, rank] next to `<m>.weight`; grouped-expert adapters are `<m>.lora_gate_and_up_{A,B}` and
`<m>.lora_down_{A,B}` of shapes \[experts, in, rank] and \[experts, rank, out] next to
`<m>.gate_and_up_projs` and `<m>.down_projs`.
**Parameters:**
**`state`** `dict[str, Tensor]`
Module state dict with the adapters.
---
**`scales`** `Mapping[str, float]`
The scale `alpha / rank` of each adapted module, by module name (`adapter_scales`).
---
**Returns:** `dict[str, Tensor]`
A new state dict with `weight + scale * B @ A` (`scale * A @ B` per expert) in the weight's dtype.
```python
nemo_automodel.components.models.praxis.lora_export.save_checkpoint(
path: str | pathlib.Path,
native_state: dict[str, torch.Tensor],
adapter: nemo_automodel.components.models.praxis.state_dict_adapter.PraxisStateDictAdapter,
source_config: str | pathlib.Path,
max_shard_bytes: int = MAX_SHARD_BYTES
) -> None
```
Write `native_state` as a Praxis checkpoint directory.
Tensors are stored in the checkpoint dtype (`torch_dtype` of the config), except the
names matching `fp32_tensors` which stay float32.
**Parameters:**
**`path`** `str | Path`
Output directory (created; existing shards are replaced).
---
**`native_state`** `dict[str, Tensor]`
Full module state dict (CPU tensors) in module key names.
---
**`adapter`** `PraxisStateDictAdapter`
Converts module keys to checkpoint keys.
---
**`source_config`** `str | Path`
`config.json` of the checkpoint the weights derive from; copied verbatim.
---
**`max_shard_bytes`** `int` — default: MAX_SHARD_BYTES
Shard size limit.
---
```python
nemo_automodel.components.models.praxis.lora_export.ADAPTER_FILE = 'adapter_model.safetensors'
```
```python
nemo_automodel.components.models.praxis.lora_export.MAX_SHARD_BYTES = 5 * 1024 ** 3
```
```python
nemo_automodel.components.models.praxis.lora_export._EXPERT_LORA = {'lora_gate_and_up': 'gate_and_up_projs', 'lora_down': 'down_projs'}
```
```python
nemo_automodel.components.models.praxis.lora_export._EXPERT_LORA_A = re.compile('^(?P.+)\\.(?Plora_gate_and_up|lora_down)_A$')
```
```python
nemo_automodel.components.models.praxis.lora_export._LORA_A = re.compile('^(?P.+)\\.lora_A\\.weight$')
```