nemo_automodel.components.models.llama.state_dict_adapter

State dict adapter for Llama model with combined projections.

Uses the generic CombinedProjectionStateDictAdapter from common/.

Module Contents

Classes

LlamaStateDictAdapter

State dict adapter for Llama models.

API

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

Bases: nemo_automodel.components.models.common.combined_projection.state_dict_adapter.CombinedProjectionStateDictAdapter

State dict adapter for Llama models.

Inherits from the generic CombinedProjectionStateDictAdapter, providing a clean interface specific to Llama.

.. rubric:: Example

from transformers import LlamaConfig

config = LlamaConfig.from_pretrained(“meta-llama/Llama-3-8B”) adapter = LlamaStateDictAdapter(config)

Convert HF checkpoint to custom format

custom_state_dict = adapter.from_hf(hf_state_dict)

Convert custom checkpoint back to HF format

hf_state_dict = adapter.to_hf(custom_state_dict)

Initialization

Initialize adapter with Llama config.