Hugging Face Transformers 中的 NLLB-MoE：稀疏门控 MoE 多语种翻译模型原理、配置与实战指南

NLLB-MoE（No Language Left Behind，Mixture-of-Experts 版本）是 Meta 面向 200+ 语种多语翻译发布的稀疏门控混合专家模型，其官方权重在 Transformers 中以 NllbMoeForConditionalGeneration 系列 API 提供开箱即用的翻译能力。本文以仓库内模型文档 docs/source/en/model_doc/nllb-moe.md 为主线，结合 modeling_nllb_moe.py 与 configuration_nllb_moe.py 的源码实现，系统讲解 NLLB-MoE 的 Top-2 路由机制、与 SwitchTransformers 的实现差异、NllbMoeConfig 全部关键配置项，并给出可直接运行的多语种翻译与微调示例。读完本文你将掌握 NLLB-MoE 的加载、推理、路由行为调参与从任意源语言翻译的完整方法。

模型背景：什么是 NLLB-MoE

NLLB-MoE 来自论文 No Language Left Behind: Scaling Human-Centered Machine Translation（作者包括 Marta R. Costa-jussà、James Cross、Onur Çelebi、Maha Elbayad、Kenneth Heafield、Kevin Heffernan 等数十位研究者），目标是通过机器学习翻译抹平全球语言障碍，特别关注大多数低资源语言。论文团队首先通过母语者访谈理解低资源语言翻译支持的真实需求，随后为缩小低资源与高资源语言之间的性能差距构建了数据集与模型——具体而言，他们开发了一个基于 Sparsely Gated Mixture of Experts（稀疏门控混合专家）的条件计算（conditional compute）模型，其训练数据来自为低资源语言量身定制的数据挖掘技术，并提出了多项架构与训练改进以缓解在数千个任务上训练时的过拟合。评测使用了人工翻译基准 Flores-200，覆盖超过 40,000 个不同翻译方向，并配合覆盖 Flores-200 全部语言的毒性基准评估翻译安全性。

需要特别注意模型族之间的血缘关系，这是本仓库官方文档明确的三个使用要点（Usage tips）：

• M2M100ForConditionalGeneration 同时是 NLLB 与 NLLB-MoE 的基座模型（base model）。这一点在源码中有直接证据：NLLB-MoE 的正弦位置编码模块即是从 m2m_100 中拷贝复用（见 modeling_nllb_moe.py 中 "Copied from transformers.models.m2m_100.modeling_m2m_100.M2M100SinusoidalPositionalEmbedding" 的注释）。
• NLLB-MoE 与 NLLB 模型整体非常相似，但它的 前馈层（feed forward layer）基于 SwitchTransformers 的实现思路构建——即把每个前馈层替换为"路由器 + 若干专家"的稀疏结构。
• NLLB-MoE 的 tokenizer 与 NLLB 模型完全一致（即 NllbTokenizer，基于 sentencepiece，词表大小为 128,112），因此 NLLB 的模型文档 中关于 BCP-47 语言码、src_lang、forced_bos_token_id 的用法均适用于 NLLB-MoE。

该模型在仓库中贡献的源码位置为 src/transformers/models/nllb_moe/，包含 4 个文件：配置文件 configuration_nllb_moe.py、建模代码 modeling_nllb_moe.py、fairseq 原版权重转换脚本 convert_nllb_moe_sharded_original_checkpoint_to_pytorch.py 以及包初始化文件 __init__.py。

NLLB-MoE 与 SwitchTransformers 的实现差异

虽然稀疏前馈层整体借鉴自 SwitchTransformer，但 NLLB-MoE 在令牌如何被路由这一核心机制上有两处重大差异（这是官方模型文档 "Implementation differences with SwitchTransformers" 一节的原始论述）：

