企业官网建设流程全解析

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_node=2 examples/run_simple_mcore_train_loop.py

torchrun会在单机 2 张 GPU 上各启动一个进程，--nproc_per_node=2指定每节点进程数（即 GPU 数）。该脚本不依赖任何真实数据，使用 mock 数据即可完成完整的初始化—前向—反向—参数更新—检查点保存/加载链路。

脚本做了什么（源码解析）

examples/run_simple_mcore_train_loop.py 是一个自包含的最小训练循环示例，主要步骤如下：

1. 初始化分布式环境（initialize_distributed）：读取RANK/WORLD_SIZE/LOCAL_RANK环境变量，调用torch.distributed.init_process_group(backend="nccl", ...)建立 NCCL 通信组；随后调用parallel_state.initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size)初始化 Megatron-Core 的张量并行/流水线并行通信组。脚本默认tensor_model_parallel_size=2、pipeline_model_parallel_size=1，即 2 卡上各承担一半张量切分。

2. 构造最小 GPT 模型（model_provider）：通过TransformerConfig配置num_layers=2、hidden_size=12、num_attention_heads=4，以use_cpu_initialization=True在 CPU 上初始化权重，pipeline_dtype=torch.float32；再用get_gpt_layer_local_spec()（非 Transformer Engine 的本地实现层规格）构建GPTModel，vocab_size=100、max_sequence_length=64。这是一个刻意缩小、便于快速验证的"玩具模型"。

3. mock 训练数据（get_train_data_iterator）：通过BlendedMegatronDatasetBuilder与MockGPTDataset（实现见 megatron/core/datasets/gpt_dataset.py#L933）构造 1000 条 mock 样本，GPTDatasetConfig中sequence_length=64，并用MegatronTokenizer.from_pretrained构建一个虚拟 tokenizer；compile_helpers()负责在多进程下编译数据加载所需的 C++ 辅助算子（仅在 rank 0 编译后通过 barrier 同步）。训练用DataLoader的batch_size=8。

4. 前向反向调度：get_forward_backward_func()返回 Megatron 的前向/反向调度函数，forward_backward_func(forward_step_func=..., data_iterator=..., model=..., num_microbatches=1, seq_length=64, micro_batch_size=8, forward_only=False)完成一轮前向+反向。forward_step_func中定义了loss_func，损失计算为"按loss_mask加权的 LM 损失均值"。

5. 梯度同步与优化器：模型被DistributedDataParallel包装（grad_reduce_in_fp32=False、overlap_grad_reduce=False、use_distributed_optimizer=False），每轮迭代调用finalize_model_grads([gpt_model])完成 DP 组间的梯度 all-reduce 与 TP 组内非张量并行参数（如 LayerNorm）的梯度同步，随后optim.step()。默认训练 5 个 iteration。

6. 分布式检查点：脚本演示了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_size=2，单卡运行需自行调整为initialize_distributed(tensor_model_parallel_size=1, ...)。

LLaMA-3 训练示例：8 卡 FP8 生产级训练

环境验证通过后，即可运行生产级示例。以下脚本在 8 张 Hopper/Ada/Blackwell GPU（FP8 需要这些架构，见 安装指南 的系统要求）上，使用 mock 数据训练 LLaMA-3 8B 规模模型，演示张量并行与优化 kernel 的配合。

./examples/llama/train_llama3_8b_h100_fp8.sh

注意：脚本默认DTYPE="fp8"，需要支持 FP8 的 GPU 架构；若不支持，可将DTYPE改为bf16后使用。

脚本参数逐项解读

examples/llama/train_llama3_8b_h100_fp8.sh 是自包含的 bash 脚本，分为几大参数块：

环境变量与路径：默认CHECKPOINT_PATH="checkpoints/llama3_8b_fp8"、TENSORBOARD_LOGS_PATH="tensorboard_logs/llama3_8b_fp8"、TOKENIZER_ARG="MOCK"、DATA_ARG="MOCK"，均可作为位置参数覆盖。CUDA_DEVICE_MAX_CONNECTIONS=1用于提升 kernel 并发；脚本注释中还预留了NCCL_IB_TIMEOUT、NVTE_*_LAYERNORM_SM_MARGIN等 NCCL/Transformer Engine 调优项。

分布式设置：GPUS_PER_NODE=8、NUM_NODES=1、WORLD_SIZE=8，通过MASTER_ADDR（默认localhost）、MASTER_PORT（默认29500）、NODE_RANK（默认 0）组装DISTRIBUTED_ARGS传给torchrun，多机扩展时修改这些变量即可。

模型参数（LLaMA-3 8B 结构）：

训练参数：--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 参数：当DTYPE=fp8时追加--fp8-format hybrid（混合 E4M3/E5M2 格式）、--fp8-amax-history-len 1024（amax 历史窗口）、--fp8-amax-compute-algo max、--fp8-param-gather。这些参数驱动 Transformer Engine 的 FP8 自动缩放训练路径。

并行与 DDP 参数：--tensor-model-parallel-size 1、--context-parallel-size 1、--sequence-parallel；DDP 侧启用--use-distributed-optimizer（分片优化器）、--overlap-grad-reduce、--overlap-param-gather（梯度归约/参数收集与计算重叠）。

数据参数（mock vs 真实数据切换）：当TOKENIZER_ARG或DATA_ARG为MOCK时走 mock 分支：--mock-data、--tokenizer-type NullTokenizer、--vocab-size 128256（LLaMA-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.py，HuggingFaceTokenizer实现见 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后缀）即可。

核心参数说明

底层处理流程（源码级）

preprocess_data.py 的处理管线为：

1. 分词（Encoder.initializer/encode）：每个 worker 进程通过build_tokenizer(args)构建分词器；对每行 JSON 按--json-keys抽取字段，逐句tokenize后拼接文档 token 序列；若开启--append-eod则在文档末尾追加 EOD token（代码位置：preprocess_data.py#L103-L105）。
2. 写入二进制（IndexedDatasetBuilder.add_document）：以DType.optimal_dtype(vocab_size)选择存储 dtype，逐文档写入.bin。
3. 生成索引（builders[key].finalize(output_idx_files[key])）：写入.idx索引文件，记录每个文档/样本的偏移与长度。
4. 多进程并行：通过multiprocessing.Pool(workers, initializer=encoder.initializer)并行处理，pool.imap(encoder.encode, fin, 32)以 32 行为一批流式消费输入，避免一次性加载大文件。
5. 进度统计：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