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`	-
`adapter_scales`	The scale `alpha / rank` of every adapter module of `model`, by module name.
`export_merged`	Merge the adapters of a LoRA run into the base checkpoint and write the result to `out`.
`load_adapters`	The adapter tensors of a checkpointer’s PEFT save, under the parameter names of `policy`.
`main`	Command line: merge the adapters of a LoRA run into the base checkpoint and write the result.
`merge_lora`	Fold the adapters into the weights they adapt and drop the adapter tensors.
`save_checkpoint`	Write `native_state` as a Praxis checkpoint directory.

Data

API

nemo_automodel.components.models.praxis.lora_export._natural_key(
    name: str
) -> list[int | str]

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.

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.

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.

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.

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.

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

intDefaults to MAX_SHARD_BYTES

Shard size limit.

nemo_automodel.components.models.praxis.lora_export.ADAPTER_FILE = 'adapter_model.safetensors'

nemo_automodel.components.models.praxis.lora_export.MAX_SHARD_BYTES = 5 * 1024 ** 3

nemo_automodel.components.models.praxis.lora_export._EXPERT_LORA = {'lora_gate_and_up': 'gate_and_up_projs', 'lora_down': 'down_projs'}

nemo_automodel.components.models.praxis.lora_export._EXPERT_LORA_A = re.compile('^(?P<base>.+)\\.(?P<kind>lora_gate_and_up|lora_down)_A$')

nemo_automodel.components.models.praxis.lora_export._LORA_A = re.compile('^(?P<base>.+)\\.lora_A\\.weight$')