1. Top-2 门控 vs Top-1 门控：NLLB-MoE 使用 top-2-gate，即对每个输入令牌，只根据门控网络（路由器）预测出的最高概率选出 top-2 个专家进行前向计算，其余专家被忽略。而 SwitchTransformers 只计算 top-1 概率，因此令牌被真正转发给专家的机会更少（NLLB-MoE 平均每个令牌有两个专家处理，表达容量更大）。
2. 未路由令牌的处理方式：如果一个令牌没有被路由到任何专家（例如由于专家容量耗尽），SwitchTransformers 仍然会把该令牌未经修改的隐藏状态直接加到输出上（类似残差连接）；而 NLLB 的 top-2 路由机制中，这类令牌的隐藏状态会被直接置零（masked），即专家层不会"原样透传"。

第二点在源码中得到了印证：NllbMoeTop2Router 的类注释明确写道 "There is no guarantee that each token is processed by an expert, or that each expert receives at least one token. The router combining weights are also returned to make sure that the states that are not updated will be masked."（不保证每个令牌都被某专家处理、也不保证每个专家至少收到一个令牌；返回组合权重是为了确保未被更新的状态被掩蔽），见 modeling_nllb_moe.py。在训练与推理层面，"未被更新即被掩蔽"意味着超出专家容量或 padding 位置上的令牌在稀疏 MLP 内的输出贡献为零。

源码视角：NllbMoeTop2Router 的完整路由流程

路由器是整个模型最核心的模块，其实现集中在 NllbMoeTop2Router（modeling_nllb_moe.py）。核心数据通路如下：

1. 构造一个 nn.Linear(config.hidden_size, config.num_experts, bias=config.router_bias) 分类器，把每个令牌的隐藏状态映射为 num_experts 维的路由 logits（第 181 行）。
2. 在 route_tokens 中对 logits 做 softmax 得到 router_probs，随后取 argmax 得到 top-1 专家及 one-hot 掩码 top_1_mask。
3. 把 top-1 位置填充为 -inf 后再取 argmax，得到 top-2 专家及其掩码 top_2_mask。若配置了 second_expert_policy="sampling"，会先给 logits 加 Gumbel 噪声再做上述选择；若为 "random"，则按 2 * top_2_max_probs 的概率随机丢弃 top-2 掩码。
4. padding 处理：若传入 padding_mask 且 router_ignore_padding_tokens=False（默认），padding 位置的 top-1/top-2 掩码会被清零——对应上文"padding 令牌不参与路由"。
5. 容量控制：若启用 batch_prioritized_routing，会按路由概率对令牌排序（高概率令牌优先），否则按序列顺序累积计数；结合 expert_capacity 判断 cumsum < capacity 的令牌才能真正入专家。源码对 eval 与训练使用不同容量口径：eval 且 moe_eval_capacity_token_fraction > 0 时 expert_capacity = ceil(fraction * nb_tokens)，否则 capacity = 2 * ceil(nb_tokens / num_experts)（第 269-273 行）。
6. 在容量裁剪后对 top-1/top-2 概率做归一化（normalize_router_probabilities，见 第 198-204 行），最终 router_probs = gates1 + gates2，其中 gate 为概率乘以上一步保留下来的掩码。

真正执行专家计算的模块是 NllbMoeSparseMLP（第 367-383 行），它把隐藏状态展平为 (batch*seq, hidden) 后依次调用 NllbMoeTop2Router 与 NllbMoeExperts。NllbMoeExperts 内部是一个 ModuleDict，保存 expert_0 ... expert_{num_experts-1} 共 num_experts 个 NllbMoeDenseActDense（fc1 -> relu -> dropout -> fc2 的两层 FFN，见 第 318-337 行），只对命中专家做 index_add_ 加权累加，实现"只激活少数专家"的条件计算。此外 NllbMoeExperts 还实现了 moe_token_dropout（专家输出掩蔽 EOM）：训练时用 Dropout2d 掩蔽专家输出，推理时乘以 1 - moe_token_dropout 做补偿（第 349-364 行）。

