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`	Convert between checkpoint keys (`from_hf`) and module keys (`to_hf`).

Functions

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

Data

_SHARED_DOWN_NATIVE

_SHARED_GATE

_SHARED_GATE_UP

_SHARED_UP

API

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

Bases: StateDictAdapter

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

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.

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

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

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.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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