深入解析 Ernie 4.5 MoE：百度混合专家大模型在 Transformers 中的加载、推理与源码原理

Ernie 4.5 MoE 是百度 Ernie 4.5 模型家族中的纯文本混合专家（Mixture-of-Experts, MoE）基础模型，于 2025-07-21 由社区贡献合并进本仓库（transformers）。本文以其官方文档 docs/source/en/model_doc/ernie4_5_moe.md 为核心骨架，讲解如何使用 AutoModelForCausalLM 加载模型完成文本生成、如何利用 Tensor Parallelism 在多 GPU 上分布式推理、以及如何通过 Bitsandbytes 4-bit 量化降低显存占用；同时结合本仓库的实现源码剖析其「Llama 核心 + Mixtral 风格 MoE + 共享专家」的架构细节与每个关键配置项的真实含义，读完即可直接上手运行并理解其底层原理。

Overview：Ernie 4.5 家族中的 MoE 文本模型

Ernie 4.5 MoE 是百度 Ernie 4.5 Model Family 发布的一部分。这个家族包含多种不同架构与规格的模型，本文档聚焦其中的基础纯文本 MoE 模型，公开了两个规模档位：

• 总参数 21B、激活参数约 3B 的模型（文档示例使用的正是该档位，checkpoint 名为 baidu/ERNIE-4.5-21B-A3B-PT）；
• 总参数 300B、激活参数约 47B 的更大规模模型。

从架构上看，Ernie 4.5 MoE 以标准 Llama 的解码器结构为核心，叠加了一套基于 Mixtral 风格、并额外引入共享专家（shared experts）的 MoE 前馈层。这一架构事实可以直接从源码得到印证：本仓库通过 modular 方式生成该模型，其源文件 src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py 直接 from ..llama.modeling_llama import LlamaAttention, LlamaRMSNorm、并从 ..mixtral.modeling_mixtral 引入 MoE 相关组件，然后在其基础上做定制（生成产物见 modeling_ernie4_5_moe.py）。

同家族的其他模型分别位于仓库文档 Ernie 4.5 文档 与 Ernie 4.5 VL MoE 文档（注意该文档链接原样为 ernie4_5_vl_moe.md，在本仓库中对应上面这条路径），阅读时可按需对照。

快速上手：加载模型并生成文本

文档给出了最直接的文本生成用法，即通过 AutoModelForCausalLM 与 AutoTokenizer 加载 checkpoint，并用 apply_chat_template 组装带角色（role）的对话消息后调用 model.generate 完成续写：

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "baidu/ERNIE-4.5-21B-A3B-PT"

# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
)

# prepare the model input
inputs = tokenizer("Hey, are you conscious? Can you talk to me?", return_tensors="pt").to(model.device)
prompt = "Hey, are you conscious? Can you talk to me?"
messages = [
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], add_special_tokens=False, return_tensors="pt").to(model.device)

# conduct text completion
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=32,
)
output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

# decode the generated ids
generate_text = tokenizer.decode(output_ids, skip_special_tokens=True)

几点实操提示（结合本仓库代码确认）：

• device_map="auto" 依赖 accelerate，会在多卡场景自动把各层分配到不同 GPU 上；
• apply_chat_template(...) 依赖该 checkpoint 的 tokenizer_config.json 中内置的对话模板；
• generate 能力来自 GenerationMixin——文档中列出的 Ernie4_5_MoeForCausalLM 类同时继承 Ernie4_5_MoePreTrainedModel 与 GenerationMixin（见 modeling_ernie4_5_moe.py），因此除 generate 外还支持 greedy_search、beam_search、sample 等全部标准生成接口；
• 为在解码时只对新生成的 token 计算 logits，前向中使用了 logits_to_keep 参数对 lm_head 的输出做切片（见 modeling_ernie4_5_moe.py），这对长 prompt + 多轮续写场景的显存与算力开销有明显优化作用。

分布式推理：Tensor Parallelism（张量并行）

