
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),僅供參考