nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter

Module Contents

Classes

GlmMoeDsaStateDictAdapter

Converts between HF GLM-MoE-DSA checkpoints and native format.

API

class nemo_automodel.components.models.glm_moe_dsa.state_dict_adapter.GlmMoeDsaStateDictAdapter(

config: Any,

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

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

dtype: torch.dtype = torch.float32,

)

Bases: nemo_automodel.components.models.glm4_moe.state_dict_adapter.Glm4MoeStateDictAdapter

Converts between HF GLM-MoE-DSA checkpoints and native format.

Extends Glm4MoeStateDictAdapter with handling for the DSA indexer weights that should not be quantized (k_norm, weights_proj).

Initialization

_indexer_non_quantized_keys

[‘indexer.k_norm.weight’, ‘indexer.k_norm.bias’, ‘indexer.weights_proj.weight’]

convert_single_tensor_to_hf(

fqn: str,

tensor: Any,

**kwargs,

) → list[tuple[str, Any]]