对于 21B 甚至更大规模的 MoE 模型，单卡显存往往放不下。文档给出了在加载时直接启用张量并行的写法——给 from_pretrained 传入 tp_plan="auto"：

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "baidu/ERNIE-4.5-21B-A3B-PT"

# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    tp_plan="auto",
)

# prepare the model input
inputs = tokenizer("Hey, are you conscious? Can you talk to me?", return_tensors="pt").to(model.device)
prompt = "Hey, are you conscious? Can you talk to me?"
messages = [
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], add_special_tokens=False, return_tensors="pt").to(model.device)

# conduct text completion
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=32,
)
output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

# decode the generated ids
generate_text = tokenizer.decode(output_ids, skip_special_tokens=True)

tp_plan="auto" 能够生效，是因为本仓库为 Ernie 4.5 MoE 在配置类中内置了一张完整的张量并行切分计划（见 configuration_ernie4_5_moe.py 的 base_model_tp_plan），它告诉运行时每个子模块该用何种并行策略：

模块（正则匹配路径）	并行策略	含义
layers.*.self_attn.q_proj / k_proj / v_proj	colwise	注意力 Q/K/V 投影按列切分
layers.*.self_attn.o_proj	rowwise	注意力输出投影按行切分
layers.*.mlp.experts.gate_up_proj	packed_colwise	MoE 路由专家融合的 gate/up 投影按列切分
layers.*.mlp.experts.down_proj	rowwise	MoE 路由专家 down 投影按行切分
layers.*.mlp.experts	moe_tp_experts	专家集合整体按 TP 分组
layers.*.mlp.shared_experts.gate_proj / up_proj	colwise	共享专家 gate/up 按列切分
layers.*.mlp.shared_experts.down_proj	rowwise	共享专家 down 按行切分
layers.*.mlp.gate_proj / up_proj / down_proj	colwise / colwise / rowwise	稠密 MLP 层（非 MoE 层）的标准切分
lm_head	colwise_gather_output	输出头按列切分并汇聚输出

除 TP 外，配置类还提供了流水线并行计划 base_model_pp_plan（embed_tokens、layers、norm 之间的数据流映射）与专家并行计划 base_model_ep_plan（把 gate 标记为 ep_router、专家 FFN 标为 grouped_gemm），说明该模型从设计上即面向大规模多卡部署。需要说明的是，tp_plan 功能由 transformers 的并行化能力提供并依赖较新版本的 accelerate，且要求机器实际具备多 GPU，请在满足环境前提的情况下使用。

低显存部署：Bitsandbytes 4-bit 量化

文档还提供了用 BitsAndBytesConfig 加载 4-bit 量化版本的写法，适合显存紧张但又不愿牺牲太多质量的场景：

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_name = "baidu/ERNIE-4.5-21B-A3B-PT"

# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    quantization_config=BitsAndBytesConfig(load_in_4bit=True),
)

# prepare the model input
inputs = tokenizer("Hey, are you conscious? Can you talk to me?", return_tensors="pt").to(model.device)
prompt = "Hey, are you conscious? Can you talk to me?"
messages = [
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], add_special_tokens=False, return_tensors="pt").to(model.device)

# conduct text completion
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=32,
)
output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

# decode the generated ids
generate_text = tokenizer.decode(output_ids, skip_special_tokens=True)

使用前提：本仓库的测试套件对 4-bit 量化路径有专门的覆盖（见 tests/models/ernie4_5_moe/test_modeling_ernie4_5_moe.py，其中导入了 require_bitsandbytes 与 BitsAndBytesConfig），因此需先安装 bitsandbytes；而 device_map 与量化组合需要 accelerate。对 MoE 模型尤其需要注意的是：配置中 MoE 路由相关参数（如 gate.weight、专家统计的 e_score_correction_bias）被显式声明为以 fp32 严格保留（配置类的 _keep_in_fp32_modules_strict = ["gate.weight", "moe_statics"]，见 modeling_ernie4_5_moe.py），这是为了确保路由精度不因低精度量化而受损。

Ernie4_5_MoeConfig：核心配置项逐一解读

