Megatron-LM 中基于 FP8 精度的 Llama-3-8B 分布式训练实战指南

Megatron-LM 中基于 FP8 精度的 Llama-3-8B 分布式训练实战指南

【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM

本篇技术指南以 examples/llama/README.md 为核心脉络，系统讲解如何基于 Megatron-LM（Megatron-Core）在 NVIDIA Hopper/Ada/Blackwell 架构 GPU 上以 FP8 混合精度训练 Llama-3-8B 模型，覆盖环境准备、Mock 数据与真实数据两种 Docker 启动方式、完整训练参数逐项解析、FP8 性能参考、测试数据集预处理，以及 FP8 训练不稳定（NaN）时的排查思路。读完本文，你将能直接复现train_llama3_8b_h100_fp8.sh脚本的全流程训练，并理解--fp8-format hybrid、--fp8-amax-history-len、--fp8-param-gather等关键参数在 Megatron-Core 底层的作用机制。

1. 概述

Megatron-LM 的 Llama 示例旨在展示：使用 FP8 精度配合 Megatron-Core 训练 Llama 系列模型。FP8（8 位浮点）通过 Transformer Engine（TE）的自动混合精度机制，将 GEMM 的输入张量量化为 8 位参与计算，从而在保持接近 BF16 训练精度的前提下显著提升吞吐、降低显存占用。

该示例的核心载体是 train_llama3_8b_h100_fp8.sh，它基于仓库根目录的 pretrain_gpt.py 入口启动一次完整的多卡训练。从源码结构看，该脚本是一个"性能基准 + 可运行示例"二合一的训练配方：既可以通过MOCK数据快速跑通全流程验证吞吐，也可以挂载自定义数据集与 HuggingFace tokenizer 进行真实训练。

2. 前置条件与环境准备

官方文档建议按以下步骤准备环境：

# Clone repository export HOST_MEGATRON_LM_DIR="/path/to/your/host/megatron-lm" git clone https://github.com/NVIDIA/Megatron-LM.git "$HOST_MEGATRON_LM_DIR" cd "$HOST_MEGATRON_LM_DIR" git checkout "core_r0.12.0" # Set paths export HOST_CHECKPOINT_PATH="./checkpoints/llama3_8b_fp8" export HOST_TENSORBOARD_LOGS_PATH="./tensorboard_logs/llama3_8b_fp8" # Optional: For real data # export HOST_TOKENIZER_MODEL_PATH="/path/to/host/tokenizer.model" # export HOST_DATA_PREFIX="/path/to/host/mydata_prefix"

需要特别说明的环境约束：

• 运行环境：示例以nvcr.io/nvidia/pytorch:25.03-py3NGC 容器为基准。该容器预装了 Megatron-Core 运行所需的 PyTorch、Transformer Engine、Apex 等依赖，可直接运行本仓库的训练入口。
• 硬件要求：FP8 训练需要NVIDIA Hopper（H100 等）、Ada 或 Blackwell 架构 GPU，这是 FP8 算力的硬件前提（详见第 7 节）。
• 数据路径约定：HOST_DATA_PREFIX是数据前缀（不含扩展名），例如某份预处理产物的前缀为/data/mydata，则实际文件为mydata.idx与mydata.bin。Docker 挂载时采用$(dirname "${HOST_DATA_PREFIX}"):/workspace/data_dir的写法，即挂载数据所在目录，再在容器内拼接"/workspace/data_dir/$(basename "${HOST_DATA_PREFIX}")"得到完整前缀。

3. 训练启动方式

示例提供了两种启动方式：使用 Mock 数据快速验证，以及使用自定义数据与 tokenizer 进行真实训练。两者都通过docker run挂载仓库、checkpoint、tensorboard 目录，并在容器内执行同一个训练脚本。

3.1 使用 Mock 数据

PYTORCH_IMAGE="nvcr.io/nvidia/pytorch:25.03-py3" docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \ -v "${HOST_MEGATRON_LM_DIR}:/workspace/megatron-lm" \ -v "${HOST_CHECKPOINT_PATH}:/workspace/checkpoints" \ -v "${HOST_TENSORBOARD_LOGS_PATH}:/workspace/tensorboard_logs" \ --workdir /workspace/megatron-lm \ $PYTORCH_IMAGE \ bash examples/llama/train_llama3_8b_h100_fp8.sh \ /workspace/checkpoints \ /workspace/tensorboard_logs \ 2>&1 | tee "${HOST_TENSORBOARD_LOGS_PATH}/training_mock_$(date +'%y-%m-%d_%H-%M-%S').log"

注意脚本只传了CHECKPOINT_PATH与TENSORBOARD_LOGS_PATH两个位置参数，TOKENIZER_ARG与DATA_ARG走脚本默认值MOCK，从而触发 Mock 数据分支（见第 4 节参数说明）。tee会把容器输出同时落盘到主机日志文件，便于回溯。

