> This page is for version Nightly (default).
> 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.state_dict_adapter

Weight-name conversion between the Praxis checkpoint and the Automodel module tree.

The checkpoint (`model-*.safetensors` with `model.safetensors.index.json`) stores the
trunk under `backbone.{encoder,decoder}.` and the expert under `action_expert.`. The
action expert names are identical on both sides. The trunk differs in these places:

* projections: `t_proj.VideoTensor` -> `time_proj`; `input_projections.TextEncodingTensor.proj.{0,2,3}`
  -> `text_in.{fc1,fc2,norm}`; `input_projections.VideoTensor.proj` -> `video_in`;
  `output_projections.VideoTensor.proj` -> `video_out` (encoder) or
  `output_projections.VideoTensor.proj.{0,1}` -> `video_out.{norm,linear}` (decoder);
* attention: `mha.in_proj.{q,k,v}_weight` -> `attn.{q,k,v}_proj.weight` and the other
  `mha.*` names -> `attn.*`;
* dense feed-forward: `ffn.0.linear.{a,b}_weight` -> `ffn.{gate,up}_proj.weight` and
  `ffn.1.weight` -> `ffn.down_proj.weight`;
* mixture of experts: `ffn.router.weight` -> `ffn.gate.weight`, `ffn.expert_bias` ->
  `ffn.gate.e_score_correction_bias`, `ffn.routed.{gate_up,down}_weight` ->
  `ffn.experts.{gate_and_up_projs,down_projs}` (same `[experts, in, out]` layout), and the
  packed shared expert `ffn.shared.gate_up` of shape \[hidden, 2 \* inter] is split and
  transposed into `ffn.shared_experts.{gate,up}_proj.weight` of shape \[inter, hidden];
  `ffn.shared.down` of shape \[inter, hidden] is transposed into `down_proj.weight`.

## Module Contents

### Classes