模型的全部超参数集中在配置类 Ernie4_5_MoeConfig（configuration_ernie4_5_moe.py），model_type = "ernie4_5_moe"。类上通过 auto_docstring 声明了默认参考 checkpoint baidu/ERNIE-4.5-21B-A3B-PT。除了从 Llama 继承的通用字段外，以下几组是理解该模型的关键：

通用解码器参数（默认值对应 21B-A3B 档位）

参数	默认值	说明
vocab_size	103424	词表大小
hidden_size	2560	隐藏层维度
intermediate_size	12288	稠密 MLP 的中间维度
num_hidden_layers	28	解码器层数
num_attention_heads	20	注意力头数
num_key_value_heads	4	GQA 的 KV 头数（20/4 的组大小对应注意力实现中的 num_key_value_groups）
hidden_act	"silu"	激活函数（配合门控 MLP 的 SwiGLU 结构）
max_position_embeddings	131072	最大序列长度（128K）
rms_norm_eps	1e-5	RMSNorm 的 epsilon
tie_word_embeddings	True	输入/输出词嵌入共享权重
use_bias	False	各线性投影（含注意力与 MLP）是否带 bias
initializer_range	0.02	参数初始化标准差

MoE 专项参数

文档通过 autodoc 把以下参数定义为 Ernie4_5_MoeConfig 的一部分，字段声明（含默认值与说明）来自配置类源码：

参数	默认值	说明
moe_num_experts	64	被路由的专家数量（routed experts 总数）
moe_k	6	每个 token 选中的专家数量（top-k）。注意代码内部等价映射为 num_experts_per_tok
moe_num_shared_experts	2	对所有 MoE 前向都生效的共享专家数量
moe_intermediate_size	1536	单个路由专家的 FFN 中间维度
moe_layer_start_index	1	MoE 层开始出现的层索引
moe_layer_end_index	-1	MoE 层可能出现的最大层索引；-1 会在 __post_init__ 中换算成 num_hidden_layers - 1
moe_layer_interval	1	MoE 层出现的间隔
moe_norm_min	1e-12	路由归一化时的最小除数（防除零）
output_router_logits	False	是否输出路由 logits
router_aux_loss_coef	0.001	辅助负载均衡损失的加权系数

两点由源码推断的重要细节：

1. MoE 层的放置是可配置的。在 Ernie4_5_MoeDecoderLayer 的初始化中，只有当 (layer_idx + 1) % moe_layer_interval == 0、layer_idx >= moe_layer_start_index 且 layer_idx <= moe_layer_end_index 同时成立时，该层才使用 Ernie4_5_MoeSparseMoeBlock，否则退化为普通稠密 Ernie4_5_MoeMLP。默认值（start=1、interval=1、end 为最后一层）意味着：第 0 层是稠密 MLP，第 1 层起全部为 MoE 层。
2. 字段别名映射：attribute_map 把通用名称 num_experts → moe_num_experts、num_experts_per_tok → moe_k（见 configuration_ernie4_5_moe.py），因此代码内部统一使用 config.num_experts、config.num_experts_per_tok 访问，与 Mixtral 系列保持一致。

另外配置类还声明了 base_model_tp_plan、base_model_pp_plan、base_model_ep_plan 三套并行计划，前文已展开过 TP 表；default_theta = 500000.0 作为 RoPE 的默认 base。若想按文档示例手动初始化模型：Ernie4_5_MoeConfig() 得到一个 21B 档位默认配置，再 Ernie4_5_MoeModel(configuration) 即可构造。

架构与源码实现：Llama 核心上的 MoE 变体

从 modeling_ernie4_5_moe.py 的类结构看，其整体是一个标准的 Llama 风格因果解码器，但把每一层的前馈部分按上文规则替换为「共享专家 + 稀疏路由专家」。逐模块梳理如下。

归一化与位置编码

