企业官网建设流程全解析

fairseq 中的 Megatron-11b：基于 Megatron-LM 的 110 亿参数语言模型并行训练与 Wikitext-103 评估实战

【免费下载链接】fairseqFacebook AI Research Sequence-to-Sequence Toolkit written in Python.项目地址: https://gitcode.com/gh_mirrors/fa/fairseq

导读

本文围绕 fairseq 仓库中 examples/megatron_11b/README.md 展开，系统讲解 Megatron-11b——一个基于 Megatron-LM 论文方案、参数量达 110 亿的单向（unidirectional）语言模型：从架构参数、训练超参到 8 卡模型并行训练命令，再到 Wikitext-103 上完整的"下载 → 反 token 化 → BPE 编码 → binarize → 评估与重归一化"全流程。读完本文，你将掌握如何在 fairseq 中复现 Megatron-11b 的模型并行训练，并理解为何评估时需要做困惑度（perplexity）重归一化。

一、Megatron-11b 概览

Megatron-11b 是 fairseq 仓库中一个具有11B（110 亿）参数的单向语言模型，其设计思路直接源自 Megatron-LM 论文（Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism）。与普通 Transformer 语言模型的关键区别在于：它采用层内模型并行（intra-layer model parallelism），将每一层的参数拆分到 8 张 GPU 上，前向/反向传播过程中通信激活值与梯度，从而让单卡无法容纳的超大模型得以训练与推理。

在数据与词表层面，Megatron-11b 使用与 RoBERTa 相同的数据和相同的字节级 BPE（byte-pair encoding）分词方案。因此，后续评估流程中会复用 RoBERTa 的 GPT-2 BPE 编码脚本（examples/roberta/multiprocessing_bpe_encoder.py），并且模型自带与 RoBERTa 风格一致的dict.txt词表。

前置提醒：模型并行实现依赖 fairseq 的 megatron 子模块。从 fairseq/model_parallel/models/transformer_lm.py 的源码可以看到，构建模型与vocab_parallel_cross_entropy损失都会在子模块缺失时抛出提示，需要先执行git submodule update --init fairseq/model_parallel/megatron安装该子模块。

二、模型架构与训练配置

2.1 架构参数

以上数值可在源码注册的架构函数中直接找到对应关系。fairseq/model_parallel/models/transformer_lm.py 中注册了transformer_lm_megatron_11b架构：

@register_model_architecture( "model_parallel_transformer_lm", "transformer_lm_megatron_11b" ) def transformer_lm_megatron_11b(args): args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 3072) args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 3072 * 6) args.decoder_layers = getattr(args, "decoder_layers", 72) args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 32) args.dropout = getattr(args, "dropout", 0.1) args.attention_dropout = getattr(args, "attention_dropout", 0.1) args.activation_fn = getattr(args, "activation_fn", "gelu") base_lm_architecture(args)

值得注意的细节：

• 同文件还注册了不带_11b后缀的transformer_lm_megatron架构（L146-L155），两者唯一区别是 FFN 维度：transformer_lm_megatron为3072 * 4，而transformer_lm_megatron_11b为3072 * 6。
• base_lm_architecture中强制设置args.decoder_normalize_before = True，注释明确说明"没有此项训练不稳定"（Model training is not stable without this）。
• 模型并行实现不支持character embeddings 与 adaptive input（build_model中直接raise NotImplementedError）。

2.2 训练超参

三、预训练模型下载

官方提供预训练权重megatron_11b（11B 参数、压缩包约 19GB）：

下载后解压即得到megatron_11b/目录（内含model.pt与dict.txt），具体命令见下文"评估流程"第一步。

四、模型并行训练命令（单节点 8 卡）

Megatron-11b 参数量过大，无法在单卡上训练。按照原版 Megatron 的做法，fairseq 采用层内模型并行：每层参数拆分到多张 GPU，前向/反向过程中分别通信激活值与梯度；同时使用vocab_parallel_cross_entropy损失函数把词表维度的 logits 也按模型并行切分，避免输出层在单卡上重建整个词表分布。

以下命令演示在单节点（8 张 GPU）上训练 Megatron-11b。若有多节点，可以通过增大--distributed-world-size将其与数据并行（data parallel）结合使用：

fairseq-train <DATA_PATH> \ --distributed-world-size 8 \ --memory-efficient-fp16 \ --num-workers 2 \ --model-parallel-size 8 \ --criterion vocab_parallel_cross_entropy \ --task language_modeling \ --sample-break-mode none \ --tokens-per-sample 1024 \ --arch transformer_lm_megatron_11b \ --share-decoder-input-output-embed \ --optimizer adam --adam-betas "(0.9, 0.98)" --adam-eps 1e-08 --clip-norm 0.0 \ --lr-scheduler inverse_sqrt --lr 0.00015 \ --warmup-updates 3000 --weight-decay 0.01 \ --dropout 0.1 --attention-dropout 0.1 \ --batch-size 2 \ --max-update 300000

