Skip to content

ernie4_5_moe

ernie4_5_moe

ERNIE 4.5 MoE (Ernie4_5_MoeForCausalLM).

Llama's tree and Llama's block: model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, mlp}, norm} and lm_head, the residual added in the block, attention through the shared eager forward. Blocks from moe_layer_start_index to moe_layer_end_index, every moe_layer_interval-th, have a mixture of experts (with optional shared experts, added before it returns); the others a dense MLP. The experts are moe_intermediate_size wide.

RENAME module-attribute

RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.norm': 'norm', 'gate': 'router'}

ENVOYS module-attribute

ENVOYS = {Ernie4_5_MoeDecoderLayer: Layer, Ernie4_5_MoeAttention: Attention, Ernie4_5_MoeMLP: Mlp, Ernie4_5_MoeSparseMoeBlock: Moe}