nemo_automodel.components.models.laguna.state_dict_adapter

Module Contents

Classes

Name	Description
`LagunaStateDictAdapter`	Convert Laguna HF checkpoints to Automodel’s grouped-MoE layout.

Functions

Name	Description
`_apply_renames`	-

Data

_HF_TO_NATIVE_RENAMES

_NATIVE_TO_HF_RENAMES

API

class nemo_automodel.components.models.laguna.state_dict_adapter.LagunaStateDictAdapter(
    config: typing.Any,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.bfloat16
)

Bases: MoESplitExpertsStateDictMixin, StateDictAdapter

Convert Laguna HF checkpoints to Automodel’s grouped-MoE layout.

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

nemo_automodel.components.models.laguna.state_dict_adapter.LagunaStateDictAdapter.from_hf(
    hf_state_dict: dict[str, typing.Any],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    kwargs = {}
) -> dict[str, typing.Any]

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

nemo_automodel.components.models.laguna.state_dict_adapter._apply_renames(
    key: str,
    renames: tuple[tuple[re.Pattern[str], str], ...]
) -> str

nemo_automodel.components.models.laguna.state_dict_adapter._HF_TO_NATIVE_RENAMES: tuple[tuple[Pattern[str], str], ...] = ((re.compile('\\.mlp\\.shared_expert\\.'), '.mlp.shared_experts.'), (re.compile(...

nemo_automodel.components.models.laguna.state_dict_adapter._NATIVE_TO_HF_RENAMES: tuple[tuple[Pattern[str], str], ...] = ((re.compile('\\.mlp\\.shared_experts\\.'), '.mlp.shared_expert.'), (re.compile(...