nemo_automodel.components.models.llama.state_dict_adapter

State dict adapter for Llama model.

The model uses separate q_proj / k_proj / v_proj / gate_proj / up_proj that match HuggingFace key names exactly, so the adapter is a passthrough (only tied-weight handling is applied in from_hf).

Module Contents

Classes

Name	Description
`LlamaStateDictAdapter`	State dict adapter for Llama models.

Data

logger

API

class nemo_automodel.components.models.llama.state_dict_adapter.LlamaStateDictAdapter(
    config: transformers.LlamaConfig
)

Bases: StateDictAdapter

State dict adapter for Llama models.

Subclassing the common interface makes the adapter checkpoint-compatible; load-path policy remains entirely inside the checkpoint package.

Uses separate projections that match HuggingFace key names exactly, so from_hf / to_hf are simple passthroughs (only tied-weight handling in from_hf).

nemo_automodel.components.models.llama.state_dict_adapter.LlamaStateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: typing.Any,
    kwargs = {}
) -> list[tuple[str, typing.Any]]

Return one Llama tensor under its unchanged HF state-dict key.

Parameters:

fqn

str

Fully-qualified HF state-dict key. The key determines the tensor’s rank and axis order.

tensor

Any

Tensor/value to export. This passthrough does not copy or transform it.

**kwargs

Defaults to {}

Optional controls, including exclude_key_regex to skip matching keys.

Returns: list[tuple[str, Any]]

A single (fqn, tensor) tuple, or an empty list when filtered. The returned tensor aliases the input.

nemo_automodel.components.models.llama.state_dict_adapter.LlamaStateDictAdapter.from_hf(
    hf_state_dict: dict[str, typing.Any],
    kwargs = {}
) -> dict[str, typing.Any]

nemo_automodel.components.models.llama.state_dict_adapter.LlamaStateDictAdapter.to_hf(
    state_dict: dict[str, typing.Any],
    exclude_key_regex: str | None = None,
    kwargs = {}
) -> dict[str, typing.Any]

nemo_automodel.components.models.llama.state_dict_adapter.logger = logging.getLogger(__name__)