
Megatron-11b 模型并行语言模型从 11B 参数训练到 Wikitext-103 评估的完整实战指南【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm导读本文围绕 fairseq 仓库中 examples/megatron_11b 所发布的Megatron-11b 单向语言模型Unidirectional Language Model展开完整介绍该 11B 参数模型如何基于 Megatron-LM 的层内模型并行intra-layer model parallelism思想在 8 张 GPU 上切分训练并给出从权重下载、数据预处理到困惑度Perplexity评估与归一化的端到端可复现流程。读完本文你将掌握 fairseq 中模型并行训练的核心参数组合、vocab_parallel_cross_entropy损失的计算原理以及一套可以照搬的 Wikitext-103 评测命令。1. 模型概述基于 Megatron-LM 的 11B 单向语言模型Megatron-11b 是一个拥有约11B110 亿参数的单向语言模型其设计直接沿袭了 Megatron-LM 论文中的思路采用层内模型并行将每一层的参数切分到 8 张 GPU 上在前向/反向传播过程中分别通信激活值与梯度从而在单卡显存放不下完整模型的情况下完成训练。与当时多数预训练模型不同Megatron-11b 的训练数据与分词方式与 RoBERTa 完全一致——使用相同的语料以及相同的字节级 BPEByte-Pair Encoding编码。因此它的输入要求是未分词untokenized的原始文本这也直接影响了后续评估时的数据预处理流程需要先 detokenize 再应用 BPE详见第 5 节。在仓库中README_FAIRSEQ.md 的更新日志明确记载April 2020: Initial model parallel support and 11B parameters unidirectional LM released即 fairseq 的模型并行支持与 Megatron-11b 是同期引入的二者属于同一套代码体系。2. 模型架构与训练配置2.1 架构参数Megatron-11b 的架构参数如下表所示见 examples/megatron_11b/README.md参数值embed_dim3072ffn_dim3072 × 6 18432layers72attention heads32这些默认值在源码中均有对应的注册架构。查看 fairseq/model_parallel/models/transformer_lm.py 中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)可以看到该架构同时指定了 GELU 激活函数、0.1 的 dropout 与 attention dropout并调用base_lm_architecture完成其余默认值的填充该函数位于同一文件的 L91-L148其中decoder_normalize_before True是保证模型并行训练稳定性的关键设置。2.2 训练细节官方公布的训练超参数如下见 README.md参数值batch size512num_updates300,000peak_lr1.5e-04lr schedulerinverse_sqrtclip norm0.0其中clip norm 0.0表示不做梯度裁剪。模型并行场景下梯度范数需要在模型并行组内做 all-reduce 聚合才能正确裁剪fairseq 的 MegatronTrainer 中正是通过distributed_utils.all_reduce(..., groupdistributed_utils.get_model_parallel_group())来聚合各并行分片上的梯度范数。3. 模型并行训练命令详解Megatron-11b 参数量过大无法在单张 GPU 上训练。fairseq 采用与 Megatron-LM 一致的层内模型并行方案每层参数切分到多张 GPU前向/反向过程中通信激活值与梯度同时损失计算也被切分——使用vocab_parallel_cross_entropy准则criterions/vocab_parallel_cross_entropy.py。官方给出的单节点8 卡训练命令见 README.mdfairseq-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;关键参数说明--model-parallel-size 8指定模型并行切分数这里假设每台节点有 8 张 GPU 用于切分模型参数。模型并行组内每个 rank 只持有参数的 1/8。--distributed-world-size 8总分布式规模。如果有多台节点可通过增大该值将模型并行与数据并行组合起来例如 2 节点 × 8 卡模型并行 world size 16同时具备数据并行维度。--criterion vocab_parallel_cross_entropy使用词表并行交叉熵。词表embedding 输出投影同样按模型并行组切分每个 rank 只对本地词表分片计算 logits 与损失。--arch transformer_lm_megatron_11b加载上文所述的 72 层 × 3072 维架构注册项。--share-decoder-input-output-embed共享输入/输出词嵌入。源码 transformer.py 的output_layer明确要求模型并行训练必须开启该选项否则会抛出NotImplementedError。--memory-efficient-fp16半精度训练配合 8×V100-32Gb 硬件单卡 batch size 设为 2。--sample-break-mode none与--tokens-per-sample 1024语言模型任务中样本按 1024 token 截断且不做额外分隔处理。官方提示该命令在DGX-18×V100-32Gb上测试通过。3.1 模型并行的源码实现ColumnParallel 与 RowParallel模型并行的底层实现分布在fairseq/model_parallel/目录modules/transformer_layer.py解码器层的两个 FFN 全连接分别替换为ColumnParallelLinearbuild_fc1gather_outputFalse与RowParallelLinearbuild_fc2input_is_parallelTrue实现 先切列、后切行 的标准 Megatron 张量并行范式。modules/multihead_attention.py多头注意力中 Q/K/V 投影用ColumnParallelLinear输出投影用RowParallelLinear头数按模型并行大小切分self.num_heads_partition num_heads // self.model_parallel_size并要求头数必须能被并行度整除32 个头 ÷ 8 卡 每卡 4 个头每个 rank 只计算自己负责的头分片。此外注意力 dropout 通过get_cuda_rng_tracker().fork()隔离各并行分片的随机数状态。models/transformer.pyoutput_layer先将特征copy_to_model_parallel_region(features)复制到模型并行区域再执行self.output_projection即词表并行嵌入的转置。当且仅当损失函数是vocab_parallel_cross_entropy时才不做gather_from_model_parallel_region从而让 logits 保持切分状态直接在并行组内完成交叉熵计算避免把整个词表 logits 汇聚到单卡上。models/transformer_lm.py构建模型时会调用pad_to_multiple_(args.model_parallel_size * 8)将词典大小补齐到 64 的倍数保证词表能被并行切分整除。从源码结构可以推断这套模型并行实现与 fairseq 内嵌的 Megatron 子模块fairseq/model_parallel/megatron/mpu提供VocabParallelEmbedding、ColumnParallelLinear、RowParallelLinear、vocab_parallel_cross_entropy等原语紧密耦合若未初始化该子模块模型、损失、训练器构建时会抛出安装提示。4. 预训练权重与评测结果4.1 预训练模型下载官方发布了一个模型权重见 README.md模型描述# 参数文件大小下载megatron_11b11B 单向语言模型11B19GBmegatron_11b.tar.gz下载命令文件约 19GBwget https://dl.fbaipublicfiles.com/fairseq/models/model_parallel/megatron_11b.tar.gz tar -xzvf megatron_11b.tar.gz4.2 Wikitext-103 评测结果Megatron-11b 在 Wikitext-103 上的官方结果见 README.md模型Valid PPLTest PPLmegatron_11b10.6410.54注意这里的困惑度是经过token 数归一化renormalization后的数值具体换算过程见第 6 节。5. Wikitext-103 评估的完整预处理流水线由于 Megatron-11b 使用字节级 BPE 且要求原始未分词输入而 Wikitext-103 数据集本身是已分词的因此评估前需要完成 detokenize → BPE → binarize 三步预处理。5.1 下载 Wikitext-103wget https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-103-raw-v1.zip unzip wikitext-103-raw-v1.zip5.2 Detokenize 测试集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对每行按空格切分后合并还原并额外清理 BPE/分词产生的残留标记例如去除 与 片段对应 GPT-2 BPE 的字节回退标记、规整 / 为、把 – 还原为–。同时该脚本支持fileinput.hook_compressed可直接处理压缩输入。5.3 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;BPE 编码脚本位于 examples/roberta/multiprocessing_bpe_encoder.py其输入输出 token 数与megatron_11b/dict.txt中的词表保持一致。--workers 60开启多进程并行编码以加速。5.4 Fairseq 二值化fairseq-preprocess \ --only-source \ --testpref wikitext-103-raw/wiki.test.bpe \ --srcdict megatron_11b/dict.txt \ --destdir wikitext103-bin;这里直接使用权重包内自带的megatron_11b/dict.txt作为源词典保证 token 编号与训练时一致--only-source表示只有源侧数据语言模型任务。6. 模型并行评估与困惑度归一化6.1 计算未归一化困惑度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; # Expected PPL (unnormalized_ppl): [8.46] # Note: the eval command needs to run on 8 GPUs for the released model评估要点--model-parallel-size 8与--distributed-world-size 8官方发布的权重是 8 路模型并行切分后的分片因此必须在 8 张 GPU 上运行才能正确加载与推理。--context-window 992滑动窗口上下文长度 992小于训练时的 1024配合 batch size 8 进行滑动窗口式困惑度计算。评估阶段使用普通cross_entropy准则即可此时模型输出层会通过gather_from_model_parallel_region汇聚完整 logits参见 transformer.py。官方期望的未归一化困惑度约为8.46。6.2 困惑度归一化公式由于测试集经过了 detokenize 和 BPEtoken 数量发生变化fairseq-eval-lm报告的困惑度需要按下式重新归一化见 README.md2 ^ ( log_2(unnormalized_PPL) * (new_token_cnt / orig_token_cnt) )对于 Wikitext-103 测试集原始已分词token 数245566detokenize BPE 后 token 数270847代入数值2 ^ ( log_2(8.46) * (270847 / 245566) ) 10.54该值与官方表格中报告的 Test PPL10.54完全吻合。理解这一换算很关键任何对测试集做过改动的评估直接报告的困惑度都需用 token 数比例校正后才能与其他公开结果横向比较。7. 总结与延伸Megatron-11b 是 fairseq 模型并行能力的重要里程碑它同时验证了三条技术路线Megatron 风格层内张量并行在纯语言模型上的可扩展性——72 层、3072 维嵌入、FFN 维数达 18432 的巨型模型仅靠 8×V100-32Gb 即可训练词表并行交叉熵vocab_parallel_cross_entropy避免了输出 logits 在单卡聚合带来的显存与通信瓶颈模型并行 × 数据并行组合通过--model-parallel-size与--distributed-world-size两个维度的组合可将训练扩展到多节点集群。如果希望在当前仓库中进一步探究实现细节建议按以下路径阅读模型注册与架构默认值fairseq/model_parallel/models/transformer_lm.py模型并行解码器与输出层逻辑fairseq/model_parallel/models/transformer.py并行注意力与 FFN 层fairseq/model_parallel/modules/multihead_attention.py、fairseq/model_parallel/modules/transformer_layer.py词表并行损失fairseq/model_parallel/criterions/vocab_parallel_cross_entropy.py模型并行训练器梯度范数聚合、RNG 状态保存fairseq/model_parallel/megatron_trainer.py评测预处理工具examples/megatron_11b/detok.py、examples/roberta/multiprocessing_bpe_encoder.py最后提醒训练与评估命令中的fairseq-train/fairseq-eval-lm/fairseq-preprocess均来自 decoding/IAD/fairseq 目录下的 fairseq 工具链--model-parallel-size 8的评估命令必须跑满 8 张 GPU 才能正确加载官方发布的分片权重。【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考