| Name                                                                                                           | Description                                                            |
| -------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------- |
| [`PraxisStateDictAdapter`](#nemo_automodel-components-models-praxis-state_dict_adapter-PraxisStateDictAdapter) | Convert between checkpoint keys (`from_hf`) and module keys (`to_hf`). |

### Functions

| Name                                                                                   | Description                                                                               |
| -------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- |
| [`_split_key`](#nemo_automodel-components-models-praxis-state_dict_adapter-_split_key) | Split a trunk key into `(tower prefix, block prefix or '', leaf)`; `None` for other keys. |

### Data

[`_BLOCK`](#nemo_automodel-components-models-praxis-state_dict_adapter-_BLOCK)

[`_BLOCK_LEAVES`](#nemo_automodel-components-models-praxis-state_dict_adapter-_BLOCK_LEAVES)

[`_DECODER_LEAVES`](#nemo_automodel-components-models-praxis-state_dict_adapter-_DECODER_LEAVES)

[`_DECODER_TO_HF`](#nemo_automodel-components-models-praxis-state_dict_adapter-_DECODER_TO_HF)

[`_ENCODER_LEAVES`](#nemo_automodel-components-models-praxis-state_dict_adapter-_ENCODER_LEAVES)

[`_ENCODER_TO_HF`](#nemo_automodel-components-models-praxis-state_dict_adapter-_ENCODER_TO_HF)

[`_LEAVES_TO_HF`](#nemo_automodel-components-models-praxis-state_dict_adapter-_LEAVES_TO_HF)

[`_MODULATOR_HF`](#nemo_automodel-components-models-praxis-state_dict_adapter-_MODULATOR_HF)

[`_MODULATOR_NATIVE`](#nemo_automodel-components-models-praxis-state_dict_adapter-_MODULATOR_NATIVE)

[`_PREFIX`](#nemo_automodel-components-models-praxis-state_dict_adapter-_PREFIX)

[`_SHARED_DOWN`](#nemo_automodel-components-models-praxis-state_dict_adapter-_SHARED_DOWN)

[`_SHARED_DOWN_NATIVE`](#nemo_automodel-components-models-praxis-state_dict_adapter-_SHARED_DOWN_NATIVE)

[`_SHARED_GATE`](#nemo_automodel-components-models-praxis-state_dict_adapter-_SHARED_GATE)

[`_SHARED_GATE_UP`](#nemo_automodel-components-models-praxis-state_dict_adapter-_SHARED_GATE_UP)

[`_SHARED_UP`](#nemo_automodel-components-models-praxis-state_dict_adapter-_SHARED_UP)

### API

```python
class nemo_automodel.components.models.praxis.state_dict_adapter.PraxisStateDictAdapter(
    shared_expert_inter_dim: int
)
```

**Bases:** [StateDictAdapter](/nemo/automodel/nemo-automodel/nemo_automodel/components/checkpoint/state_dict_adapter#nemo_automodel-components-checkpoint-state_dict_adapter-StateDictAdapter)

Convert between checkpoint keys (`from_hf`) and module keys (`to_hf`).

```python
nemo_automodel.components.models.praxis.state_dict_adapter.PraxisStateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: torch.Tensor,
    kwargs = {}
) -> list[tuple[str, torch.Tensor]]
```

Convert one module tensor to its checkpoint tensor(s); see `to_hf`.

```python
nemo_automodel.components.models.praxis.state_dict_adapter.PraxisStateDictAdapter.from_hf(
    hf_state_dict: dict[str, torch.Tensor],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    kwargs = {}
) -> dict[str, torch.Tensor]
```

Map checkpoint tensors to module tensors.

**Parameters:**

**`hf_state_dict`** `dict[str, Tensor]`

Checkpoint tensors keyed by their checkpoint names.

---

**Returns:** `dict[str, Tensor]`

Tensors keyed by module names. Shared-expert weights are new contiguous tensors; every

```python
nemo_automodel.components.models.praxis.state_dict_adapter.PraxisStateDictAdapter.to_hf(
    state_dict: dict[str, torch.Tensor],
    exclude_key_regex: str | None = None,
    kwargs = {}
) -> dict[str, torch.Tensor]
```

Map module tensors back to checkpoint tensors (the inverse of `from_hf`).

**Parameters:**

**`state_dict`** `dict[str, Tensor]`

Module tensors keyed by module names.

---

**Returns:** `dict[str, Tensor]`

Tensors keyed by checkpoint names; the two shared-expert input projections of a block

```python
nemo_automodel.components.models.praxis.state_dict_adapter._split_key(
    key: str
) -> tuple[str, str, str] | None
```

Split a trunk key into `(tower prefix, block prefix or '', leaf)`; `None` for other keys.

```python
nemo_automodel.components.models.praxis.state_dict_adapter._BLOCK = re.compile('^(?P<block>core\\.blocks\\.\\d+\\.)(?P<leaf>.+)$')
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._BLOCK_LEAVES = {'mha.in_proj.q_weight': 'attn.q_proj.weight', 'mha.in_proj.k_weight': 'attn.k_p...
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._DECODER_LEAVES = {'t_proj.VideoTensor.weight': 'time_proj.weight', 't_proj.VideoTensor.bias': 'ti...
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._DECODER_TO_HF = {v: k for k, v in (_DECODER_LEAVES.items())}
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._ENCODER_LEAVES = {'t_proj.VideoTensor.weight': 'time_proj.weight', 't_proj.VideoTensor.bias': 'ti...
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._ENCODER_TO_HF = {v: k for k, v in (_ENCODER_LEAVES.items())}
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._LEAVES_TO_HF = {v: k for k, v in (_BLOCK_LEAVES.items())}
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._MODULATOR_HF = re.compile('^core\\.modulators\\.(?P<i>\\d+)\\.(?P<p>weight|bias)$')
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._MODULATOR_NATIVE = re.compile('^core\\.blocks\\.(?P<i>\\d+)\\.modulator\\.(?P<p>weight|bias)$')
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._PREFIX = re.compile('^(?P<tower>backbone\\.(?:encoder|decoder)\\.)(?P<rest>.+)$')
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._SHARED_DOWN = 'ffn.shared.down'
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._SHARED_DOWN_NATIVE = 'ffn.shared_experts.down_proj.weight'
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._SHARED_GATE = 'ffn.shared_experts.gate_proj.weight'
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._SHARED_GATE_UP = 'ffn.shared.gate_up'
```

```python
nemo_automodel.components.models.praxis.state_dict_adapter._SHARED_UP = 'ffn.shared_experts.up_proj.weight'
```