• Ernie4_5_MoeRMSNorm：等价于 T5LayerNorm 的实现，在 fp32 下计算方差并 rsqrt 归一化（modeling_ernie4_5_moe.py）。解码器每个子层前后（input_layernorm、post_attention_layernorm）各放一个。
• Ernie4_5_MoeRotaryEmbedding：支持从 rope_parameters 读取 rope_type，非默认类型时调用 ROPE_INIT_FUNCTIONS[rope_type] 计算逆频率；默认按 rope_theta（default_theta=500000.0）计算（L83-L121）。
• apply_rotary_pos_emb 采用 glm 风格的整维旋转与交错拼接（interleave），并强制在 fp32 下完成旋转后再转回原精度（L147-L178）。

注意力

Ernie4_5_MoeAttention 是标准 GQA 多头注意力：Q 投影到 num_attention_heads × head_dim，K/V 投影到 num_key_value_heads × head_dim，缩放因子为 head_dim ** -0.5，支持分组 KV 复用的 repeat_kv。关键点在于其注意力后端是可插拔的：它通过 ALL_ATTENTION_FUNCTIONS.get_interface(config._attn_implementation, eager_attention_forward) 分派到 eager / SDPA / FlashAttention / FlexAttention 实现。配置类 Ernie4_5_MoePreTrainedModel 上显式声明了 _supports_flash_attn = True、_supports_sdpa = True、_supports_flex_attn = True 与 _supports_attention_backend = True（见 L461-L463），因此用户可以自由切换注意力实现。这也解释了模型卡片中的 FlashAttention 与 SDPA 徽标。

MoE 前馈：路由 + 稀疏专家 + 共享专家

这是与 Mixtral 的核心差异所在，由 Ernie4_5_MoeSparseMoeBlock（L373-L400）串起三个子组件：

1. 门控/路由器 Ernie4_5_MoeTopKRouter（L345-L370）：
   • 对 hidden_states 做一次 fp32 的线性打分得到 router_logits，接 softmax；
   • 路由打分前会先加上一个专家打分修正偏置（e-score correction bias），由 Ernie4_5_MoeStatics（L278-L302）保存：一个不可训练的零初始化偏置张量，其 forward 是一个恒等加零的“技巧”，目的是让只含参数的模块在张量并行下也能正常 isolate 成 DTensor；
   • 取 top-k 索引后 gather 对应权重，再用 torch.clamp(sum, min=moe_norm_min) 做归一化（防除零），最后转回输入 dtype。
2. 路由专家 Ernie4_5_MoeExperts（L305-L342）：所有专家权重以 3D 参数张量存储——gate_up_proj 形状为 (num_experts, 2 * moe_intermediate_size, hidden_size)（融合了 gate 与 up 两个投影），down_proj 形状为 (num_experts, hidden_size, moe_intermediate_size)。前向时先用 one-hot 构造专家掩码找出实际命中的专家，再逐专家做 SiLU 门控乘加与加权累加；实现类上还标有 @use_experts_implementation，表示可被 experiments 框架级自定义实现替换。
3. 共享专家：当 moe_num_shared_experts > 0 时，块内会额外构造一个 Ernie4_5_MoeMLP，其中间维度为 moe_intermediate_size × moe_num_shared_experts（L382-L384），即共享专家其实是一个更宽的稠密 FFN。前向计算是 路由专家的加权输出 + 共享专家输出，再把结果加回残差。这种“共享专家”设计让每个 token 在不增加路由负担的情况下，总能从常开专家处获得稳定的表示能力。

整层数据流与输出头

• Ernie4_5_MoeDecoderLayer 沿用了「Pre-Norm → Self-Attention → 残差 → Post-Norm → MLP → 残差」的标准 Llama 顺序（L422-L451），MoE 只是替换其中 MLP 一步。
• Ernie4_5_MoeModel 负责词嵌入、28 层解码器堆叠、末层 RMSNorm，并共享全模型的 RoPE；forward 中通过 create_causal_mask 构建因果掩码，KV 缓存默认为 DynamicCache。
• Ernie4_5_MoeForCausalLM 在其上叠了 lm_head。由于 tie_word_embeddings=True，_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}（L634）把输出头与词嵌入权重绑定，减小存储开销。为兼容更多同族权重，_keys_to_ignore_on_load_unexpected = ["mtp"]，即暂不支持多 token 预测（MTP）头，加载含 MTP 头权重时会自动忽略相关键。

