Megatron-LM 推理实战指南基于 Megatron Core 高层 API 的离线推理与 OpenAI 兼容服务【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本指南以examples/inference/下的官方示例为骨架系统讲解 Megatron-LMMegatron Core的两大推理入口面向批量离线生成的offline_inference.py以及面向 OpenAI 兼容 HTTP 服务的launch_inference_server.py并延伸介绍低层 API 进阶脚本与 MoE 路由轨迹分析工具。读完本文你将掌握 sync/async、direct/coordinator 四种执行模式的取舍、Qwen 2.5 与 Nemotron 混合 MoE 两类真实启动脚本的每个关键参数以及如何采集并分析 MoE 路由分布数据来验证“按层静态缓存”等优化假设。一、从整体到入口高层推理 API 与两个顶层脚本Megatron Core 的推理能力由megatron/core/inference/apis/下的两个高层类封装MegatronLLM同步提供 vLLM 风格的generate(prompts, sampling_params)API返回list[DynamicInferenceRequest]同时支持 direct直接与 coordinator协调器两种执行模式MegatronAsyncLLM异步基于 asyncio 的封装generate为协程支持serve(...)启动 OpenAI 兼容 HTTP 前端必须使用 coordinator 模式其构造器在use_coordinatorFalse时会直接抛出ValueError原因是引擎内部的 asyncio 原语会绑定到调用方事件循环与同步engine.generate()路径冲突见 async_llm.py。两个类内部统一管理底层引擎流水线DynamicInferenceContext、GPTInferenceWrapper、TextGenerationController、DynamicInferenceEngine并提供pause/unpause/suspend/resume/shutdown等生命周期控制。完整 API 心智模型见 megatron/core/inference/README.md更全面的用户手册见 docs/mcore-inference-user-guide.md。examples/inference/目录下两个顶层 Python 入口覆盖了全部常见工作流入口脚本定位替代的旧路径offline_inference.py批量离线生成支持 syncdirect、synccoordinator、asynccoordinator 三种模式组合通过 CLI 标志切换gpt_dynamic_inference.py、gpt_dynamic_inference_with_coordinator.pylaunch_inference_server.py基于MegatronAsyncLLM.serve(...)的 OpenAI 兼容 HTTP 服务器tools/run_dynamic_text_generation_server.py共享工具函数集中在 examples/inference/utils.pyRequest单条请求的状态机记录 prompt 文本、token、到达/开始/结束时间与ttft、build_requests按--prompts/--prompt-file/ 合成请求三种来源构造请求列表、build_dynamic_engine_setup_prefix输出动态批处理配置摘要、输出格式化与 JSON dump 函数。其中合成请求的到达时间用simpy模拟泊松到达过程get_time_offsets因此运行示例前需要pip install simpy。二、环境前置条件已按官方方式安装 Megatron-LM含 MCore 依赖Python 环境可运行torchrunpip install simpyutils.py中合成请求到达模拟所需脚本注释明确要求一个 Megatron 格式的 checkpoint离线示例为 Qwen 2.5-1.5B服务示例为 Nemotron-6 3B 混合 MoE与对应的 Hugging Face tokenizer需要--hf-token下载;具备多卡环境离线脚本默认 8 进程服务脚本默认 8 进程TP2、EP8。三、离线推理offline_inference.py 的三种运行模式offline_inference.py在 Megatron 模型上执行合成负载推理输出三部分内容setup-prefix 配置摘要行、“Unique prompts outputs” 表格、吞吐量总结还可通过--output-path输出 JSON dump 用于回归测试。3.1 三种模式与 shell 封装run_offline_inference.sh 封装了典型的 Qwen 2.5-1.5B 配置必需参数--hf-tokenHugging Face token用于下载 tokenizer、--checkpointcheckpoint 路径透传为--load可选参数--mode sync|async默认sync选择MegatronLLM还是MegatronAsyncLLM、--use-coordinator默认关闭即 direct 模式、--nproc n默认8。# sync direct默认 bash examples/inference/run_offline_inference.sh \ --hf-token HF_TOKEN --checkpoint /path/to/qwen-1.5b # sync coordinator bash examples/inference/run_offline_inference.sh \ --hf-token HF_TOKEN --checkpoint /path/to/qwen-1.5b --use-coordinator # async coordinator bash examples/inference/run_offline_inference.sh \ --hf-token HF_TOKEN --checkpoint /path/to/qwen-1.5b --mode async --use-coordinator注意async direct 目前不受支持——MegatronAsyncLLM构造器强制要求use_coordinatorTrue见 async_llm.py而MegatronLLM可同时用于 sync 的 direct 与 coordinator 模式。文档明确说明所有可行模式产生数值相同的生成文本。3.2 脚本底层的参数校验与运行流阅读 offline_inference.py 源码可以发现两类关键校验_validate_high_level_api_args--use-coordinator与--inference-repeat-n 1互斥。原因是engine.reset()与 coordinator 模式下运行在 runtime 线程上的engine_loop_task存在竞争同时--prompt-file与--num-tokens-from-file组合也会被拒绝——高层 API 每次generate()调用只接受一个sampling_params无法表达逐请求的生成长度应改用统一的--num-tokens-to-generate。_validate_prompt_lengths未开启 chunked prefill 时会断言所有 prompt 长度不超过llm.context.max_tokens。运行流方面sync 与 async 分支结构对称在with MegatronLLM(...) as llm/async with MegatronAsyncLLM(...)上下文内仅 primary rank 提交任务并打印 setup 前缀随后循环inference_repeat_n次generate用torch.cuda.reset_peak_memory_stats()归零显存统计、按输出 token 数计算吞吐量tok/s。coordinator 模式下时间测量不做 rank0 广播do_broadcastnot args.use_coordinatorworker 进程在__exit__中阻塞直到收到 STOP 传播。最后在引擎关闭后对所有 rank 做峰值显存 MAX 归约get_global_peak_memory_stats_bytes输出形如~~~ setup prefix … throughput: X.XXX tok/s … total time: X.XXXs … mem A/B GB … steps: N … capture -- ~~~3.3 内嵌的模型配置参数脚本末尾的torchrun命令段完整指定了 Qwen 2.5-1.5B 的架构参数逐项对应模型超参--bf16、--tensor-model-parallel-size 1、--micro-batch-size 64、--dist-ckpt-strictness log_unexpected、--inference-rng-tracker推理 RNG 追踪保证采样可复现、--cuda-graph-impl local本地 CUDA 图实现、--decode-only-cuda-graphs仅解码阶段捕获图、--tokenizer-type HuggingFaceTokenizer、--tokenizer-model Qwen/Qwen2.5-1.5B、--no-use-tokenizer-model-from-checkpoint-args以及--num-layers 28、--hidden-size 1536、--num-attention-heads 12、--num-query-groups 2GQA、--swiglu、--normalization RMSNorm、--disable-bias-linear、--position-embedding-type rope、--rotary-percent 1.0、--rotary-base 1000000、--seq-length 32768、--ffn-hidden-size 8960。更换模型时需同步替换这些架构参数与 checkpoint。四、OpenAI 兼容推理服务器launch_inference_server.pylaunch_inference_server.py通过MegatronAsyncLLM.serve(blockingTrue)在 coordinator 引擎之上启动 HTTP 前端暴露/v1/completions与/v1/chat/completions两个端点仅 global rank 0 提供服务其余 rank 跳过 HTTP 安装但仍遵守blocking语义以保证所有进程同步返回见 llm.py 与 async_llm.py。4.1 启动命令run_inference_server.sh 封装了 Nemotron-6 3B混合 MoE配置TP2、EP8、PP1bash examples/inference/run_inference_server.sh \ --hf-token HF_TOKEN \ --hf-home /path/to/hf_home \ --checkpoint /path/to/nemotron-3b-hybrid-moe必需参数--hf-token、--hf-homeHF 缓存目录、--checkpoint可选--nproc n默认8脚本会导出CUDA_DEVICE_MAX_CONNECTIONS1Megatron 在使用张量或上下文并行时的必需项。就绪后约 2 分钟Nemotron-6 3B会出现横幅INFO:root:Inference co-ordinator is ready to receive requests! INFO:hypercorn.error:Running on http://0.0.0.0:5000 (CTRL C to quit)4.2 模型与批处理相关参数速览服务脚本中的关键配置项及其作用--tensor-model-parallel-size 2、--expert-model-parallel-size 8、--pipeline-model-parallel-size 1并行切分方案TP2、EP8、PP1配合--sequence-parallel与--moe-token-dispatcher-type alltoall分发 MoE token--transformer-impl inference_optimized使用推理优化后的 transformer 实现--attention-backend flash、--enable-chunked-prefillflash attention 与分块 prefill--cuda-graph-impl local、--cuda-graph-scope full_iteration_inference、--inference-dynamic-batching-num-cuda-graphs -1CUDA 图覆盖整个推理迭代--inference-dynamic-batching-buffer-size-gb 20、--inference-dynamic-batching-max-tokens 2048、--inference-dynamic-batching-max-requests 256动态批处理缓冲区上限--inference-max-seq-length 4096、--inference-logging-step-interval 50、--return-log-probs、--moe-router-dtype fp32。4.3 服务端采样默认值、eval-mode 与请求行为采样默认值--default-temperature默认1.0、--default-top-p默认1.0、--default-top-k默认0用于填充请求中缺失的采样字段请求级取值永远优先于默认值。这些参数对应 ServeConfig 中的同名字段。--eval-mode对评估等纯 serving 负载默认不返回 prompt token ID除非请求显式要求chat 请求默认prevent_retokenizationfalse单个请求仍可通过prevent_retokenization或return_tokenized_data自行开启。--parsers启用的响应解析器名如json、tool_use透传给底层 text-generation server--verbose打开逐请求 HTTP 日志--frontend-replicas控制主 rank 上派生的 HTTP 前端进程数默认 4。模型名字段动态服务器当前返回model: EMPTY且不校验请求的model字段客户端可随意传任意值。服务端代码将ServeConfig透传给start_text_gen_server见 llm.py其中coordinator_host与host语义不同前者是 coordinator 内部 ZMQ 流量地址后者是 HTTP 对外监听地址。任何 OpenAI 兼容客户端均可直接请求。五、进阶示例直接驱动低层 APIexamples/inference/advanced/下的脚本绕过高层 API直接驱动megatron.core.inference的低层接口gpt_dynamic_inference.py手动add_request/step_modern步进循环offline 示例的底层原型gpt_dynamic_inference_with_coordinator.py显式管理 coordinator 与InferenceClient生命周期gpt_static_inference.py静态引擎推理simple_t5_batch_inference.pyT5 批量推理run_prefix_cache_lru_resume_repro.sh前缀缓存 LRU 恢复复现脚本。适用场景需要步级调度控制、自定义 forward-step / 采样集成或正在迁移既有推理流水线。对于典型工作流官方明确建议优先使用offline_inference.py与launch_inference_server.py。CI 配方tests/test_utils/recipes/下的 h100{gpt,moe,mamba}-*-inference.yaml目前仍针对这些进阶脚本运行。六、MoE 路由分析工具从轨迹采集到分布可预测性分析tools/moe_routing/下的analyze_routing.py与analyze_routing_*.py脚本分析 MoE 模型各层的 top-K 路由决策。JSONL 轨迹格式与分析脚本对训练与推理两种场景通用。6.1 两种采集路径Sink vs Hook路径开启方式捕获内容CUDA graphsSink--moe-enable-routing-replay legacy 调度仅 top-K 索引开启Hook不开启 replay --cuda-graph-impl none索引 隐状态 路由权重必须关闭关键差异只有 Hook 路径捕获analyze_routing_predictability.py所需的隐状态与路由权重Sink 填充的是一个进程内管道缓冲区只保存索引。因此做路由集中度 / 负载均衡分析廉价且图安全用 Sink需要保存隐状态、权重等昂贵数据时用 Hook。6.2 采集命令训练Hook 路径--moe-routing-trace-path /path/to/trace_dir # 开启追踪 --moe-routing-trace-max-training-iters 500 # 可选N 次迭代后停止 --moe-routing-trace-capture-hidden-states # 供 predictability 分析 --moe-routing-trace-dump-weights # 供 predictability 分析注意Python forward hook 在 CUDA 图重放期间不会触发因此训练时必须禁用 MoE cudagraph否则被图捕获的层会被静默跳过。推理——Sink仅路由索引图开启路由 replay 需要 legacy 调度由于异步调度默认开启开启 routing replay sink 时必须显式选择 legacy 模式--moe-routing-trace-path /path/to/trace_dir --moe-routing-trace-max-inference-steps 200 --moe-enable-routing-replay --inference-dynamic-batching-async-sched-mode legacy推理——Hook为 predictability 增加隐状态与权重--moe-routing-trace-path /path/to/trace_dir --moe-routing-trace-max-inference-steps 200 --cuda-graph-impl none --moe-routing-trace-capture-hidden-states --moe-routing-trace-dump-weights所有路径都写出router_trace_rank{N}.jsonl每个 rank 一个文件--moe-routing-trace-capture-hidden-states额外写出hidden_states_rank{N}.bin--moe-routing-trace-dump-weights写出router_state_rank{N}.pt后两者是analyze_routing_predictability.py的必需输入。6.3 轨迹格式与底层实现轨迹由 megatron/core/transformer/moe/router_trace.py 中的RouterTracer类实现每 (step, block, layer) 一条记录{step: 0, stage: pre_dispatch, block: decoder, layer: 3, rank: 0, num_tokens: 128, topk: 22, top_indices: [[12, 45, ...], ...]}MTP 记录额外携带mtp_idx字段以避免与共享层号的 decoder 层冲突sidecar 二进制文件中每个 JSONL 记录会获得hs_offset/hs_bytes/hs_shape隐状态与logit_offset/logit_bytes/logit_shapepre-topk 路由 logits字段可用load_hidden_states_for_record/load_logits_for_record读取。6.4 运行分析python tools/moe_routing/analyze_routing.py /path/to/trace_dir --num-experts 512dispatcher 按顺序执行以下分析脚本核心问题作用analyze_routing_concentration.py路由有多集中hot-set 大小假设检验按层静态缓存是否可行高集中度比率 2×支持该方案近均匀分布则排除analyze_routing_predictability.pyL-1 层的隐状态能在多大程度上预测 L 层的路由分布肯定信号高余弦/斯皮尔曼相关意味着分布级路由可提前一层预测6.5 解读分布可预测性输出analyze_routing_predictability.py将 L 层的路由权重应用于 L-1 层的隐状态得到“预测的逐专家 token 数分布”再与实际路由结果比较。这提供了一个衡量“上一层 MoE 层的隐状态信号是否足以预测下一层的专家负载聚合分布”的示例数值接近 0 说明该层对之间的跨层信号很弱。注意这是分布级结果逐 token 的分配误差会在聚合计数直方图中相互抵消。6.6 新增路由指标的正确姿势要增加新的路由指标应把捕获逻辑放进 megatron/core/transformer/moe/router_trace.py作为RouterTracer类的一部分使其同时服务于训练与推理避免在 megatron/training/activation_logging.py 中为路由指标添加专属日志流——该文件负责轻量计数监控tokens_per_expert输出格式不同。七、测试与更多资源功能测试tests/functional_tests/test_cases/gpt/下的gpt_offline_inference_*与gpt_inference_server_smoke_*用例覆盖离线推理与服务器冒烟场景单元测试tests/unit_tests/inference/high_level_api/针对高层 APIMegatronLLM/MegatronAsyncLLM同目录还包含test_dynamic_text_generation_server_cli.py、test_openai_streaming.py、test_chat_completions.py等服务器相关测试API 参考megatron/core/inference/README.md低层引擎megatron/core/inference/完整用户手册docs/mcore-inference-user-guide.md含支持特性、direct 与 coordinator 模式对比、已知限制与路线图。八、小结Megatron-LM 的高层推理 API 让同步/异步、direct/coordinator 四类执行模式通过极简 CLI 标志即可切换并将动态批处理、分块 prefill、分页注意力、CUDA 图、MoE 专家并行等底层机制封装在引擎内部配合--output-path的 JSON dump 与 MoE 路由轨迹分析既可以作为训练/评估/RL 的数值一致生成后端也能用于路由行为的离线研究与优化假设验证。实际部署时请根据“是否需要 HTTP 服务”coordinator 必选、“是否需要异步”MegatronAsyncLLM与“是否需要步级控制”advanced 低层脚本三个问题选择入口。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考