Megatron-LM × ModelOpt 高级配置与 Slurm 多节点实战：从变量体系到 Kimi-K2 EAGLE3 训练-北京尧图网络

Megatron-LM × ModelOpt 高级配置与 Slurm 多节点实战从变量体系到 Kimi-K2 EAGLE3 训练【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本篇技术指南基于 Megatron-LM 仓库中的 examples/post_training/modelopt/ADVANCED.md系统讲解 NVIDIA Model OptimizerModelOpt集成示例的高级配置机制与Slurm 多节点执行方式。你将掌握model_conf位置参数与${HF_MODEL_CKPT}、${TP}、${PP}、${EP}等环境变量的协作体系理解sbatch包装脚本的变量传递约束并能直接复刻一套 BF16 Kimi-K2-Instruct EAGLE3 的 8 节点训练与导出流程。所有结论均以当前仓库的实际脚本源码为佐证可直接对照 examples/post_training/modelopt 目录验证。背景ModelOpt 集成示例的定位ModelOptnvidia-modelopt为 NVIDIA 硬件提供端到端模型优化能力涵盖量化真实或模拟、知识蒸馏、剪枝、投机解码等。仓库中的 examples/post_training/modelopt/README.md 明确指出这套示例的核心特性直接优化Megatron-Core 分布式 checkpoint支持各类模型并行度TP、EP、ETP、PP可导出为 TensorRT-LLM、vLLM、SGLang 就绪的统一 checkpoint。ADVANCED.md 正是这套示例中面向复杂场景的补充文档聚焦两大主题如何正确组织配置变量本地与 Slurm 两种模式以及如何用 sbatch 在多节点集群上跑通大型 MoE 模型的 EAGLE3 训练。每个脚本都以bash script.sh [model_conf]形式调用第一个位置参数[model_conf]对应 conf 目录下的模型配置文件例如meta-llama/Llama-3.1-8B-Instruct、deepseek-ai/DeepSeek-R1、moonshotai/Kimi-K2-Instruct。高级配置理解整套变量体系位置参数与${MODEL_ARGS}的映射为保持简洁示例统一采用 Shell 脚本环境变量的传参方式。每个脚本至少要求 1 个位置参数[model_conf]部分脚本需要更多参数例如量化脚本quantize.sh的第二个位置参数[qformat]\ HF_MODEL_CKPTpretrained_model_name_or_path \ bash quantize.sh [model_conf] [qformat]❗ IMPORTANT:model_conf用于获取对应的 Megatron-LM${MODEL_ARGS}。例如meta-llama/Llama-3.1-8B-Instruct或deepseek-ai/DeepSeek-R1均已支持。其底层映射逻辑位于 conf/arguments.sh脚本把第一个参数存入MLM_MODEL_CFG随后执行source ${SCRIPT_DIR}/conf/${MLM_MODEL_CFG}.sh从而把该文件内定义的MODEL_ARGS模型结构参数和TOKENIZER_MODEL注入当前环境。也就是说model_conf本质上是一个配置文件查找键决定了加载哪个模型架构定义。例如conf/meta-llama/Llama-3.1-8B-Instruct.sh 中定义了--num-layers 32 --hidden-size 4096 --ffn-hidden-size 14336 --num-attention-heads 32 --group-query-attention --num-query-groups 8 --rotary-base 500000 --use-rope-scaling等 Llama-3.1 结构参数conf/openai/gpt-oss-120b.sh 则定义了--num-layers 36 --num-experts 128 --moe-ffn-hidden-size 2880 --moe-router-topk 4 --enable-gpt-oss等 GPT-OSS MoE 结构参数该文件注释还提示enable-gpt-oss是启用默认 GPT-OSS 配置的临时 workaround。预训练 checkpoint 通过变量${HF_MODEL_CKPT}提供可以在命令行直接赋值也可以写进配置 Shell 脚本。更多变量如${TP}、${EP}等同样支持命令行传入但文档明确建议把全部变量集中放到一个独立的 Shell 配置脚本中便于复用与维护。使用配置脚本本地环境当${HF_MODEL_CKPT}未在命令行设置时可以使用模板脚本 slurm/env_setup_template.sh 传递所有变量如果你有自己的脚本则通过${SANDBOX_ENV_SETUP}指定\ SANDBOX_ENV_SETUPpath_to_your_script \ bash quantize.sh [model_conf] [qformat]env_setup_template.sh展示了这些变量的典型写法默认 TP1、ETP1、EP1、PP1HF_MODEL_CKPT/workspace/scratch/meta-llama/Llama-3.2-1B-Instruct TP1 ETP1 EP1 PP1从 conf/arguments.sh 的源码可以确认这套变量的默认值体系当SANDBOX_ENV_SETUP已设置时会先source它TP缺省为 1ETP缺省时继承 TP 的值ETP${TP}EP、PP、CP、DP缺省均为 1启动器LAUNCH_SCRIPT缺省为torchrun --nproc_per_node$((ETP * EP * PP * CP * DP))——即进程数由各并行度乘积决定。此外若未设置MLM_SKIP_INSTALL脚本会自动执行pip install -r ${SCRIPT_DIR}/requirements.txt安装依赖并导出TOKENIZERS_PARALLELISMFalse、OMP_NUM_THREADS1、NCCL_IB_SL1、NCCL_IB_TIMEOUT22、CUDA_DEVICE_MAX_CONNECTIONS1等运行环境变量。常见配置变量一览变量含义HF_MODEL_CKPT预训练模型 checkpoint 路径TP张量并行度PP流水线并行度EP专家并行度用于 MoE 模型ETP专家张量并行度用于 MoE 模型MLM_MODEL_SAVEMegatron-LM checkpoint 保存路径MLM_MODEL_LOADMegatron-LM checkpoint 加载路径MLM_EXTRA_ARGS附加 Megatron-LM 参数例如用于不均匀 PP 的--decoder-first-pipeline-num-layers/--decoder-last-pipeline-num-layers以量化脚本 quantize.sh 为例可以看到变量如何被消费脚本基于${MLM_MODEL_CKPT}是否存在决定走--pretrained-model-path ${HF_MODEL_CKPT}从 HF 直接量化还是--load ${MLM_MODEL_CKPT}量化已有 Megatron 分布式 checkpoint分支同时把TP/ETP/EP/PP/CP映射为--tensor-model-parallel-size/--expert-tensor-parallel-size/--expert-model-parallel-size/--pipeline-model-parallel-size/--context-parallel-size。第二个位置参数[qformat]的处理也相当讲究case ${QUANT_CFG}分支传入auto/AUTO/auto_quantize跳过量化配置参数完全由MLM_EXTRA_ARGS中的--auto-quantize-bits驱动混合精度搜索传入含/或以.yaml/.yml结尾的路径通过--recipe传入自定义 recipe其余情况视为内置量化配置全大写名称如FP8_DEFAULT_CFG、NVFP4_DEFAULT_CFG通过--export-quant-cfg传入。Slurm 多节点执行sbatch 包装脚本对于需要多节点的模型示例脚本支持通过一个 sbatch 包装脚本以 Slurm 方式运行。文档给出的起点是 slurm/sbatch.sh可在此基础上做少量修改也可以直接复用自己的现有 sbatch 脚本。与本地执行的关键差异文档强调了一个硬性约束与本地环境不同Slurm 模式下只允许通过 Shell 脚本默认env_setup_template.sh传递变量命令行变量透传不受支持。原因是sbatch与srun不会自动继承你当前 Shell 里的普通环境变量。从 slurm/sbatch.sh 源码可以看到它只显式透传了三个变量--export HF_MODEL_CKPT${HF_MODEL_CKPT},SANDBOX_ENV_SETUP${SANDBOX_ENV_SETUP},LAUNCH_SCRIPT${LAUNCH_SCRIPT}其余所有配置TP、PP、EP、MLM_MODEL_SAVE、MLM_EXTRA_ARGS等都必须由${SANDBOX_ENV_SETUP}指向的脚本在作业内部source后生效——这正是 conf/arguments.sh 中source ${SANDBOX_ENV_SETUP}一行在 Slurm 场景下变得至关重要的原因。sbatch.sh 的必需变量与默认值启动脚本对三个关键变量做了强制检查其余给出默认值并打印告警变量是否必填默认值说明USER_FSW必填缺失即报错退出—可读写的 scratch 空间挂载到容器内/workspace/scratchSANDBOX_DIR可选当前工作目录$(pwd)挂载到容器内/workspace/nmm-sandboxSANDBOX_ENV_SETUP可选./env_setup_template.sh环境配置脚本Slurm 模式下变量传递的唯一途径CONTAINER_IMAGE可选nvidia-modelopt-megatron:latest容器镜像LAUNCH_SCRIPT可选python启动命令作业提交后脚本通过srun -l --mpipmix启动容器--container-workdir指向/workspace/nmm-sandbox--container-mounts把${SANDBOX_DIR}与${USER_FSW}分别挂载最终执行bash ${1}——即命令行传入的第一个参数就是要运行的示例脚本名如eagle3.sh moonshotai/Kimi-K2-Instruct。--output%x_%j_$DATETIME.log会为每次运行生成带作业名、作业 ID 和时间戳的独立日志文件。默认的 SBATCH 头为--nodes1 --ntasks-per-node8 --gpus-per-node8多节点作业需通过sbatch --nodesN覆盖。实战案例BF16 Kimi-K2-Instruct EAGLE3 训练8 节点这是 ADVANCED.md 中最具参考价值的完整案例使用 8 节点 DGX H100每节点 8 卡共 64 张 H100以 BF16 精度训练 Kimi-K2-Instruct 的 EAGLE3 投机解码头。训练配置TP8、ETP1、EP64conf/moonshotai/kimi_k2_instruct.sh 是一个已验证过的配置内容极其精简#!/bin/bash HF_MODEL_CKPT/workspace/scratch/moonshotai/Kimi-K2-Instruct TP8 ETP1 EP64解读TP8表示张量并行度为 8与单节点内 8 卡对应ETP1表示专家权重不做张量切分EP64表示 64 个专家并行组——因为 Kimi-K2 是 MoE 模型其架构参数可见于 conf/moonshotai/Kimi-K2-Instruct.sh--num-experts 384 --moe-layer-freq [0]*1[1]*60 --moe-ffn-hidden-size 2048 --moe-router-topk 8 --moe-token-dispatcher-type alltoall --multi-latent-attention --num-layers 61 --hidden-size 7168等。64 个专家并行组跨 8 个节点分布配合alltoalltoken 派发器实现大 MoE 的并行推理与训练。使用前只需把HF_MODEL_CKPT更新为容器内真实的 checkpoint 路径。提交训练作业export USER_FSWpath_to_scratch_space export CONTAINER_IMAGEpath_to_container_image export SANDBOX_ENV_SETUP./conf/moonshotai/kimi_k2_instruct.sh sbatch --nodes8 slurm/sbatch.sh eagle3.sh moonshotai/Kimi-K2-Instruct注意这里的变量写法USER_FSW、CONTAINER_IMAGE、SANDBOX_ENV_SETUP由 sbatch 包装脚本消费并随srun透传而训练用的并行度变量全部封装在kimi_k2_instruct.sh内、由作业内部 source从而规避了 Slurm 变量透传限制。提交后实际执行的 eagle3.sh 是一个两段式工作流在线 EAGLE3 训练转换阶段若${MLM_MODEL_SAVE}目录不存在先调用convert_model.py以--algorithm eagle3把 HF 格式的 Kimi-K2 checkpoint 转换为 Megatron 格式的 EAGLE3 模型--pretrained-model-path ${HF_MODEL_CKPT}→--save ${MLM_MODEL_SAVE}训练阶段调用finetune.py以--load/--save ${MLM_MODEL_SAVE}加载并持续保存 EAGLE3 模型。在线模式下目标冻结模型与 draft 模型同时驻留内存训练所需的hidden_states即时生成训练过程中会周期性在 MT-Bench 提示词上评估acceptance lengthAL越高越好。eagle3.sh还封装了一整套可直接覆盖的训练超参默认值数据参数默认使用Magpie-Align/Magpie-Llama-3.1-Pro-MT-300K-Filtered数据集、--train-samples 128000、--split 100,0,0优化参数默认--lr 5.0e-5、cosine 衰减、--weight-decay 0.0、--adam-beta2 0.95评估参数默认--eval-interval 1000 --save-interval 1000 --log-interval 100。这些默认值均可通过${MLM_DATA_ARGS}、${MLM_OPTIM_ARGS}、${MLM_TRAIN_ARGS}、${MLM_EVAL_ARGS}覆盖。导出训练好的 EAGLE3 模型PP-only要导出训练好的 EAGLE3 模型切换到导出配置脚本。文档明确导出仅支持流水线并行PP模式。此案例使用 2 个节点PP16export USER_FSWpath_to_scratch_space export CONTAINER_IMAGEpath_to_container_image export SANDBOX_ENV_SETUP./conf/moonshotai/kimi_k2_instruct_export.sh sbatch --nodes2 slurm/sbatch.sh export.sh moonshotai/Kimi-K2-Instructconf/moonshotai/kimi_k2_instruct_export.sh 展示了如何用MLM_EXTRA_ARGS实现不均匀 PPHF_MODEL_CKPT/workspace/scratch/moonshotai/Kimi-K2-Instruct MLM_EXTRA_ARGS \ --decoder-first-pipeline-num-layers 3 \ --decoder-last-pipeline-num-layers 2 \ --init-model-with-meta-device \ --use-cpu-initialization \ # Layer distribution over PP: 3, [4] * 14, 2. PP16注释清楚地写出了层分配方案Kimi-K2 共 61 层PP16 时按3, [4]×14, 2分布——首流水段 3 层、中间 14 段各 4 层、末段 2 层3 4×14 2 61。这就是 ADVANCED.md 中MLM_EXTRA_ARGS: Additional Megatron-LM arguments (e.g., for uneven PP)的具体用法。导出本身由 export.sh 执行其源码有几个值得注意的强制行为导出时强制TP1if [ ${TP} ! 1 ]; then TP1; ... fi无论训练时用了多大的张量并行默认参数含--use-cpu-initialization与导出配置脚本中的--init-model-with-meta-device配合避免显存峰值使用--export-dir ${EXPORT_DIR}输出统一格式 checkpoint供 vLLM、TensorRT-LLM、SGLang 直接加载推理。支持范围参照关于该集成套件中哪些模型支持 EAGLE3可参照 README.md 的支持矩阵deepseek-ai/DeepSeek-R1、meta-llama/Llama-3.1-8B/405B/3.2-1B-Instruct、meta-llama/Llama-4-{Scout,Maverick}、moonshotai/Kimi-K2-Instruct、openai/gpt-oss-{20b,120b}Online以及Qwen/Qwen3系列均支持 EAGLE3而nvidia/NVIDIA-Nemotron-3-{Nano,Super,Ultra}与NVIDIA-Nemotron-Nano-9B-v2等则不支持。使用前务必确认目标模型在 conf 目录下存在对应配置。Checkpoint Resume当前状态说明ADVANCED.md 中Checkpoint Resumecheckpoint 恢复一节目前标记为WIP工作中即官方尚未在高级文档中给出专门的断点续训说明。不过这并不意味着示例不支持恢复能力——从脚本默认参数可以推断其底层机制quantize.sh、eagle3.sh 等脚本的默认参数都包含--finetune --auto-detect-ckpt-format其中--auto-detect-ckpt-format用于自动识别旧版/新版 checkpoint 格式eagle3.sh 训练阶段以--load ${MLM_MODEL_SAVE} --save ${MLM_MODEL_SAVE}加载并回写同一路径天然具备可恢复性训练过程中按--save-interval 1000周期性落盘${MLM_RESUME_ARGS}预留为恢复参数扩展点。因此在实际使用中可以通过设置${MLM_MODEL_CKPT}或训练脚本的--load路径指向已保存的 Megatron-LM 分布式 checkpoint 来续跑等官方 WIP 章节更新后建议优先以仓库最新文档为准。实践建议本地调试先走模板用 slurm/env_setup_template.sh 起步把所有并行度变量写进自己的配置脚本命令行只保留SANDBOX_ENV_SETUP与位置参数保持可复现。Slurm 提交牢记三条铁律USER_FSW必填变量只能走SANDBOX_ENV_SETUP脚本sbatch的命令行参数如--nodes与脚本内 SBATCH 头--ntasks-per-node、--gpus-per-node需按集群实际资源对齐。导出与训练解耦训练阶段可自由使用 TP/EP/ETP 组合但导出阶段会被强制TP1且仅支持 PP 模式大型模型务必配合--init-model-with-meta-device --use-cpu-initialization降低导出内存压力。对齐模型支持矩阵启动任何脚本前先确认model_conf在 conf 下有对应文件避免${MODEL_ARGS}映射失败。如需进一步了解量化配置、Auto Quantize 混合精度搜索、QAT/QAD 蒸馏以及剪枝等具体特性的用法可继续阅读 README.md、distillation.md 与 speculative.md它们与本文的高级配置体系共用同一套变量与脚本框架。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