稀疏层的摆放：encoder/decoder_sparse_step

NLLB-MoE 并不是把每一层都做成 MoE，而是按固定步长在 Transformer 层中穿插稀疏层。从 NllbMoeEncoder/NllbMoeDecoder 的构造函数可以看出，判定逻辑是：

is_sparse = (i + 1) % sparse_step == 0 if sparse_step > 0 else False

即 encoder_sparse_step=4 时，编码器中每 4 层有 1 层是稀疏 MoE 层（第 3、7、11 层，取决于 i+1 是否整除 4）；同样地 decoder_sparse_step=4 意味着解码器中每 4 层有 1 层稀疏。未被标记为 is_sparse 的层仍然使用普通的稠密 NllbMoeDenseActDense FFN，参见 NllbMoeEncoderLayer（第 527-530 行）与 NllbMoeDecoderLayer（第 584-587 行）。这种"稀疏步长"设计直接控制了整个模型的计算量与激活参数量——步长越大，MoE 层越少，激活的专家参数越少，这也正是官方 54B 总参数量下每个 token 实际只激活少量参数的原因。

其余宏观结构（编码器-解码器 seq2seq 架构、自注意力/交叉注意力、正弦位置编码、LayerNorm、残差连接）与 NLLB/M2M100 保持一致；自注意力通过 ALL_ATTENTION_FUNCTIONS 分发，相关源码见 NllbMoeAttention。有一点值得注意：NllbMoePreTrainedModel 目前显式声明 _supports_flash_attn = False、_supports_sdpa = False、_supports_flex_attn = False，并留有 TODO 说明 Flash Attention 因 mask 准备方式与 eager 不一致、SDPA 存在 logits 抖动而暂未启用（第 651-656 行），即当前实现默认走 eager 注意力路径。

NllbMoeConfig 关键配置参数详解

NllbMoeConfig（configuration_nllb_moe.py）继承自 PreTrainedConfig，model_type = "nllb-moe"，并提供与 num_attention_heads/hidden_size/num_hidden_layers 等通用命名的 attribute_map 映射。除常规 seq2seq 配置外，它包含了一批专门控制 MoE 路由行为的参数，按仓库内默认值汇总如下：

参数	默认值	含义
vocab_size	128112	词表大小，与 NLLB 一致（含 200+ 语言的 BCP-47 语言标记 token）
d_model / hidden_size	1024	编码器与解码器隐藏维度（可经 attribute_map 用 hidden_size 访问）
encoder_layers / decoder_layers	12 / 12	编码器 / 解码器 Transformer 层数
encoder_attention_heads / decoder_attention_heads	16 / 16	多头注意力头数
encoder_ffn_dim / decoder_ffn_dim	4096 / 4096	稠密前馈层中间维度（专家内部 FFN 维度）
encoder_layerdrop / decoder_layerdrop	0.05 / 0.05	LayerDrop 概率
activation_function	"relu"	前馈层与专家内激活函数
max_position_embeddings	1024	最大位置编码长度
dropout / attention_dropout	0.1 / 0.1	全连接 / 注意力 dropout
activation_dropout	0.0	激活层 dropout
scale_embedding	True	为 True 时词嵌入乘以 sqrt(d_model) 缩放，见 NllbMoeScaledWordEmbedding
pad_token_id / bos_token_id / eos_token_id	1 / 0 / 2	特殊 token id
decoder_start_token_id	2	解码起始 token
num_experts	128	每个稀疏层中的专家总数
router_bias	False	路由器分类器是否带 bias
router_dtype	"float32"	路由器的计算精度；论文讨论的 selective precision 建议保持 float32（按 第 30-32 行 的说明）
router_ignore_padding_tokens	False	路由时是否忽略 padding 令牌；为 False 时 padding 令牌不会被路由到任何专家
expert_capacity	64	每个专家最多能接收的令牌数；容量耗尽后超出的令牌被丢弃（掩蔽）
encoder_sparse_step	4	编码器中稀疏层的间隔：每隔 4 层出现 1 层 MoE
decoder_sparse_step	4	解码器中稀疏层的间隔：每隔 4 层出现 1 层 MoE
second_expert_policy	"all"	第二个专家的采样策略："all" 表示每个令牌必然拿到 top-2，"sampling" 会加 Gumbel 噪声，"random" 会按概率随机保留 top-2
normalize_router_prob_before_dropping	False	是否在按容量丢弃（capacity dropping）之前先归一化路由概率
batch_prioritized_routing	False	是否在容量裁剪前按路由概率对令牌排序，使高概率令牌优先进入专家
moe_eval_capacity_token_fraction	1.0	验证阶段专家容量按总令牌数的比例计算（取值区间 (0.0, 1.0]）；为负数则与训练一致
moe_token_dropout	0.2	MoE 专家输出掩蔽（EOM）的比率，实现为专家输出上的 Dropout2d
router_z_loss_coef	0.001	路由 z-loss 的系数（z-loss 鼓励路由器 logits 保持较小幅值）
router_aux_loss_coef	0.001	路由辅助负载均衡 loss 的系数
output_router_logits	False	是否在输出中返回 encoder/decoder 各层的路由 logits
tie_word_embeddings	True	是否共享编码器 / 解码器 / LM head 的词嵌入（源码中三处权重通过 _tied_weights_keys 相互绑定）

