nemo_automodel.components.models.common.bidirectional

Module Contents

Classes

Name	Description
`EncoderStateDictAdapter`	Adapter for encoder model state dicts.

API

class nemo_automodel.components.models.common.bidirectional.EncoderStateDictAdapter(
    backbone_adapter: nemo_automodel.components.checkpoint.state_dict_adapter.StateDictAdapter | None = None
)

Bases: StateDictAdapter

Adapter for encoder model state dicts.

Internal format uses a model. prefix on all keys. HF format does not. This adapter strips or adds the model. prefix as needed, including for PEFT-wrapped keys (base_model.model.model.X <-> base_model.model.X).

_MODEL_PREFIX

= 'model.'

_PEFT_MODEL_PREFIX

= _PEFT_PREFIX + _MODEL_PREFIX

_PEFT_PREFIX

= 'base_model.model.'

checkpoint_key_aliases

dict[str, str]

Forward backbone aliases, which are already expressed in exported HF names.

nemo_automodel.components.models.common.bidirectional.EncoderStateDictAdapter._add_model_prefix(
    key
)

nemo_automodel.components.models.common.bidirectional.EncoderStateDictAdapter._strip_model_prefix(
    key
)

nemo_automodel.components.models.common.bidirectional.EncoderStateDictAdapter.convert_single_tensor_to_hf(
    fqn,
    tensor,
    kwargs = {}
)

Convert a wrapped parameter using the backbone’s per-tensor contract.

Parameters:

fqn

Wrapped native parameter name.

tensor

Tensor in its native model-specific shape, dtype, device and distributed placement.

**kwargs

Defaults to {}

Options passed to the backbone adapter.

Returns:

HF name/tensor pairs following the backbone adapter’s layouts; without an adapter the tensor aliases input.

nemo_automodel.components.models.common.bidirectional.EncoderStateDictAdapter.forced_hf_dtype_mapping(
    state_dict: dict[str, torch.Tensor]
) -> dict[str, str]

Preserve backbone export dtype requirements through the encoder wrapper.

Parameters:

state_dict

dict[str, torch.Tensor]

HF parameter names and tensors in their unchanged model-specific layouts.

Returns: dict[str, str]

HF parameter names mapped to safetensors dtype names required by the backbone adapter.

nemo_automodel.components.models.common.bidirectional.EncoderStateDictAdapter.from_hf(
    hf_state_dict,
    device_mesh = None,
    kwargs = {}
)

Apply the backbone conversion before restoring the wrapper prefix.

Parameters:

hf_state_dict

HF tensors in the backbone adapter’s documented layouts.

device_mesh

Defaults to None

Mesh forwarded to the backbone adapter without modification.

**kwargs

Defaults to {}

Options passed to the backbone adapter.

Returns:

Native tensors in the wrapped model’s layouts and distributed placements.

nemo_automodel.components.models.common.bidirectional.EncoderStateDictAdapter.to_hf(
    state_dict,
    kwargs = {}
)

Remove the wrapper prefix, then apply the backbone conversion.

Parameters:

state_dict

Native tensors in the wrapped model’s documented layouts and distributed placements.

**kwargs

Defaults to {}

Options passed to the backbone adapter.

Returns:

HF tensors in the backbone adapter’s documented layouts, or unchanged tensors without an adapter.