Megatron-LM × ModelOpt 高级配置与 Slurm 多节点实战：从变量体系到 Kimi-K2 EAGLE3 训练

相关推荐

C语言快递管理系统：链表操作与文件持久化实战解析

神经网络优化共享单车调度路径：从供需预测到路径规划的关键技术

SPD-Conv与CRNN增强的YOLOv7排水管道裂缝检测方法

更多资讯

Simulink构建可验证CDMA扩频通信系统

PHP协程编程：从原理到Hyperf实战优化

脑机接口EMC设计三大致命坑：共模噪声、电极谐振与外壳缝隙

Zabbix邮件报警配置全指南：从SMTP到Action的完整链路

STM32CubeProgrammer安装本质：嵌入式底层信任链重建

Tolaria 中的 AppImage 音频/视频预览外部回退机制：运行时能力门控的设计与实践

今日更新

AI驱动的自动化专利分析平台技术解析

LeetCode-Go 题解：1252. Cells with Odd Values in a Matrix（奇数值单元格计数）双解法剖析

在线答疑系统Java毕设实战：Spring Boot前后端分离与状态流转

本周精选

AI驱动的自动化专利分析平台技术解析

LeetCode-Go 题解：1252. Cells with Odd Values in a Matrix（奇数值单元格计数）双解法剖析

在线答疑系统Java毕设实战：Spring Boot前后端分离与状态流转

本月热门

基于YOLO的焊缝缺陷检测：从数据标注到工业部署全流程实践

分布式服务的决策记录

创业产品的首版范围

读完文章，想聊聊您的网站？