3.2 使用自定义数据与 Tokenizer

PYTORCH_IMAGE="nvcr.io/nvidia/pytorch:25.03-py3" docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \ -v "${HOST_MEGATRON_LM_DIR}:/workspace/megatron-lm" \ -v "${HOST_CHECKPOINT_PATH}:/workspace/checkpoints" \ -v "${HOST_TENSORBOARD_LOGS_PATH}:/workspace/tensorboard_logs" \ -v "${HOST_TOKENIZER_MODEL_PATH}:/workspace/tokenizer_model" \ -v "$(dirname "${HOST_DATA_PREFIX}"):/workspace/data_dir" \ --workdir /workspace/megatron-lm \ $PYTORCH_IMAGE \ bash examples/llama/train_llama3_8b_h100_fp8.sh \ /workspace/checkpoints \ /workspace/tensorboard_logs \ /workspace/tokenizer_model \ "/workspace/data_dir/$(basename "${HOST_DATA_PREFIX}")" \ 2>&1 | tee "${HOST_TENSORBOARD_LOGS_PATH}/training_custom_$(date +'%y-%m-%d_%H-%M-%S').log"

与 Mock 模式相比，这里额外传入了第 3、4 个位置参数（tokenizer 模型路径与数据前缀），并多挂载了 tokenizer 模型文件和真实数据目录。脚本检测到这两个参数非MOCK后会自动切换到真实数据分支。

4. 配置详解

本节逐一拆解训练脚本中的默认配置。文档明确说明：这些参数可直接在 train_llama3_8b_h100_fp8.sh 中修改，其配置基线参考了 NeMo Framework 性能脚本（NeMo 底层同样使用 Megatron-Core）。

4.1 默认并行策略

对应脚本中的TP_SIZE=1、CP_SIZE=1、PP_SIZE=1（注：脚本实际默认CP_SIZE=1，README 中给出的是该配方文档层面的默认策略，具体以脚本为准）。脚本同时始终开启--sequence-parallel，并注释说明"TP_SIZE=2 时总是启用序列并行"。对于 8 卡单节点的 Llama-3-8B，数据并行 + 序列并行即可承载模型，无需张量/流水并行切分。

4.2 Llama-3-8B 模型架构

--num-query-groups 8配合--group-query-attention实现了 Llama-3 的GQA（Grouped Query Attention）结构，--kv-channels 128对应每组的 KV 头维度。其余架构细节包括：--attention-dropout 0.0/--hidden-dropout 0.0（训练不启用 dropout）、--attention-backend fused（使用融合注意力内核）、--apply-layernorm-1p（在 QKV 投影前额外应用 LayerNorm，对应 Llama-3 的真实实现）、--untie-embeddings-and-output-weights（embedding 与 LM head 权重不共享）、--disable-bias-linear（线性层去 bias）、--init-method-std 0.0134（权重初始化标准差）。

4.3 关键训练参数

脚本同时开启了--bf16（优化器主权重与梯度仍以 BF16/FP32 形式维护，FP8 仅作用于计算路径）、--grad-reduce-in-bf16、--cross-entropy-loss-fusion、--calculate-per-token-loss、--manual-gc --empty-unused-memory-level 1（显存管理）等参数。

值得一提的细节：脚本为 decoupled AdamW 单独设置了--decoupled-lr 5.0e-4与--decoupled-min-lr 4.5e-5，并注明"需确保优化器兼容"。这是 Llama-3 官方配方中权重（embedding 等）使用更高学习率、与主体参数解耦的策略。

4.4 FP8 相关参数（DTYPE_ARGS）

当DTYPE == "fp8"时，脚本追加以下参数：

"--fp8-format hybrid" # FP8 格式：hybrid "--fp8-amax-history-len 1024" # amax 历史窗口长度 "--fp8-amax-compute-algo max" # 取历史窗口内最大 amax "--fp8-param-gather" # 以 FP8 执行参数 all-gather

对照源码可确认每个参数的语义：

• --fp8-format：对应 transformer_config.py 中的fp8字段。hybrid表示激活与权重张量统一使用 e4m3、输出激活梯度使用 e5m2（e5m2 的动态范围更大，适合梯度这类分布剧烈的张量），而e4m3则全部使用 e4m3。从 enums.py 看，fp8_recipe还支持delayed（延迟缩放）、tensorwise、mxfp8（仅 Blackwell）、blockwise等配方，默认delayed。
• --fp8-amax-history-len与--fp8-amax-compute-algo：分别对应 transformer_config.py 的fp8_amax_history_len与fp8_amax_compute_algo字段，二者最终被透传给 Transformer Engine（见 transformer_engine.py 中amax_history_len、amax_compute_algo的构造调用）。延迟缩放配方基于历史窗口内的amax估计缩放因子，max策略取窗口内最大值以更保守地抑制溢出，1024 的长窗口可显著提升数值稳定性。
• --fp8-param-gather：定义于 arguments.py，其说明为"将计算参数保持为 FP8，并以 FP8 执行参数 all-gather"。常规做法下主权重先由 FP32 转成 BF16/FP16 再参与 all-gather 与 GEMM；开启后权重直接在 FP8 下收集，据 MoE 文档（moe/README.md）可将参数通信开销降低约 50%。fp8_utils.py 中封装了该特性的核心逻辑。需注意 arguments.py 的约束：--fp8-param-gather仅支持分布式优化器、FSDP 或推理模式。

