nemo_automodel.components.models.ernie4_5.state_dict_adapter

Module Contents

Classes

Name	Description
`Ernie4_5StateDictAdapter`	Passthrough adapter for dense ERNIE 4.5 checkpoints.
`Ernie4_5_MoeStateDictAdapter`	Convert ERNIE 4.5 MoE HF checkpoints to AutoModel grouped-expert format.

Data

logger

API

class nemo_automodel.components.models.ernie4_5.state_dict_adapter.Ernie4_5StateDictAdapter(
    config: typing.Any
)

Bases: StateDictAdapter

Passthrough adapter for dense ERNIE 4.5 checkpoints.

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

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

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

class nemo_automodel.components.models.ernie4_5.state_dict_adapter.Ernie4_5_MoeStateDictAdapter(
    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 ERNIE 4.5 MoE HF checkpoints to AutoModel grouped-expert format.

nemo_automodel.components.models.ernie4_5.state_dict_adapter.Ernie4_5_MoeStateDictAdapter._hf_key_to_native(
    key: str
) -> str

nemo_automodel.components.models.ernie4_5.state_dict_adapter.Ernie4_5_MoeStateDictAdapter._native_key_to_hf(
    key: str
) -> str

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

nemo_automodel.components.models.ernie4_5.state_dict_adapter.Ernie4_5_MoeStateDictAdapter.from_hf(
    hf_state_dict: dict[str, typing.Any],
    device_mesh: 'DeviceMesh' | None = None,
    kwargs = {}
) -> dict[str, typing.Any]

nemo_automodel.components.models.ernie4_5.state_dict_adapter.Ernie4_5_MoeStateDictAdapter.to_hf(
    state_dict: dict[str, typing.Any],
    exclude_key_regex: str | None = None,
    quantization: bool = False,
    kwargs = {}
) -> dict[str, typing.Any]

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