注：以上命令在DGX-1（8×V100-32GB）上验证通过。

4.1 关键参数逐项解读

4.2 底层实现：为什么需要这些参数

从 fairseq/model_parallel/models/transformer_lm.py 源码可以看出模型并行语言模型的核心机制：

1. 词表对齐：build_model中执行task.source_dictionary.pad_to_multiple_(args.model_parallel_size * 8)与task.target_dictionary.pad_to_multiple_(...)，把词表补齐到8 * 模型并行度的整数倍，便于在 GPU 间均匀切分词表。
2. 并行词嵌入：build_embedding使用VocabParallelEmbedding构建词嵌入（L74-L83），并按embed_dim ** -0.5初始化、将 padding 行置零。
3. 输出层切分：在 fairseq/model_parallel/models/transformer.py#L107-L121 的output_layer中，特征先经copy_to_model_parallel_region复制到模型并行区域再做投影；只有当 criterion 是vocab_parallel_cross_entropy时才不需要 gather 回完整词表——否则仍需gather_from_model_parallel_region，这正是命令中必须配套使用该损失函数的原因。同时，该实现要求共享输入输出嵌入，否则直接raise NotImplementedError。
4. 并行层结构：fairseq/model_parallel/modules/transformer_layer.py 中，FFN 的两个线性层分别使用ColumnParallelLinear（gather_output=False）与RowParallelLinear（input_is_parallel=True），注意力使用ModelParallelMultiheadAttention，从而把注意力头与 FFN 参数都拆分到多卡。
5. 损失计算：fairseq/model_parallel/criterions/vocab_parallel_cross_entropy.py 注册了vocab_parallel_cross_entropy损失，调用 megatron 子模块中的vocab_parallel_cross_entropy在切分的词表分片上并行计算交叉熵，随后按非 padding token 求和，并支持reduce_metrics汇总输出 nll_loss 与 ppl。

五、评估结果

模型在Wikitext-103语言建模基准上的结果：

六、在 Wikitext-103 上评估megatron_11b：完整实操

评估链路的核心难点在于：Megatron-11b 使用字节级 BPE，要求输入为原始（未 token 化）文本，而 Wikitext-103 自带的是已 token 化的数据。因此需要先反 token 化，再做 BPE 编码，最后 binarize 并评估。整个流程共 6 步。

第 1 步：下载 Megatron-11b

# 警告：该文件有 19GB wget https://dl.fbaipublicfiles.com/fairseq/models/model_parallel/megatron_11b.tar.gz tar -xzvf megatron_11b.tar.gz

解压后得到megatron_11b/目录，内含评估所需的model.pt（权重）与dict.txt（词表）。

第 2 步：下载 Wikitext-103

wget https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-103-raw-v1.zip unzip wikitext-103-raw-v1.zip

第 3 步：反 token 化测试集

python -m examples.megatron_11b.detok wikitext-103-raw/wiki.test.raw > wikitext-103-raw/wiki.test.detok

该脚本实现在 examples/megatron_11b/detok.py 中：它用sacremoses.MosesDetokenizer对按空格切分的 token 做 Moses 反 token 化，并额外清理 BPE 风格标记与空格：移除@/@，将=/=还原为=，将–还原为–。这一步把 Wikitext-103 已 token 化的测试集恢复成接近原始文本的形式，供字节级 BPE 使用。

第 4 步：BPE 编码

wget -N 'https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/encoder.json' wget -N 'https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/vocab.bpe' python -m examples.roberta.multiprocessing_bpe_encoder \ --encoder-json encoder.json \ --vocab-bpe vocab.bpe \ --inputs "wikitext-103-raw/wiki.test.detok" \ --outputs "wikitext-103-raw/wiki.test.bpe" \ --workers 60

这一步复用 RoBERTa 的多进程 BPE 编码器 examples/roberta/multiprocessing_bpe_encoder.py：脚本基于fairseq.data.encoders.gpt2_bpe.get_encoder加载 GPT-2 BPE 的encoder.json与vocab.bpe，用multiprocessing.Pool并行编码（--workers指定进程数，默认 20，示例中调高到 60 以加速），并支持--keep-empty保留空行。

第 5 步：fairseq binarize

fairseq-preprocess \ --only-source \ --testpref wikitext-103-raw/wiki.test.bpe \ --srcdict megatron_11b/dict.txt \ --destdir wikitext103-bin

