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-LMMegatron-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 系列模型。FP88 位浮点通过 Transformer EngineTE的自动混合精度机制将 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 HopperH100 等、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_IMAGEnvcr.io/nvidia/pytorch:25.03-py3 docker run --rm --gpus all --ipchost --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 \ 21 | 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 使用自定义数据与 TokenizerPYTORCH_IMAGEnvcr.io/nvidia/pytorch:25.03-py3 docker run --rm --gpus all --ipchost --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}) \ 21 | 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 默认并行策略并行维度默认值Tensor ParallelTP1Pipeline ParallelPP1Context ParallelCP2对应脚本中的TP_SIZE1、CP_SIZE1、PP_SIZE1注脚本实际默认CP_SIZE1README 中给出的是该配方文档层面的默认策略具体以脚本为准。脚本同时始终开启--sequence-parallel并注释说明TP_SIZE2 时总是启用序列并行。对于 8 卡单节点的 Llama-3-8B数据并行 序列并行即可承载模型无需张量/流水并行切分。4.2 Llama-3-8B 模型架构配置项值层数num-layers32Hidden size4096FFN hidden size14336注意力头数num-attention-heads32Query groupsnum-query-groups8KV channelskv-channels128序列长度seq-length8192归一化RMSNorm--normalization RMSNorm激活函数SwiGLU--swiglu位置编码RoPE--position-embedding-type rope--rotary-base 1000000--rotary-percent 1.0--num-query-groups 8配合--group-query-attention实现了 Llama-3 的GQAGrouped 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-weightsembedding 与 LM head 权重不共享、--disable-bias-linear线性层去 bias、--init-method-std 0.0134权重初始化标准差。4.3 关键训练参数参数值说明微批大小micro-batch-size1每步每个 GPU 的样本数全局批大小global-batch-size128有效批大小学习率lr1.5e-4最小学习率min-lr1.0e-5cosine 衰减下限权重衰减weight-decay0.1FP8 格式fp8-formathybride4m3 e5m2 混合优化器AdamWbeta10.9beta20.95使用 decoupled 变体梯度裁剪clip-grad1.0训练样本数train-samples1,953,125,000约 15T token8B 模型惯例规模预热样本lr-warmup-samples3,906,252衰减样本lr-decay-samples1,949,218,748退出时长exit-duration-in-mins235性能基准模式约 4 小时自动退出脚本同时开启了--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、输出激活梯度使用 e5m2e5m2 的动态范围更大适合梯度这类分布剧烈的张量而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是项目内置的合成 tokenizernull_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 数/并行配置均为完整事实摘录如下Model#-GPUsGBSMBSSeq LengthTPPPCPVPEPGATokens/sec/GPUTFLOP/sec/GPULLAMA3-8B812818192112113213812800LLAMA3-70B641281819248151641621780LLAMA3-405B1024512181928828164315834图例LegendGBSGlobal Batch Size全局批大小MBSMicro Batch Size微批大小TPTensor Parallel size张量并行PPPipeline Parallel size流水并行CPContext Parallel size上下文并行VPVirtual Pipeline stages虚拟流水阶段数EPExpert Parallel size专家并行GAGradient Accumulation steps梯度累积步数可以观察到规模扩展规律8B 模型在 8 卡单节点即可跑满TP/PP 均为 1仅靠 CP 与数据并行70B 需要 64 卡并引入 TP4、PP8、VP5 的组合切分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 的.bintoken id 数据与.idx索引/元数据格式。预处理完成后把--output-prefix指定的前缀作为--data-path传入训练脚本即可如第 3.2 节的HOST_DATA_PREFIX。--append-eod会在每个文档末尾追加 end-of-document token配合--split 99,1,0train/valid/test 按 99%/1%/0% 划分使用。更复杂的语料混合、FIM 等场景可参考 contenteditable="false">【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考