一点使用提醒：normalize_router_prob_before_dropping 与 batch_prioritized_routing 在配置类的类属性兜底默认值均为 False，但 NllbMoeConfig 的 docstring 按论文原始行为把其语义默认描述为 True。实际加载官方 checkpoint（如 facebook/nllb-moe-54b）时，config.json 中保存的值会覆盖这些兜底默认值，因此路由行为请以 model.config 的实际加载结果为准。

用 NLLB-MoE 做多语种翻译

硬件与内存前提

官方可用的 NLLB-MoE checkpoint（如 54B 版本）需要约 350GB 的存储空间，单机内存也不足以一次性放下全部权重。官方文档明确建议：如果本机 RAM 不够，请务必配合 accelerate 使用——例如用 device_map="auto" 让库自动把各层权重分布到 GPU / CPU / 磁盘（需要安装 accelerate），从而用多卡显存加内存卸载的方式完成加载与生成。

示例一：英语翻译成法语

模型生成时通过 forced_bos_token_id 指定目标语言：把该参数设为目标语言的 token id 即可强制解码器从目标语言的句首标记开始生成。这里需要用到 NLLB 系列语言标记的 BCP-47 代码，例如法语的 fra_Latn。Transformers 的 tokenizer 会把所有语言码映射到一个字典中，可通过 tokenizer.lang_code_to_id["fra_Latn"] 直接取得其 id（也可用 tokenizer.convert_tokens_to_ids("fra_Latn")，两种写法等价）。以下完整示例来自官方文档，可直接运行：

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-moe-54b")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-moe-54b", device_map="auto")

article = "Previously, Ring's CEO, Jamie Siminoff, remarked the company started when his doorbell wasn't audible from his shop in his garage."
inputs = tokenizer(article, return_tensors="pt").to(model.device)

translated_tokens = model.generate(
    **inputs, forced_bos_token_id=tokenizer.lang_code_to_id["fra_Latn"], max_length=50
)
tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
# "Auparavant, le PDG de Ring, Jamie Siminoff, a fait remarquer que la société avait commencé
#  lorsque sa sonnette n'était pas audible depuis son magasin dans son garage."