辅助负载均衡损失

MoE 稀疏路由若不加以约束，容易产生“专家坍缩”（少数专家被频繁选中）。为此 Ernie4_5_MoeForCausalLM 在 output_router_logits=True 时计算负载均衡损失：load_balancing_loss_func（L563-L629）按 Switch Transformer 论文中的公式，分层累加每个专家被 top-k 选中的 token 占比与路由概率之和，最终 loss = 总项 × num_experts；有 attention_mask 时所有统计都用 mask 加权以排除 padding token。最终损失为 LM loss + router_aux_loss_coef × aux_loss（默认系数 0.001，见 L699-L708）。返回结构为 MoeCausalLMOutputWithPast（loss、logits、aux_loss、past_key_values、router_logits 等），并支持通过 labels 直接训练。

输出可观测性与推理约束

• Ernie4_5_MoePreTrainedModel 声明了 _can_record_outputs，可在前向时按需记录 router_logits、每层 hidden_states 与 attentions（L467-L471），调试 MoE 路由行为时很有用。
• 由于 use_cache=True 默认开启、keys_to_ignore_at_inference = ["past_key_values"]，解码期会自然复用 KV 缓存，配合上一节提到的 logits_to_keep 实现高效增量生成。

测试与验证：质量保障的落点

仓库在 tests/models/ernie4_5_moe/test_modeling_ernie4_5_moe.py 中为该模型提供了完整测试套件，可作为验证与二次开发的参考基准：

• Ernie4_5_MoeModelTester(CausalLMModelTester) 与 Ernie4_5_MoeModelTest(CausalLMModelTest) 表明它复用了 transformers 统一因果 LM 测试框架，覆盖前向、反向、KV 缓存、梯度检查点等通用行为；
• 专门的 test_load_balancing_loss 会构造 num_experts=3、output_router_logits=True 的配置，校验 router_logits 形状为 (batch × seq, num_experts)、aux_loss ≈ 2，并验证加入被 mask 掉的 padding token 不会改变负载均衡损失——这正是论文与实现中排除 padding 语义的直接回归测试；
• test_flash_attn_2_equivalence（标记 @require_flash_attn、@require_torch_accelerator 与 @slow）用 bf16 对比 flash_attention_2 与 eager 两种注意力后端的输出一致性（容差 atol=rtol=1e-2），印证了前文所述多注意力后端支持的实现结论；
• 测试文件还导入了 BitsAndBytesConfig 与 require_bitsandbytes，与文档中的 4-bit 量化用法一一对应。

小结与后续阅读

总而言之，Ernie 4.5 MoE 在 transformers 中的集成可以概括为：以 21B/300B 两档 MoE 参数规模面世，采用 Llama 解码器骨架 + Mixtral 式 top-k 稀疏路由 + 百度自研共享专家与路由打分修正偏置的组合架构。实操层面只需三行核心调用（AutoTokenizer + AutoModelForCausalLM.from_pretrained(device_map="auto") + model.generate）即可完成单卡推理，再按需叠加 tp_plan="auto" 多卡张量并行或 BitsAndBytesConfig(load_in_4bit=True) 显存优化；进阶用户可直接阅读 configuration_ernie4_5_moe.py 调整 MoE 层密度、专家数与路由拓扑，并结合 tests/models/ernie4_5_moe/test_modeling_ernie4_5_moe.py 验证改动。

若想横向了解同家族的另外两支架构，可继续阅读仓库文档 Ernie 4.5（稠密文本模型） 与 Ernie 4.5 VL MoE（多模态 MoE）。另外，本文中反复出现的建模文件 modeling_ernie4_5_moe.py 与 modular_ernie4_5_moe.py 是一对「自动生成产物 ↔ modular 源」关系：前者由后者生成、不可手改（文件头部有 CI 强制提示），对模型做任何源码级修改都应落在 modular 源文件上，再运行转换脚本重新生成——这也是本仓库所有 modular 模型的标准开发流程。