4.5 分布式优化器与通信重叠

脚本的DDP_ARGS启用：

• --use-distributed-optimizer：分布式优化器，将优化器状态按数据并行维度分片（对应项目 dist_optimizer 特性）。
• --overlap-grad-reduce/--overlap-param-gather：梯度归约与参数收集分别与反向、前向计算重叠，隐藏通信延迟。

这与--fp8-param-gather协同工作：开启后参数 all-gather 走 FP8 路径（见 distrib_optimizer.py 关于启用/未启用--fp8-param-gather时权重存储方式的差异说明）。

4.6 数据与日志相关参数

Mock 分支：

--mock-data # 跳过真实数据加载，合成数据 --tokenizer-type NullTokenizer --vocab-size 128256 # Llama-3 词表大小 --data-cache-path ${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

真实数据分支：

--data-path $DATA_ARG --tokenizer-type HuggingFaceTokenizer --tokenizer-model $TOKENIZER_ARG --data-cache-path ${DATA_CACHE_PATH} --split '99,1,0' --no-create-attention-mask-in-dataloader --no-mmap-bin-files --num-workers 1 --vocab-size 128256 # 脚本注释：HuggingFaceTokenizer 可能自动推断，也可能需要显式指定

NullTokenizer是项目内置的合成 tokenizer（null_tokenizer.py），用于"性能基准与调试"场景，其eod_id默认取vocab_size - 1、pad_id默认 -1（无 pad token）。--mock-data在 arguments.py 中的定义为跳过数据加载与校验、改用人工合成的 Mock 数据。日志与评估部分还包含：--log-interval 1（每步打印）、--eval-iters 32、--eval-interval 100、--save-interval 1000、--log-throughput、--profile --profile-step-start 4 --profile-step-end 6（在 4~6 步间做 profiler 采样）、--ckpt-format torch_dist、--distributed-timeout-minutes 60、--save/--load与--tensorboard-dir。

5. FP8 性能参考

README 给出了该配方在官方基准环境下的吞吐数据（模型/GPU 数/并行配置均为完整事实，摘录如下）：

图例（Legend）：

• GBS：Global Batch Size（全局批大小）
• MBS：Micro Batch Size（微批大小）
• TP：Tensor Parallel size（张量并行）
• PP：Pipeline Parallel size（流水并行）
• CP：Context Parallel size（上下文并行）
• VP：Virtual Pipeline stages（虚拟流水阶段数）
• EP：Expert Parallel size（专家并行）
• GA：Gradient Accumulation steps（梯度累积步数）

可以观察到规模扩展规律：8B 模型在 8 卡单节点即可跑满（TP/PP 均为 1，仅靠 CP 与数据并行）；70B 需要 64 卡并引入 TP=4、PP=8、VP=5 的组合切分；405B 则扩展到 1024 卡，GBS 提升至 512 以维持吞吐。三个规模下单卡 TFLOP/s 均在 780~834 区间，说明 FP8 计算路径在各级并行下都保持了接近峰值的利用率。由于 NeMo 与 Megatron-Core 共享底层训练栈，最新的性能基准建议以官方 NeMo 性能文档为准（该链接为文档内部引用，见 examples/llama/README.md）。

6. 测试数据集与预处理

文档推荐使用WikiText-103作为测试数据集。预处理命令如下：

python "${HOST_MEGATRON_LM_DIR}/tools/preprocess_data.py" \ --input your_dataset.json \ --output-prefix test_dataset \ --tokenizer-type HuggingFaceTokenizer \ --tokenizer-model /path/to/tokenizer.model \ --append-eod

该命令调用仓库中的 tools/preprocess_data.py 将 JSON 文本语料按 tokenizer 切分并序列化为 Megatron 的.bin（token id 数据）与.idx（索引/元数据）格式。预处理完成后，把--output-prefix指定的前缀作为--data-path传入训练脚本即可（如第 3.2 节的HOST_DATA_PREFIX）。--append-eod会在每个文档末尾追加 end-of-document token，配合--split '99,1,0'（train/valid/test 按 99%/1%/0% 划分）使用。更复杂的语料混合、FIM 等场景可参考 contenteditable="false">【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM