nemo_automodel.components.models.ernie4_5.state_dict_adapter

Module Contents

Classes

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

nemo_automodel.components.models.ernie4_5.state_dict_adapter.logger

‘getLogger(…)’

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

Bases: nemo_automodel.components.checkpoint.state_dict_adapter.StateDictAdapter

Passthrough adapter for dense ERNIE 4.5 checkpoints.

Initialization

from_hf(

hf_state_dict: dict[str, Any],

**kwargs,

) → dict[str, Any]

to_hf(

state_dict: dict[str, Any],

exclude_key_regex: Optional[str] = None,

**kwargs,

) → dict[str, Any]

convert_single_tensor_to_hf(

fqn: str,

tensor: Any,

**kwargs,

) → list[tuple[str, Any]]

class nemo_automodel.components.models.ernie4_5.state_dict_adapter.Ernie4_5_MoeStateDictAdapter(

config: Any,

moe_config: nemo_automodel.components.moe.config.MoEConfig,

backend: nemo_automodel.components.models.common.BackendConfig,

dtype: torch.dtype = torch.bfloat16,

)

Bases: nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin, nemo_automodel.components.checkpoint.state_dict_adapter.StateDictAdapter

Convert ERNIE 4.5 MoE HF checkpoints to AutoModel grouped-expert format.

Initialization

_hf_key_to_native(key: str) → str

_native_key_to_hf(key: str) → str

from_hf(

hf_state_dict: dict[str, Any],

device_mesh: Optional[torch.distributed.device_mesh.DeviceMesh] = None,

**kwargs,

) → dict[str, Any]

to_hf(

state_dict: dict[str, Any],

exclude_key_regex: Optional[str] = None,

quantization: bool = False,

**kwargs,

) → dict[str, Any]

convert_single_tensor_to_hf(

fqn: str,

tensor: Any,

**kwargs,

) → list[tuple[str, Any]]