Megatron-LM BERT 分布式预训练实战Docker 启动、340M 基线脚本与 4B/20B 扩展配置全解析【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本指南聚焦 Megatron-LM 仓库中 BERT 预训练的完整实战链路如何用 Docker 容器启动训练、examples/bert/train_bert_340m_distributed.sh中 340M 参数基线脚本的每一项参数含义以及如何通过改写模型并行配置将其扩展到 4B 与 20B 规模。读完本文你将能独立配置数据路径、校验参数组合并把该示例脚本迁移到自己的多机多卡环境中运行。1. 前置准备数据、词表与运行环境BERT 预训练需要三类关键输入词汇表文件vocab fileMegatron-LM 的 BERT 训练脚本默认采用BertWordPieceLowerCase分词器见 pretrain_bert.py 入口处的args_defaults{tokenizer_type: BertWordPieceLowerCase}对应的词表文件通常命名为bert-vocab.txt脚本注释中亦写为bert-vocab.json以实际下载的词表为准。预处理后的数据文件data path需传入形如xxx_text_document的前缀路径这是经过 Megatron 数据预处理工具如tools/preprocess_data.py生成二进制格式后的文件前缀脚本会据此找到xxx_text_document.idx与xxx_text_document.bin。检查点与日志目录训练过程中模型的保存checkpoint与 TensorBoard 日志输出目录。运行环境建议直接使用 NVIDIA PyTorch 官方容器如nvcr.io/nvidia/pytorch:24.01-py3它已预装 Megatron-LM 运行所需的 PyTorch、NCCL 及 Transformer Engine 等依赖。2. Docker 方式启动训练examples/bert/README.md给出了最直接的容器化启动方式。首先在宿主机上定义四个环境变量PYTORCH_IMAGEnvcr.io/nvidia/pytorch:24.01-py3 CHECKPOINT_PATH # 指定检查点保存路径 TENSORBOARD_LOGS_PATH # 指定 TensorBoard 日志路径 VOCAB_FILE # 指定词表文件路径如 /path/to/bert-vocab.txt DATA_PATH # 指定数据文件前缀如 /path/to/xxx_text_document随后运行容器把数据和代码目录挂载进容器并直接执行训练脚本docker run \ --gpusall \ --ipchost \ --workdir /workspace/megatron-lm \ -v /path/to/data:/path/to/data \ -v /path/to/megatron-lm:/workspace/megatron-lm \ megatron-lm nvcr.io/nvidia/pytorch:24.01-py3 \ bash examples/bert/train_bert_340m_distributed.sh $CHECKPOINT_PATH $TENSORBOARD_LOGS_PATH $VOCAB_FILE $DATA_PATH关键点说明--gpusall将宿主机全部 GPU 暴露给容器--ipchost使用宿主机的进程间通信命名空间这是多进程torchrun训练稳定运行的必要条件否则可能因共享内存不足报错。--workdir /workspace/megatron-lm将容器工作目录设为挂载进来的仓库根目录保证torchrun pretrain_bert.py的相对路径引用正确。训练脚本接收四个位置参数$1~$4分别对应检查点路径、TensorBoard 日志路径、词表文件路径和数据路径。原文档亦提示根据实际环境是否使用 SLURM、是否有自定义镜像等上述命令可能略有差异例如将--gpusall替换为--gpus device0,1,2,3或改为docker run --rm --shm-size...等。3. 340M 基线脚本逐段解析examples/bert/train_bert_340m_distributed.sh是本示例的核心它完整演示了如何用torchrun启动一个 340M 参数BERT-Large 规模的分布式预训练。下面按段拆解。3.1 分布式环境变量与启动器export CUDA_DEVICE_MAX_CONNECTIONS1 GPUS_PER_NODE8 MASTER_ADDRlocalhost MASTER_PORT29500 NUM_NODES1 NODE_RANK0 WORLD_SIZE$(($GPUS_PER_NODE*$NUM_NODES))CUDA_DEVICE_MAX_CONNECTIONS1限制每个 GPU 与 host 的最大连接数是 Megatron-LM 官方脚本的标准做法可避免部分场景下的通信死锁。单节点 8 卡GPUS_PER_NODE8NUM_NODES1多节点时需将MASTER_ADDR改为主节点 IP、NODE_RANK改为对应节点序号并相应扩大NUM_NODES。启动器参数通过 Bash 数组组织便于追加到torchrunDISTRIBUTED_ARGS( --nproc_per_node $GPUS_PER_NODE --nnodes $NUM_NODES --master_addr $MASTER_ADDR --master_port $MASTER_PORT )3.2 BERT 模型结构参数BERT_MODEL_ARGS( --num-layers 24 --hidden-size 1024 --num-attention-heads 16 --seq-length 512 --max-position-embeddings 512 --attention-backend auto # 可选 flash / fused / unfused / local )--num-layers 24、--hidden-size 1024、--num-attention-heads 16对应 BERT-Large 的标准结构24 层、1024 维、16 头模型参数约 340M故脚本注释称其为 340M parameter model。--seq-length 512与--max-position-embeddings 512输入序列长度与最大位置编码长度均为 512这是 BERT 预训练的经典设置如要训练更长序列两者需同步调整。--attention-backend auto注意力后端自动选择。从 pretrain_bert.py 的model_provider可以看出默认使用 Transformer Engine 版层结构bert_layer_with_transformer_engine_spec其注意力模块定义在 bert_layer_specs.py掩码类型为AttnMaskType.padding。如需切换到纯本地实现可在启动命令中追加--spec local对应bert_layer_local_spec。3.3 训练超参数TRAINING_ARGS( --micro-batch-size 4 --global-batch-size 32 --train-iters 1000000 --weight-decay 1e-2 --clip-grad 1.0 --fp16 --lr 0.0001 --lr-decay-iters 990000 --lr-decay-style linear --min-lr 1.0e-5 --lr-warmup-fraction .01 )--micro-batch-size 4单次前向/反向的微批次大小--global-batch-size 32全局批次大小所有数据并行组在一次迭代中消费的样本总数。二者与流水线/张量并行度共同决定梯度累积步数。--train-iters 1000000训练 100 万次迭代--lr-decay-iters 990000与--lr-warmup-fraction 0.01配合--lr-decay-style linear构成先 warmup 1% 迭代、再线性衰减到--min-lr 1.0e-5的经典 BERT 学习率调度。--fp16混合精度训练可显著降低显存占用并加速训练。3.4 模型并行配置MODEL_PARALLEL_ARGS( --tensor-model-parallel-size 8 --pipeline-model-parallel-size 16 )340M 脚本采用张量并行 8 流水线并行 16 的组合配合单节点 8 卡torchrun时实际会由框架在多节点间编排真实多节点运行需同步修改 3.1 节的NUM_NODES与NODE_RANK。张量并行将每个 Transformer 层的矩阵切分到 8 个 GPU流水线并行则把 24 层网络按阶段划分到 16 个流水线级二者正交组合是 Megatron-LM 支撑超大模型的核心手段。3.5 数据与评估日志参数DATA_ARGS( --data-path $DATA_PATH --vocab-file $VOCAB_FILE --split 949,50,1 ) EVAL_AND_LOGGING_ARGS( --log-interval 100 --save-interval 10000 --eval-interval 1000 --save $CHECKPOINT_PATH --load $CHECKPOINT_PATH --eval-iters 10 --tensorboard-dir $TENSORBOARD_LOGS_PATH )--split 949,50,1数据按 94.9% / 5% / 0.1% 划分为训练、验证、测试三个子集。--save-interval 10000每万次迭代保存一次检查点--eval-interval 1000每千次迭代跑一次验证每次验证跑--eval-iters 10个批次。训练从 pretrain_bert.py 的train_valid_test_datasets_provider可以看到数据构建的完整细节它使用BERTMaskedWordPieceDataset实现于 bert_dataset.py默认进行 3-gram 全词掩码masking_max_ngram3、masking_do_full_wordTrue并基于--split与--data-path构建混合数据集。4. 更大规模4B 与 20B 配置原文档指出该示例文件夹展示的是 340M 模型同时给出了升级到更大规模的参数组合可直接替换脚本中的BERT_MODEL_ARGS与MODEL_PARALLEL_ARGS两段。4.1 4B 配置单卡组可运行--num-layers 48 \ --hidden-size 2560 \ --num-attention-heads 32 \ --tensor-model-parallel-size 1 \ --pipeline-model-parallel-size 1 \48 层、2560 维、32 头约 40 亿参数并行度均设为 1意味着可先用单卡或纯数据并行跑通训练流程用于验证数据与代码链路。4.2 20B 配置大规模并行--num-layers 48 \ --hidden-size 6144 \ --num-attention-heads 96 \ --tensor-model-parallel-size 4 \ --pipeline-model-parallel-size 4 \48 层、6144 维、96 头约 200 亿参数采用张量并行 4 流水线并行 4共 16 个模型并行 GPU若开启数据并行还需乘以其维度是仓库文档中给出的典型大规模组合。注意切换配置时--seq-length、--micro-batch-size等训练参数也需根据显存容量相应调整必要时打开激活重计算--recompute-activations以降低显存峰值。5. 训练入口与数据流速览整个训练由 pretrain_bert.py 驱动其main入口先以BertWordPieceLowerCase为默认分词器解析参数随后调用megatron.training.pretrain进入通用训练循环。核心数据流如下model_provider根据--bert-no-binary-head默认开启二分类头即 BERT 的下一句预测/NSP 目标确定num_tokentypes构造 BertModelTransformer 层默认采用 Transformer Engine spec支持 FP8 训练。get_batch/forward_step从数据迭代器取出text / types / labels / is_random / loss_mask / padding_mask六个字段经张量并行广播后送入模型loss_func计算带loss_mask的掩码语言模型损失若启用二分类头则叠加 SOP句子顺序预测交叉熵损失最终通过average_losses_across_data_parallel_group得到跨数据并行组平均的lm loss与sop loss日志指标。从源码结构看pretrain_bert.py与 GPT 等模型共享同一套megatron.training.pretrain训练循环差异仅在于模型构建、数据集的 BERT 掩码逻辑与损失函数这正是 Megatron-LM多模型共用一套训练基础设施的设计体现。6. 常见问题与排查提示共享内存不足多进程数据加载时若报No space left on device请在 Docker 运行时显式增加共享内存如--shm-size1g。词表路径不匹配VOCAB_FILE指向的词表必须与BertWordPieceLowerCase分词器匹配否则 token 化结果错乱会导致损失异常。检查点/数据目录为空脚本要求CHECKPOINT_PATH、TENSORBOARD_LOGS_PATH、DATA_PATH四个位置参数一个都不能少且数据前缀必须指向已用tools/preprocess_data.py预处理好的_text_document文件。多节点扩展将 340M 脚本用于多节点时除修改NUM_NODES与NODE_RANK外还需保证所有节点能通过MASTER_ADDR:MASTER_PORT互通通常配合 SLURM/PBS 等调度器注入环境变量。按此流程你可以从 340M 基线跑通完整链路再依据第 4 节的配置逐步向 4B、20B 乃至更大规模扩展并结合 pretrain_bert.py 与 bert_layer_specs.py 深入定制自己的 BERT 变体。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
Megatron-LM BERT 分布式预训练实战：Docker 启动、340M 基线脚本与 4B/20B 扩展配置全解析
Megatron-LM BERT 分布式预训练实战Docker 启动、340M 基线脚本与 4B/20B 扩展配置全解析【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本指南聚焦 Megatron-LM 仓库中 BERT 预训练的完整实战链路如何用 Docker 容器启动训练、examples/bert/train_bert_340m_distributed.sh中 340M 参数基线脚本的每一项参数含义以及如何通过改写模型并行配置将其扩展到 4B 与 20B 规模。读完本文你将能独立配置数据路径、校验参数组合并把该示例脚本迁移到自己的多机多卡环境中运行。1. 前置准备数据、词表与运行环境BERT 预训练需要三类关键输入词汇表文件vocab fileMegatron-LM 的 BERT 训练脚本默认采用BertWordPieceLowerCase分词器见 pretrain_bert.py 入口处的args_defaults{tokenizer_type: BertWordPieceLowerCase}对应的词表文件通常命名为bert-vocab.txt脚本注释中亦写为bert-vocab.json以实际下载的词表为准。预处理后的数据文件data path需传入形如xxx_text_document的前缀路径这是经过 Megatron 数据预处理工具如tools/preprocess_data.py生成二进制格式后的文件前缀脚本会据此找到xxx_text_document.idx与xxx_text_document.bin。检查点与日志目录训练过程中模型的保存checkpoint与 TensorBoard 日志输出目录。运行环境建议直接使用 NVIDIA PyTorch 官方容器如nvcr.io/nvidia/pytorch:24.01-py3它已预装 Megatron-LM 运行所需的 PyTorch、NCCL 及 Transformer Engine 等依赖。2. Docker 方式启动训练examples/bert/README.md给出了最直接的容器化启动方式。首先在宿主机上定义四个环境变量PYTORCH_IMAGEnvcr.io/nvidia/pytorch:24.01-py3 CHECKPOINT_PATH # 指定检查点保存路径 TENSORBOARD_LOGS_PATH # 指定 TensorBoard 日志路径 VOCAB_FILE # 指定词表文件路径如 /path/to/bert-vocab.txt DATA_PATH # 指定数据文件前缀如 /path/to/xxx_text_document随后运行容器把数据和代码目录挂载进容器并直接执行训练脚本docker run \ --gpusall \ --ipchost \ --workdir /workspace/megatron-lm \ -v /path/to/data:/path/to/data \ -v /path/to/megatron-lm:/workspace/megatron-lm \ megatron-lm nvcr.io/nvidia/pytorch:24.01-py3 \ bash examples/bert/train_bert_340m_distributed.sh $CHECKPOINT_PATH $TENSORBOARD_LOGS_PATH $VOCAB_FILE $DATA_PATH关键点说明--gpusall将宿主机全部 GPU 暴露给容器--ipchost使用宿主机的进程间通信命名空间这是多进程torchrun训练稳定运行的必要条件否则可能因共享内存不足报错。--workdir /workspace/megatron-lm将容器工作目录设为挂载进来的仓库根目录保证torchrun pretrain_bert.py的相对路径引用正确。训练脚本接收四个位置参数$1~$4分别对应检查点路径、TensorBoard 日志路径、词表文件路径和数据路径。原文档亦提示根据实际环境是否使用 SLURM、是否有自定义镜像等上述命令可能略有差异例如将--gpusall替换为--gpus device0,1,2,3或改为docker run --rm --shm-size...等。3. 340M 基线脚本逐段解析examples/bert/train_bert_340m_distributed.sh是本示例的核心它完整演示了如何用torchrun启动一个 340M 参数BERT-Large 规模的分布式预训练。下面按段拆解。3.1 分布式环境变量与启动器export CUDA_DEVICE_MAX_CONNECTIONS1 GPUS_PER_NODE8 MASTER_ADDRlocalhost MASTER_PORT29500 NUM_NODES1 NODE_RANK0 WORLD_SIZE$(($GPUS_PER_NODE*$NUM_NODES))CUDA_DEVICE_MAX_CONNECTIONS1限制每个 GPU 与 host 的最大连接数是 Megatron-LM 官方脚本的标准做法可避免部分场景下的通信死锁。单节点 8 卡GPUS_PER_NODE8NUM_NODES1多节点时需将MASTER_ADDR改为主节点 IP、NODE_RANK改为对应节点序号并相应扩大NUM_NODES。启动器参数通过 Bash 数组组织便于追加到torchrunDISTRIBUTED_ARGS( --nproc_per_node $GPUS_PER_NODE --nnodes $NUM_NODES --master_addr $MASTER_ADDR --master_port $MASTER_PORT )3.2 BERT 模型结构参数BERT_MODEL_ARGS( --num-layers 24 --hidden-size 1024 --num-attention-heads 16 --seq-length 512 --max-position-embeddings 512 --attention-backend auto # 可选 flash / fused / unfused / local )--num-layers 24、--hidden-size 1024、--num-attention-heads 16对应 BERT-Large 的标准结构24 层、1024 维、16 头模型参数约 340M故脚本注释称其为 340M parameter model。--seq-length 512与--max-position-embeddings 512输入序列长度与最大位置编码长度均为 512这是 BERT 预训练的经典设置如要训练更长序列两者需同步调整。--attention-backend auto注意力后端自动选择。从 pretrain_bert.py 的model_provider可以看出默认使用 Transformer Engine 版层结构bert_layer_with_transformer_engine_spec其注意力模块定义在 bert_layer_specs.py掩码类型为AttnMaskType.padding。如需切换到纯本地实现可在启动命令中追加--spec local对应bert_layer_local_spec。3.3 训练超参数TRAINING_ARGS( --micro-batch-size 4 --global-batch-size 32 --train-iters 1000000 --weight-decay 1e-2 --clip-grad 1.0 --fp16 --lr 0.0001 --lr-decay-iters 990000 --lr-decay-style linear --min-lr 1.0e-5 --lr-warmup-fraction .01 )--micro-batch-size 4单次前向/反向的微批次大小--global-batch-size 32全局批次大小所有数据并行组在一次迭代中消费的样本总数。二者与流水线/张量并行度共同决定梯度累积步数。--train-iters 1000000训练 100 万次迭代--lr-decay-iters 990000与--lr-warmup-fraction 0.01配合--lr-decay-style linear构成先 warmup 1% 迭代、再线性衰减到--min-lr 1.0e-5的经典 BERT 学习率调度。--fp16混合精度训练可显著降低显存占用并加速训练。3.4 模型并行配置MODEL_PARALLEL_ARGS( --tensor-model-parallel-size 8 --pipeline-model-parallel-size 16 )340M 脚本采用张量并行 8 流水线并行 16 的组合配合单节点 8 卡torchrun时实际会由框架在多节点间编排真实多节点运行需同步修改 3.1 节的NUM_NODES与NODE_RANK。张量并行将每个 Transformer 层的矩阵切分到 8 个 GPU流水线并行则把 24 层网络按阶段划分到 16 个流水线级二者正交组合是 Megatron-LM 支撑超大模型的核心手段。3.5 数据与评估日志参数DATA_ARGS( --data-path $DATA_PATH --vocab-file $VOCAB_FILE --split 949,50,1 ) EVAL_AND_LOGGING_ARGS( --log-interval 100 --save-interval 10000 --eval-interval 1000 --save $CHECKPOINT_PATH --load $CHECKPOINT_PATH --eval-iters 10 --tensorboard-dir $TENSORBOARD_LOGS_PATH )--split 949,50,1数据按 94.9% / 5% / 0.1% 划分为训练、验证、测试三个子集。--save-interval 10000每万次迭代保存一次检查点--eval-interval 1000每千次迭代跑一次验证每次验证跑--eval-iters 10个批次。训练从 pretrain_bert.py 的train_valid_test_datasets_provider可以看到数据构建的完整细节它使用BERTMaskedWordPieceDataset实现于 bert_dataset.py默认进行 3-gram 全词掩码masking_max_ngram3、masking_do_full_wordTrue并基于--split与--data-path构建混合数据集。4. 更大规模4B 与 20B 配置原文档指出该示例文件夹展示的是 340M 模型同时给出了升级到更大规模的参数组合可直接替换脚本中的BERT_MODEL_ARGS与MODEL_PARALLEL_ARGS两段。4.1 4B 配置单卡组可运行--num-layers 48 \ --hidden-size 2560 \ --num-attention-heads 32 \ --tensor-model-parallel-size 1 \ --pipeline-model-parallel-size 1 \48 层、2560 维、32 头约 40 亿参数并行度均设为 1意味着可先用单卡或纯数据并行跑通训练流程用于验证数据与代码链路。4.2 20B 配置大规模并行--num-layers 48 \ --hidden-size 6144 \ --num-attention-heads 96 \ --tensor-model-parallel-size 4 \ --pipeline-model-parallel-size 4 \48 层、6144 维、96 头约 200 亿参数采用张量并行 4 流水线并行 4共 16 个模型并行 GPU若开启数据并行还需乘以其维度是仓库文档中给出的典型大规模组合。注意切换配置时--seq-length、--micro-batch-size等训练参数也需根据显存容量相应调整必要时打开激活重计算--recompute-activations以降低显存峰值。5. 训练入口与数据流速览整个训练由 pretrain_bert.py 驱动其main入口先以BertWordPieceLowerCase为默认分词器解析参数随后调用megatron.training.pretrain进入通用训练循环。核心数据流如下model_provider根据--bert-no-binary-head默认开启二分类头即 BERT 的下一句预测/NSP 目标确定num_tokentypes构造 BertModelTransformer 层默认采用 Transformer Engine spec支持 FP8 训练。get_batch/forward_step从数据迭代器取出text / types / labels / is_random / loss_mask / padding_mask六个字段经张量并行广播后送入模型loss_func计算带loss_mask的掩码语言模型损失若启用二分类头则叠加 SOP句子顺序预测交叉熵损失最终通过average_losses_across_data_parallel_group得到跨数据并行组平均的lm loss与sop loss日志指标。从源码结构看pretrain_bert.py与 GPT 等模型共享同一套megatron.training.pretrain训练循环差异仅在于模型构建、数据集的 BERT 掩码逻辑与损失函数这正是 Megatron-LM多模型共用一套训练基础设施的设计体现。6. 常见问题与排查提示共享内存不足多进程数据加载时若报No space left on device请在 Docker 运行时显式增加共享内存如--shm-size1g。词表路径不匹配VOCAB_FILE指向的词表必须与BertWordPieceLowerCase分词器匹配否则 token 化结果错乱会导致损失异常。检查点/数据目录为空脚本要求CHECKPOINT_PATH、TENSORBOARD_LOGS_PATH、DATA_PATH四个位置参数一个都不能少且数据前缀必须指向已用tools/preprocess_data.py预处理好的_text_document文件。多节点扩展将 340M 脚本用于多节点时除修改NUM_NODES与NODE_RANK外还需保证所有节点能通过MASTER_ADDR:MASTER_PORT互通通常配合 SLURM/PBS 等调度器注入环境变量。按此流程你可以从 340M 基线跑通完整链路再依据第 4 节的配置逐步向 4B、20B 乃至更大规模扩展并结合 pretrain_bert.py 与 bert_layer_specs.py 深入定制自己的 BERT 变体。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考