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 仓库中一个具有11B110 亿参数的单向语言模型其设计思路直接源自 Megatron-LM 论文Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism。与普通 Transformer 语言模型的关键区别在于它采用层内模型并行intra-layer model parallelism将每一层的参数拆分到 8 张 GPU 上前向/反向传播过程中通信激活值与梯度从而让单卡无法容纳的超大模型得以训练与推理。在数据与词表层面Megatron-11b 使用与 RoBERTa 相同的数据和相同的字节级 BPEbyte-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 架构参数参数取值embed_dimdecoder_embed_dim3072ffn_dimdecoder_ffn_embed_dim3072 × 6 18432layersdecoder_layers72attention headsdecoder_attention_heads32以上数值可在源码注册的架构函数中直接找到对应关系。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 inputbuild_model中直接raise NotImplementedError。2.2 训练超参参数取值batch sizebsz512num_updates300,000peak_lr1.5e-04lr schedulerinverse_sqrtclip norm0.0三、预训练模型下载官方提供预训练权重megatron_11b11B 参数、压缩包约 19GB模型描述参数量文件大小下载megatron_11bmegatron_11b 单向语言模型11B19GBmegatron_11b.tar.gz下载后解压即得到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-18×V100-32GB上验证通过。4.1 关键参数逐项解读参数作用与说明--model-parallel-size 8模型并行度本场景把每层参数拆分到 8 张 GPU--distributed-world-size 8全局分布式规模。单节点 8 卡时与模型并行度一致多节点时可继续增大以叠加数据并行--criterion vocab_parallel_cross_entropy词表并行交叉熵损失配合输出层切分使用详见下文--arch transformer_lm_megatron_11b选择已注册的模型并行语言模型架构--share-decoder-input-output-embed共享输入/输出词嵌入。模型并行下必须开启否则会报错--memory-efficient-fp16显存高效 FP16 训练缓解 11B 参数带来的显存压力--task language_modeling/--sample-break-mode none/--tokens-per-sample 1024语言建模任务配置样本不按句子边界截断每样本最长 1024 token--batch-size 2每张 GPU 上的 batch size配合 8 卡与梯度累积等效实现整体 bsz 5124.2 底层实现为什么需要这些参数从 fairseq/model_parallel/models/transformer_lm.py 源码可以看出模型并行语言模型的核心机制词表对齐build_model中执行task.source_dictionary.pad_to_multiple_(args.model_parallel_size * 8)与task.target_dictionary.pad_to_multiple_(...)把词表补齐到8 * 模型并行度的整数倍便于在 GPU 间均匀切分词表。并行词嵌入build_embedding使用VocabParallelEmbedding构建词嵌入L74-L83并按embed_dim ** -0.5初始化、将 padding 行置零。输出层切分在 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。并行层结构fairseq/model_parallel/modules/transformer_layer.py 中FFN 的两个线性层分别使用ColumnParallelLineargather_outputFalse与RowParallelLinearinput_is_parallelTrue注意力使用ModelParallelMultiheadAttention从而把注意力头与 FFN 参数都拆分到多卡。损失计算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语言建模基准上的结果模型Valid perplexityTest perplexitymegatron_11b10.6410.54六、在 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-103wget 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 binarizefairseq-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_PATHwikitext103-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的输出对齐到原始文本口径。八、常见问题与注意事项必须先安装 megatron 子模块git submodule update --init fairseq/model_parallel/megatron否则ModelParallelTransformerLanguageModel.build_model与VocabParallelCrossEntropyCriterion都会抛出 ImportError详见 transformer_lm.py 与 vocab_parallel_cross_entropy.py。训练必须开启--share-decoder-input-output-embed模型并行的输出层实现依赖共享嵌入output_layer中直接对非共享情况抛出 NotImplementedError。模型并行实现暂不支持character embeddings、adaptive inputadaptive softmax以及 quantization noiseq_noise 0时在并行层直接抛NotImplementedError。评估与训练损失函数不同训练用vocab_parallel_cross_entropy词表分片内计算损失评估用cross_entropygather 完整词表后计算。显存与硬件前提README 中验证环境为 DGX-18×V100-32GB训练需开启--memory-efficient-fp16评估同样要求 8 卡才能加载模型。困惑度口径所有报告值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.pyLM 评估入口fairseq_cli/eval_lm.py【免费下载链接】fairseqFacebook AI Research Sequence-to-Sequence Toolkit written in Python.项目地址: https://gitcode.com/gh_mirrors/fa/fairseq创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考