注意：NLLB-MoE 模型的对应 checkpoint 是 facebook/nllb-moe-54b（仓库集成测试即针对该模型执行英法翻译并断言输出与 fairseq 原实现一致，见 tests/models/nllb_moe/test_modeling_nllb_moe.py）；如需轻量验证路由与生成逻辑，单元测试使用的是 2 个专家的随机小模型 hf-internal-testing/random-nllb-moe-2-experts。Flores-200 基准集中包含全部可用语言及其 BCP-47 代码（语言名采用 语言_文字 形式，如 eng_Latn、fra_Latn、ron_Latn、deu_Latn），需要其它语言时按同样格式替换即可。

示例二：从任意非英语语言出发翻译

NLLB 系 tokenizer 默认把 英语（eng_Latn） 当作源语言。若要翻译其它源语言，必须在 tokenizer 初始化时通过 src_lang 关键字参数指定源语言的 BCP-47 代码，这样语言标记会作为前缀注入源序列。下面是从罗马尼亚语翻译成德语的官方示例：

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-moe-54b", src_lang="ron_Latn")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-moe-54b", device_map="auto")

article = "Şeful ONU spune că nu există o soluţie militară în Siria"
inputs = tokenizer(article, return_tensors="pt").to(model.device)

translated_tokens = model.generate(
    **inputs, forced_bos_token_id=tokenizer.lang_code_to_id["deu_Latn"], max_length=30
)
tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]

两个示例的共同要点可归纳为一张"快速卡片"：

• 源语言 → tokenizer(..., src_lang="<BCP-47>")（不写则默认为 eng_Latn）；
• 目标语言 → model.generate(..., forced_bos_token_id=tokenizer.lang_code_to_id["<BCP-47>"])；
• 解码输出 → tokenizer.batch_decode(tokens, skip_special_tokens=True)。

将翻译扩展为 pipeline 或加载量化版本

由于 NLLB-MoE 的 tokenizer 与 NLLB 完全一致，NLLB 模型文档 中介绍的两种更轻量的用法思路也适用于 MoE 家族：

• 直接使用 pipeline(task="translation", model=..., src_lang=..., tgt_lang=..., device=...) 完成封装式翻译；
• 对参数规模更大的 checkpoint，若显存依然吃紧，可尝试 BitsAndBytesConfig(load_in_8bit=True) 等量化加载路径（配合 quantization_config 与 device_map="auto"）以压缩驻留内存。若想了解仓库内可用的量化后端细节，可参考 量化总览文档。

训练 / 微调 NLLB-MoE：路由损失与 MoE 输出

虽然官方文档的重心在推理生成，但源码 NllbMoeForConditionalGeneration 明确支持 seq2seq 训练。理解下面几个点，微调时才能正确开启 MoE 特有的训练信号：

1. 标签右移与解码输入：传入 labels 时，模型内部调用 shift_tokens_right（第 1018-1031 行）自动构造 decoder_input_ids（首位置填入 decoder_start_token_id，-100 的位置被替换为 pad_token_id），无需手动准备。

2. 路由 logits 开关：前向时把 output_router_logits=True（或配置类中置为 True），输出中就会出现 encoder_router_logits 与 decoder_router_logits。这一行为有专门的单测覆盖：测试断言设置 output_router_logits=True 后，模型输出中的两个字段均不为 None，见 tests/models/nllb_moe/test_modeling_nllb_moe.py。

3. 两项路由辅助损失：当同时满足"传入 labels"与"开启 output_router_logits"时，模型会计算：

   • 负载均衡辅助损失：由 load_balancing_loss_func（第 936-1015 行）实现，按 Switch Transformer 论文公式 (4)-(6) 惩罚专家路由不均衡（各层 logits 拼接后对 top-2 计算"各专家收到的令牌比例 × 平均路由概率"之和），top_k=2 与 NLLB-MoE 的路由数一致；
   • router z-loss：配置中通过 router_z_loss_coef 引入（该系数存储在模型实例上，见 第 1050-1051 行）。

   最终 loss = CrossEntropyLoss(lm_logits, labels) + router_aux_loss_coef * (encoder_aux_loss + decoder_aux_loss)，其中 encoder_aux_loss/decoder_aux_loss 也会随 Seq2SeqMoEOutput 一并返回，方便监控路由分布。

