nemo_automodel.components.models.ernie4_5

Submodules

Package Contents

Classes

Name	Description
`Ernie4_5ForCausalLM`	Dense ERNIE 4.5 causal language model.
`Ernie4_5_MoeForCausalLM`	ERNIE 4.5 MoE causal language model with AutoModel EP support.

API

class nemo_automodel.components.models.ernie4_5.model.Ernie4_5ForCausalLM(
    config: transformers.models.ernie4_5.configuration_ernie4_5.Ernie4_5Config,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

Bases: HFCheckpointingMixin, Module

Dense ERNIE 4.5 causal language model.

_nemo_tied_weights_keys

= {'lm_head.weight': 'model.embed_tokens.weight'}

_pp_plan

= {'lm_head': (['hidden_states'], ['logits'])}

_tp_plan

= {'lm_head': 'colwise_rep'}

backend

= backend or BackendConfig()

lm_head

model

= Ernie4_5Model(config, self.backend)

state_dict_adapter

= Ernie4_5StateDictAdapter(config)

tie_word_embeddings_support

TieSupport = TieSupport.TIED_ONLY

vocab_size

= config.vocab_size

nemo_automodel.components.models.ernie4_5.model.Ernie4_5ForCausalLM.forward(
    input_ids: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    logits_to_keep: typing.Union[int, torch.Tensor] = 0,
    output_hidden_states: bool | None = None,
    attn_kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast

nemo_automodel.components.models.ernie4_5.model.Ernie4_5ForCausalLM.from_config(
    config: transformers.models.ernie4_5.configuration_ernie4_5.Ernie4_5Config,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

classmethod

nemo_automodel.components.models.ernie4_5.model.Ernie4_5ForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args = (),
    kwargs = {}
)

classmethod

nemo_automodel.components.models.ernie4_5.model.Ernie4_5ForCausalLM.get_input_embeddings()

nemo_automodel.components.models.ernie4_5.model.Ernie4_5ForCausalLM.get_output_embeddings()

nemo_automodel.components.models.ernie4_5.model.Ernie4_5ForCausalLM.set_input_embeddings(
    value
)

nemo_automodel.components.models.ernie4_5.model.Ernie4_5ForCausalLM.set_output_embeddings(
    new_embeddings
)

nemo_automodel.components.models.ernie4_5.model.Ernie4_5ForCausalLM.tie_weights()

class nemo_automodel.components.models.ernie4_5.model.Ernie4_5_MoeForCausalLM(
    config: transformers.models.ernie4_5_moe.configuration_ernie4_5_moe.Ernie4_5_MoeConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin

ERNIE 4.5 MoE causal language model with AutoModel EP support.

_keep_in_fp32_modules

= ['rotary_emb']

_keep_in_fp32_modules_strict

_nemo_tied_weights_keys

= {'lm_head.weight': 'model.embed_tokens.weight'}

_pp_plan

= {'lm_head': (['hidden_states'], ['logits'])}

_tp_plan

= {'lm_head': 'colwise_rep'}

backend

= backend or BackendConfig()

lm_head

model

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.TIED_ONLY

vocab_size

= config.vocab_size

nemo_automodel.components.models.ernie4_5.model.Ernie4_5_MoeForCausalLM.forward(
    input_ids: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    logits_to_keep: typing.Union[int, torch.Tensor] = 0,
    output_hidden_states: bool | None = None,
    attn_kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast

nemo_automodel.components.models.ernie4_5.model.Ernie4_5_MoeForCausalLM.from_config(
    config: transformers.models.ernie4_5_moe.configuration_ernie4_5_moe.Ernie4_5_MoeConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

classmethod

nemo_automodel.components.models.ernie4_5.model.Ernie4_5_MoeForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args = (),
    kwargs = {}
)

classmethod

nemo_automodel.components.models.ernie4_5.model.Ernie4_5_MoeForCausalLM.get_capabilities(
    config
) -> nemo_automodel._transformers.model_capabilities.ModelCapabilities

classmethod

Return parallelism capabilities for a specific ERNIE-4.5 config.

ERNIE-4.5 ships in two flavors that share this class file but exercise different code paths:

baidu/ERNIE-4.5-21B-A3B-PT — MoE variant (this NeMo custom class). moe_num_experts > 0 in the HF config. Demonstrated by examples/llm_finetune/ernie4_5/ernie4_5_21b_a3b_hellaswag.yaml (ep_size=8).
baidu/ERNIE-4.5-0.3B-PT — dense variant. No expert config. Demonstrated by examples/llm_finetune/ernie4_5/ernie4_5_0p3b_hellaswag.yaml (tp/cp/pp/ep all 1).

nemo_automodel.components.models.ernie4_5.model.Ernie4_5_MoeForCausalLM.get_input_embeddings()

nemo_automodel.components.models.ernie4_5.model.Ernie4_5_MoeForCausalLM.get_output_embeddings()

nemo_automodel.components.models.ernie4_5.model.Ernie4_5_MoeForCausalLM.set_input_embeddings(
    value
)

nemo_automodel.components.models.ernie4_5.model.Ernie4_5_MoeForCausalLM.set_output_embeddings(
    new_embeddings
)

nemo_automodel.components.models.ernie4_5.model.Ernie4_5_MoeForCausalLM.tie_weights()