Megatron-LM 首次训练实战指南：从最小分布式循环到 LLaMA-3 FP8 训练与数据预处理

Megatron-LM 首次训练实战指南从最小分布式循环到 LLaMA-3 FP8 训练与数据预处理【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本文是 Megatron-LM 的首次训练运行技术指南覆盖从环境验证、最小分布式训练循环到 8 卡 FP8 精度 LLaMA-3 8B 生产级训练再到自有数据集 JSONL 预处理的全流程。读者学完后将掌握torchrun多卡启动方式、Megatron-Core 训练循环的关键组件并行状态、DDP、前向反向调度、分布式检查点、LLaMA-3 训练脚本中的模型/训练/并行/数据参数含义以及preprocess_data.py生成.bin/.idx二进制数据的具体操作方法。前置条件请先完成 安装指南PyPI 安装、源码安装或 NGC 容器任一方式均可再继续本文。本文所有命令默认在 Megatron-LM 仓库根目录执行。最小训练示例验证分布式环境本节目标用最简单的方式跑通一个分布式训练循环确认环境配置正确为后续真实模型训练铺路。启动命令torchrun --nproc_per_node2 examples/run_simple_mcore_train_loop.pytorchrun会在单机 2 张 GPU 上各启动一个进程--nproc_per_node2指定每节点进程数即 GPU 数。该脚本不依赖任何真实数据使用 mock 数据即可完成完整的初始化—前向—反向—参数更新—检查点保存/加载链路。脚本做了什么源码解析examples/run_simple_mcore_train_loop.py 是一个自包含的最小训练循环示例主要步骤如下初始化分布式环境initialize_distributed读取RANK/WORLD_SIZE/LOCAL_RANK环境变量调用torch.distributed.init_process_group(backendnccl, ...)建立 NCCL 通信组随后调用parallel_state.initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size)初始化 Megatron-Core 的张量并行/流水线并行通信组。脚本默认tensor_model_parallel_size2、pipeline_model_parallel_size1即 2 卡上各承担一半张量切分。构造最小 GPT 模型model_provider通过TransformerConfig配置num_layers2、hidden_size12、num_attention_heads4以use_cpu_initializationTrue在 CPU 上初始化权重pipeline_dtypetorch.float32再用get_gpt_layer_local_spec()非 Transformer Engine 的本地实现层规格构建GPTModelvocab_size100、max_sequence_length64。这是一个刻意缩小、便于快速验证的玩具模型。mock 训练数据get_train_data_iterator通过BlendedMegatronDatasetBuilder与MockGPTDataset实现见 megatron/core/datasets/gpt_dataset.py#L933构造 1000 条 mock 样本GPTDatasetConfig中sequence_length64并用MegatronTokenizer.from_pretrained构建一个虚拟 tokenizercompile_helpers()负责在多进程下编译数据加载所需的 C 辅助算子仅在 rank 0 编译后通过 barrier 同步。训练用DataLoader的batch_size8。前向反向调度get_forward_backward_func()返回 Megatron 的前向/反向调度函数forward_backward_func(forward_step_func..., data_iterator..., model..., num_microbatches1, seq_length64, micro_batch_size8, forward_onlyFalse)完成一轮前向反向。forward_step_func中定义了loss_func损失计算为按loss_mask加权的 LM 损失均值。梯度同步与优化器模型被DistributedDataParallel包装grad_reduce_in_fp32False、overlap_grad_reduceFalse、use_distributed_optimizerFalse每轮迭代调用finalize_model_grads([gpt_model])完成 DP 组间的梯度 all-reduce 与 TP 组内非张量并行参数如 LayerNorm的梯度同步随后optim.step()。默认训练 5 个 iteration。分布式检查点脚本演示了dist_checkpointing.save保存 sharded state dict 到./ckpt与dist_checkpointing.load加载并load_state_dict回填的完整闭环并打印Successfully loaded the model确认成功。验证要点若启动后看到Iteration 0..4: Losses reduced: {...}日志并最终输出Successfully loaded the model说明环境NCCL、CUDA、Megatron-Core 安装全部就绪。单卡环境下可将--nproc_per_node改为 1 快速冒烟但脚本内部默认初始化tensor_model_parallel_size2单卡运行需自行调整为initialize_distributed(tensor_model_parallel_size1, ...)。LLaMA-3 训练示例8 卡 FP8 生产级训练环境验证通过后即可运行生产级示例。以下脚本在 8 张 Hopper/Ada/Blackwell GPUFP8 需要这些架构见 安装指南 的系统要求上使用 mock 数据训练 LLaMA-3 8B 规模模型演示张量并行与优化 kernel 的配合。./examples/llama/train_llama3_8b_h100_fp8.sh注意脚本默认DTYPEfp8需要支持 FP8 的 GPU 架构若不支持可将DTYPE改为bf16后使用。脚本参数逐项解读examples/llama/train_llama3_8b_h100_fp8.sh 是自包含的 bash 脚本分为几大参数块环境变量与路径默认CHECKPOINT_PATHcheckpoints/llama3_8b_fp8、TENSORBOARD_LOGS_PATHtensorboard_logs/llama3_8b_fp8、TOKENIZER_ARGMOCK、DATA_ARGMOCK均可作为位置参数覆盖。CUDA_DEVICE_MAX_CONNECTIONS1用于提升 kernel 并发脚本注释中还预留了NCCL_IB_TIMEOUT、NVTE_*_LAYERNORM_SM_MARGIN等 NCCL/Transformer Engine 调优项。分布式设置GPUS_PER_NODE8、NUM_NODES1、WORLD_SIZE8通过MASTER_ADDR默认localhost、MASTER_PORT默认29500、NODE_RANK默认 0组装DISTRIBUTED_ARGS传给torchrun多机扩展时修改这些变量即可。模型参数LLaMA-3 8B 结构参数值说明--use-mcore-models—使用 Megatron-Core 模型实现--num-layers3232 层 decoder--hidden-size4096隐藏维度--ffn-hidden-size14336FFN 中间维度SwiGLU--num-attention-heads32注意力头数--group-query-attention/--num-query-groups8GQAKV 头 8 个--kv-channels128每 KV 头通道数--seq-length/--max-position-embeddings8192序列长度与最大位置--position-embedding-type rope/--rotary-base1000000RoPE 位置编码base 为 1e6--swiglu—SwiGLU 激活--normalization RMSNorm—归一化层--attention-backend fused—融合 attention kernel--untie-embeddings-and-output-weights—输入/输出 embedding 不共享权重训练参数--micro-batch-size 1、--global-batch-size 128数据并行下梯度累积micro×DP×accumulation global--train-samples/--lr-decay-samples/--lr-warmup-samples定义 cosine 调度曲线--lr-decay-style cosine--lr 0.00015、--min-lr 0.00001还包含 decoupled AdamW 专用的--decoupled-lr/--decoupled-min-lr--clip-grad 1.0、--weight-decay 0.1、--adam-beta1 0.9、--adam-beta2 0.95为常规优化器设置精度相关--bf16、--grad-reduce-in-bf16、--cross-entropy-loss-fusion、--calculate-per-token-loss性能相关--manual-gc、--empty-unused-memory-level 1--exit-duration-in-mins 235表示训练满 235 分钟自动退出benchmark 场景。FP8 参数当DTYPEfp8时追加--fp8-format hybrid混合 E4M3/E5M2 格式、--fp8-amax-history-len 1024amax 历史窗口、--fp8-amax-compute-algo max、--fp8-param-gather。这些参数驱动 Transformer Engine 的 FP8 自动缩放训练路径。并行与 DDP 参数--tensor-model-parallel-size 1、--context-parallel-size 1、--sequence-parallelDDP 侧启用--use-distributed-optimizer分片优化器、--overlap-grad-reduce、--overlap-param-gather梯度归约/参数收集与计算重叠。数据参数mock vs 真实数据切换当TOKENIZER_ARG或DATA_ARG为MOCK时走 mock 分支--mock-data、--tokenizer-type NullTokenizer、--vocab-size 128256LLaMA-3 词表大小、--data-cache-path benchmark_cache_llama3_8b_fp8、--tiktoken-pattern v2、--split 99,1,0、--no-create-attention-mask-in-dataloader、--no-mmap-bin-files、--num-workers 1。若传入真实 tokenizer 路径与数据前缀则走真实数据分支--data-path $DATA_ARG、--tokenizer-type HuggingFaceTokenizer、--tokenizer-model $TOKENIZER_ARG。NullTokenizer实现见 megatron/core/tokenizers/text/libraries/null_tokenizer.pyHuggingFaceTokenizer实现见 megatron/core/tokenizers/text/libraries/huggingface_tokenizer.py。评估与日志--log-interval 1、--eval-iters 32、--eval-interval 100、--save-interval 1000、--log-throughput、--profile--profile-step-start 4/--profile-step-end 6指定 profiling 步区间、--ckpt-format torch_dist分布式检查点格式、--distributed-timeout-minutes 60、--save/--load指向同一 checkpoint 目录支持断点续训、--tensorboard-dir。入口校验脚本启动前检查pretrain_gpt.py是否存在于仓库根目录该文件位于 pretrain_gpt.py不存在则报错并提示在仓库根目录运行。运行产物checkpoint 目录默认checkpoints/llama3_8b_fp8/torch_dist格式的分布式检查点可被--load复用续训TensorBoard 日志默认tensorboard_logs/llama3_8b_fp8/loss、吞吐等指标benchmark_cache_llama3_8b_fp8/mock/真实数据构建过程中的数据缓存目录。数据准备把自有数据转成 Megatron 二进制格式训练自有数据前Megatron 期望输入预处理后的二进制文件.bin与.idx配对.bin存放 token id 序列IndexedDataset以最优 dtype 存储见 tools/preprocess_data.py 中DType.optimal_dtype(tokenizer.vocab_size)的自动选择逻辑.idx是文档/样本的索引文件供训练时按文档或样本快速定位。第 1 步准备 JSONL 文件每一行是一个 JSON 对象必须包含text字段对应--json-keys的默认值text{text: Your training text here...} {text: Another training sample...}如需自定义字段名可用--json-keys指定如--json-keys text title预处理脚本会为每个 key 分别生成对应的.bin/.idx文件。第 2 步运行预处理脚本python tools/preprocess_data.py \ --input data.jsonl \ --output-prefix processed_data \ --tokenizer-type HuggingFaceTokenizer \ --tokenizer-model /path/to/tokenizer.model \ --workers 8 \ --append-eod执行后会在--output-prefix基础上生成processed_data_text_document.bin与processed_data_text_document.idx未开启--split-sentences时为 document 级别开启后为_sentence级别。训练时通过--data-path指向该前缀不含.bin/.idx后缀即可。核心参数说明参数作用补充说明源码依据--input输入 JSON/JSONL 文件路径必填支持--partitions分片输入--output-prefix输出二进制文件前缀必填实际生成{prefix}_{key}_{level}.bin/.idx--tokenizer-type分词器类型如HuggingFaceTokenizer、GPT2BPETokenizer、NullTokenizer等--tokenizer-model分词器模型文件路径对 HuggingFace 类分词器通常为 tokenizer 路径--workers并行处理进程数源码提示良好默认值为workers × partitions 可用 CPU 核数--append-eod在文档末尾追加eodtoken源码中doc_ids.append(Encoder.tokenizer.eod)开启后每个文档以 EOD 结尾--json-keys从 JSON 中抽取的字段默认[text]可传多个--split-sentences是否按句子切分依赖 NLTKBERT 类任务建议开启GPT 类任务通常关闭--find-optimal-num-workers自动探测最优 worker 数会依次用--workers-to-check默认 16/32/64跑小作业并输出性能排序--partitions文件分区数配合多文件 glob 输入与--keep-sequential-samples使用底层处理流程源码级preprocess_data.py 的处理管线为分词Encoder.initializer/encode每个 worker 进程通过build_tokenizer(args)构建分词器对每行 JSON 按--json-keys抽取字段逐句tokenize后拼接文档 token 序列若开启--append-eod则在文档末尾追加 EOD token代码位置preprocess_data.py#L103-L105。写入二进制IndexedDatasetBuilder.add_document以DType.optimal_dtype(vocab_size)选择存储 dtype逐文档写入.bin。生成索引builders[key].finalize(output_idx_files[key])写入.idx索引文件记录每个文档/样本的偏移与长度。多进程并行通过multiprocessing.Pool(workers, initializerencoder.initializer)并行处理pool.imap(encoder.encode, fin, 32)以 32 行为一批流式消费输入避免一次性加载大文件。进度统计print_processing_stats按--log-interval默认 1000输出已处理文档数与吞吐docs/s、MB/s。整个流程支持--find-optimal-num-workers自动调优脚本会对候选 worker 数各跑一轮小规模处理受--max-documents限制默认 100000 条最终打印按吞吐排序的最优 worker 数。训练时引用已处理数据以 examples/llama/train_llama3_8b_h100_fp8.sh 的真实数据分支为例将DATA_ARG指向预处理输出的前缀、TOKENIZER_ARG指向 tokenizer 路径后重新运行即可接入自有数据。更通用的数据加载说明见 数据准备最佳实践。下一步扩展训练规模阅读 并行策略指南了解张量并行、流水线并行、上下文并行与序列并行的组合方式深入学习数据加载阅读 数据准备掌握数据集混合、采样与缓存细节进阶特性浏览 高级特性总览其中涵盖 FP8 训练、上下文并行、分布式优化器、多 token 预测等能力。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考