4. 输出结构：NllbMoeForConditionalGeneration.forward 返回 Seq2SeqMoEOutput（来自 modeling_outputs.py 同目录的 modeling_outputs.py），除标准 seq2seq 字段（logits、past_key_values、各层 hidden states 与 attentions）外，还包含 encoder_aux_loss、decoder_aux_loss、encoder_router_logits、decoder_router_logits 等 MoE 专属字段，便于做路由可视化与分析。

如果希望加载 fairseq 原版 NLLB-MoE 分片 checkpoint 进行转换，仓库提供了 convert_nllb_moe_sharded_original_checkpoint_to_pytorch.py，内部实现 key 重命名（rename_fairseq_keys）、num_experts 分片（shard_on_the_fly）等逻辑，可把 fairseq 权重转换为 Transformers 可直接加载的格式。

验证与可信度：单元 / 集成测试

仓库用三层测试保证了 NLLB-MoE 实现的正确性（测试文件 test_modeling_nllb_moe.py）：

• 单元级路由测试：NllbMoeRouterTest 用 num_experts=4、expert_capacity=4 的最小复现配置，直接以固定随机种子喂入 NllbMoeTop2Router.route_tokens，把"掩码 + 专家线性层 + 按路由概率加权"的手写参考实现与 HF 实现逐元素比对，断言与 fairseq 参考均值张量在 atol=1e-4 内一致（见 第 452-504 行）；test_batch_prioritized_routing 则验证 batch_prioritized_routing=True 与 second_expert_policy="random" 组合下的排序路由路径。
• 模型级测试：NllbMoeModelTest 覆盖 NllbMoeModel 与 NllbMoeForConditionalGeneration 两类模型，并通过 output_router_logits=True 断言 encoder/decoder 路由 logits 的输出链路。
• 端到端集成测试：NllbMoeModelIntegrationTests 加载真实的 facebook/nllb-moe-54b checkpoint 进行批量英→法翻译，将 tokenizer.batch_decode 结果与 EXPECTED_FAIRSEQ_TRANSLATION 常量逐一比对，确保与 fairseq 原实现逐字对齐。

这些测试既证明了"Top-2 路由 → 专家加权聚合 → masked 未更新状态"的实现细节，也为读者提供了从轻量随机 checkpoint（2 experts）到完整 54B 权重的验证路径。

延伸学习资源

• 本模型对应的任务指南：Translation 任务指南——翻译任务的数据准备、评估与 pipeline 用法；Summarization 任务指南——由于 NllbMoeForConditionalGeneration 是带语言建模头的条件生成模型，同样可用于摘要生成。
• 本模型的稠密（非 MoE）姊妹版本见 NLLB 模型文档，两者共享 tokenizer 与大多数结构；对照阅读可更直观地理解 MoE 化带来的差异。
• 深入源码：模型实现与路由细节见 modeling_nllb_moe.py，配置全量参数见 configuration_nllb_moe.py，fairseq 权重转换见 convert_nllb_moe_sharded_original_checkpoint_to_pytorch.py，验证逻辑见 test_modeling_nllb_moe.py。

综上，NLLB-MoE 在 Transformers 中的实现精髓可概括为一句话：以 M2M100/NLLB 的 seq2seq 骨架为底、以 SwitchTransformer 风格的稀疏前馈为形、以 fairseq 的 top-2 容量裁剪路由为魂。理解 NllbMoeTop2Router.route_tokens 中的"top-2 选择 → padding 掩蔽 → 容量裁剪 → 概率归一化 → 加权聚合"五步流水线，再加上对 NllbMoeConfig 中路由器相关参数的把握，无论是做多语种翻译推理、路由行为分析还是自定义稀疏步长微调，都能做到有的放矢。