使用模型自带的dict.txt作为源词表，仅处理测试集（--only-source），输出到wikitext103-bin/。注意因为 Megatron-11b 需要模型并行（词表按 8 卡切分），词表应能被模型并行度对齐，使用模型自带词表即可保证一致。

第 6 步：评估困惑度

由于我们对测试集做了反 token 化与 BPE 处理，fairseq-eval-lm报出的困惑度是未归一化的，需要按 token 数比例重归一化（详见下一节）。

DATA_PATH=wikitext103-bin/ fairseq-eval-lm \ $DATA_PATH \ --path megatron_11b/model.pt \ --task language_modeling \ --gen-subset test \ --batch-size 8 \ --criterion cross_entropy \ --context-window 992 \ --distributed-world-size 8 \ --model-parallel-size 8 # 期望 PPL（未归一化）：[8.46] # 注：评估命令需要在 8 张 GPU 上运行（发布的模型依赖模型并行）

参数说明：

• --context-window 992：设置 LM 评估的上下文窗口长度。Megatron-11b 训练时tokens-per-sample为 1024，评估窗口取 992 是为留出预测目标的空间（默认的最大目标位置数在源码中为DEFAULT_MAX_TARGET_POSITIONS = 1024，见 transformer_lm.py）。
• --distributed-world-size 8 --model-parallel-size 8：必须在 8 张 GPU 上运行，因为发布模型按 8 卡模型并行切分保存，权重需要重新分布到各卡。
• --criterion cross_entropy：评估阶段使用普通交叉熵（配合gather_from_model_parallel_region汇总完整词表分布），训练阶段才使用vocab_parallel_cross_entropy。

评估入口为 fairseq_cli/eval_lm.py，它加载 checkpoint 后用SequenceScorer逐批计算语言模型困惑度。

七、困惑度重归一化（renormalization）详解

评估得到的未归一化 PPL 为8.46，但发布论文与 README 中报告的测试 PPL 是10.54。差距来自 token 数变化：

• Wikitext-103 原始测试集 token 数为245566；
• 反 token 化 + BPE 编码后 token 数变为270847。

由于困惑度是几何平均的 token 概率（2^(-avg_log2_prob)），当 token 数改变时，总对数概率需要按 token 比例缩放。重归一化公式为：

2 ^ ( log_2(unnormalized_PPL) * (new_token_cnt / orig_token_cnt) )

代入数值：

2 ^ ( log_2(8.46) * (270847 / 245566) ) = 10.54

更一般地，无论你准备的是哪个测试集，只要记下"处理前原始 token 数"与"处理后 BPE token 数"，套用上述公式即可把fairseq-eval-lm的输出对齐到原始文本口径。

八、常见问题与注意事项

1. 必须先安装 megatron 子模块：git submodule update --init fairseq/model_parallel/megatron，否则ModelParallelTransformerLanguageModel.build_model与VocabParallelCrossEntropyCriterion都会抛出 ImportError（详见 transformer_lm.py 与 vocab_parallel_cross_entropy.py）。
2. 训练必须开启--share-decoder-input-output-embed：模型并行的输出层实现依赖共享嵌入（output_layer中直接对非共享情况抛出 NotImplementedError）。
3. 模型并行实现暂不支持：character embeddings、adaptive input（adaptive softmax）以及 quantization noise（q_noise > 0时在并行层直接抛NotImplementedError）。
4. 评估与训练损失函数不同：训练用vocab_parallel_cross_entropy（词表分片内计算损失），评估用cross_entropy（gather 完整词表后计算）。
5. 显存与硬件前提：README 中验证环境为 DGX-1（8×V100-32GB）；训练需开启--memory-efficient-fp16，评估同样要求 8 卡才能加载模型。
6. 困惑度口径：所有报告值（valid 10.64 / test 10.54）均为重归一化后的结果，直接运行fairseq-eval-lm得到的是未归一化值 8.46，二者不可混用。

九、延伸阅读

• 模型并行语言模型完整实现：fairseq/model_parallel/models/transformer_lm.py
• 词表并行交叉熵损失：fairseq/model_parallel/criterions/vocab_parallel_cross_entropy.py
• 模型并行 Transformer（含输出层切分逻辑）：fairseq/model_parallel/models/transformer.py
• 模型并行层（Column/RowParallelLinear、并行注意力）：fairseq/model_parallel/modules/transformer_layer.py
• 反 token 化脚本：examples/megatron_11b/detok.py
• 多进程 BPE 编码脚本：examples/roberta/multiprocessing_bpe_encoder.py
• LM 评估入口：fairseq_cli/eval_lm.py

【免费下载链接】fairseqFacebook AI Research Sequence-to-Sequence Toolkit written in Python.项目地址: https://gitcode.com/gh_mirrors/fa/fairseq