Megatron-LM 静态推理回归测试全解析:Gold Standard Prompts 多提示词 Token 匹配验证实战
【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM
本指南围绕 Megatron-LM 功能测试用例gpt_static_inference_tp1_pp1_16b_multiprompt_tokensmatch展开,深入剖析其"金标准提示词(Gold Standard Prompts)"的设计动机、16B MoE 模型静态推理的完整配置、以及从 CI 配方到 golden values 比对的全链路实现。读完本文,你将理解 Megatron-LM 如何用训练语料中天然高频出现的许可证文本构造"已知补全",从而在不依赖具体设备输出的前提下,对静态推理引擎做确定性回归验证,并掌握在 A100/H100 上维护该类测试黄金值的注意事项。
测试用例概览:定位与目录结构
该测试用例位于 tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_16b_multiprompt_tokensmatch,是 Megatron-LM 功能测试集中一个典型的frozen-start+inference模式用例:加载一份预训练的 16B MoE 检查点,运行静态推理引擎,对多个提示词逐一生成补全,再与预先录制好的 golden values 逐 token 比对,以此守护推理代码路径的正确性与确定性。
目录中共有 5 个文件,职责非常清晰:
| 文件 | 作用 |
|---|---|
| README.md | 阐述金标准提示词的设计哲学与维护约束 |
| model_config.yaml | 完整模型参数、环境变量、指标声明,测试的"灵魂" |
| test_prompts.jsonl | 两条测试提示词(许可证文本片段) |
| golden_values_dev_dgx_a100.json | A100 环境录制的黄金值 |
| golden_values_dev_dgx_h100.json | H100 环境录制的黄金值(含更完整的逐 token 指标) |
Gold Standard Prompts:用"已知补全"而非"设备相关补全"来验证模型
测试用例的 README.md 全文只有 9 行,却浓缩了该测试最核心的方法论,是理解整个用例的钥匙。其核心观点可以概括为三点:
其一,验证目标是"已知补全(known completion)",而不是"设备特定补全(device specific completion)"。回归测试若依赖设备相关的输出,golden values 会随硬件、库版本波动而失效,丧失回归保护价值。因此这里刻意选择模型几乎必然能精确续写的文本,让补全结果由模型"记住的知识"决定,而非由运行时环境决定。
其二,提示词选自训练语料中高频出现的许可证文本。README 明确指出,这些提示词取自已发布的两类常见许可证文本(Creative Commons 与 GNU GPL)的开头部分——这类文本在训练数据集中被大量重复收录,属于"过表征(overrepresented)"内容,因此一个能力合格的模型应当能够精确完成其至少下一整段的内容。这正是把"能否原样续写许可证下一段"作为正确性判据的原因。
其三,段落边界\n\n是"身份保持"的最后观测点。README 特别说明:\n\n(空行分隔的段落边界)似乎是若干模型保持逐字一致(identity)的最后位置,即第一个\n\n出现之前的输出是可靠的。这意味着黄金值记录的是模型从记忆复现的"标准答案",而不是任意生成结果。
README 末尾还给出了一条强约束,测试维护者必须遵守:
Please do not change the gold standard results for a100/h100 for this test without carefully considering if the result is still "correct". These are not arbitrary outputs conditional on a device, they are specific outputs based on a common text that should be overrepresented in training so should be easy for a relatively competent model to complete exactly.
翻译过来即:不得在未经仔细评估"结果是否仍然正确"的情况下修改本测试的 A100/H100 黄金值。这些输出并非设备条件性的任意输出,而是基于训练中过表征的常见文本的确定性输出,一个相对称职的模型应当能精确补全。这条约定保证了:若未来某次改动导致推理结果与黄金值不一致,第一反应应当是排查回归,而不是顺手改黄金值掩盖问题。
提示词样本解析:训练语料中高频出现的许可证文本
test_prompts.jsonl 为 JSONL 格式,每行一个提示词,共两条,分别对应两种许可证:
- Creative Commons Attribution-ShareAlike 4.0 International Public License(CC BY-SA 4.0):
Creative Commons Attribution-ShareAlike 4.0 International Public License By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-ShareAlike 4.0 International Public License ("Public License").- GNU GENERAL PUBLIC LICENSE Version 3(GPL v3):
GNU GENERAL PUBLIC LICENSE Version 3, 29 June 2007 Preamble The GNU General Public License is a free, copyleft license for software and other kinds of works.选择这两条提示词的用意与 README 完全呼应:它们都是互联网与代码仓库中数亿次重复出现的固定文本,模型对此类文本的记忆非常稳固。以 H100 黄金值为例,请求0(CC 许可证)的预期补全为许可证正文第一段的精确续写:
To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions.
请求1(GPL v3)的预期补全为:
The licenses for most software and other practical works are designed to take away your freedom to share and change the works. By contrast, the GNU General Public License is intended to guarantee your freedom to share and change all versions of a program--to make sure it remains free software for all its users. We, the Free Software Foundation, use the GNU General Public License for most of our software; it applies ...
注意两个黄金值文件的补全开头都有一个前导空格,这是 tokenizer 解码结果的一部分,也是逐字比对时不能忽略的细节——这也解释了为什么验证采用逐 token / 逐字符前缀比对而非语义近似比对。
模型与推理配置:16B MoE 静态推理的完整 model_config.yaml
model_config.yaml 是测试的执行清单,包含ENV_VARS、TEST_TYPE、MODE、MODEL_ARGS、METRICS与LAUNCHER六大部分。下面按功能分组逐一解读。
环境变量:为确定性回归测试锁定的运行环境
ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 1 NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE: 1 NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 NCCL_ALGO: Ring CUBLAS_WORKSPACE_CONFIG: :4096:8CUDA_DEVICE_MAX_CONNECTIONS: 1:限制单设备并发 CUDA 连接数,避免多流调度带来的时序抖动。NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0:禁止 TransformerEngine 使用非确定性算法,是"确定性输出"的硬件侧保障。NCCL_ALGO: Ring:固定 NCCL 集合通信算法为 Ring,消除通信算法选择的不确定性(本用例 TP/PP/EP 均为 1,通信面极小,但统一约定更稳妥)。CUBLAS_WORKSPACE_CONFIG: :4096:8:为 cuBLAS 分配确定性 workspace,配合 PyTorch 的确定性模式使用。NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE: 1:允许在加载 TE 检查点时恢复不安全 pickle 扩展状态,保证旧检查点可正常载入。
这些环境变量与后文的--deterministic-mode: true一起,构成了"一次运行、处处可复现"的硬性前提。
测试类型与运行模式
TEST_TYPE: frozen-start MODE: inferencefrozen-start是 run_ci_test.sh 中定义的一组测试类型之一(支持的类型见脚本第 81 行的regular / ckpt-resume / frozen-resume / frozen-start / checkpoint-consistency / release)。对于frozen-start,脚本会将CHECKPOINT_LOAD_PATH指向外部预置的冻结检查点(见脚本第 264-268 行的分支逻辑),而不是先训练再自加载。MODE: inference则决定了后续的验证走"推理管线"分支:结果从推理脚本写出的 JSON 读取,而非从 TensorBoard 提取(见脚本第 398-411 行的注释与分支)。
模型架构参数:DeepSeek 风格 16B MoE
MODEL_ARGS: --transformer-impl: transformer_engine --tensor-model-parallel-size: 1 --pipeline-model-parallel-size: 1 --expert-model-parallel-size: 1 --use-mcore-models: true --num-experts: 64 --moe-router-topk: 6 --moe-grouped-gemm: true --moe-token-dispatcher-type: alltoall --moe-router-load-balancing-type: seq_aux_loss --moe-aux-loss-coeff: 1e-3 --moe-router-score-function: sigmoid --moe-z-loss-coeff: 0 --num-layers: 27 --hidden-size: 2048 --moe-ffn-hidden-size: 1408 --moe-shared-expert-intermediate-size: 2816 --ffn-hidden-size: 10944 --num-attention-heads: 16 --kv-channels: 128 --normalization: RMSNorm --swiglu: true --position-embedding-type: rope --rotary-base: 1000000 --rotary-percent: 1.0 --untie-embeddings-and-output-weights: true --disable-bias-linear: true --init-method-std: 0.014 --attention-dropout: 0.0 --hidden-dropout: 0.0 --seq-length: 4096 --max-position-embeddings: 4096这是一份完整的 DeepSeek 风格 MoE 模型定义,关键点包括:
- 并行配置:TP=1、PP=1、EP=1,整个 16B 模型单卡即可容纳,是回归测试追求"最小资源、最快反馈"的典型取舍。
--moe-token-dispatcher-type: alltoall+--moe-grouped-gemm是当前主流的 MoE 高效执行组合。 - MoE 结构:64 个专家、Top-6 路由、共享专家中间维度 2816,非 MoE 的 FFN 维度 10944。路由采用
seq_aux_loss负载均衡 +sigmoid评分函数(DeepSeek 风格 softmax-free 路由)。 - Transformer 结构:27 层、hidden 2048、16 注意力头、KV 通道 128、RMSNorm、SwiGLU、RoPE(rotary base 1e6)——与检查点路径中的
deepseek_16b_pyt命名一致。 - 推理无关但必须一致的架构参数:
--init-method-std、dropout 等训练期参数在此仅为与检查点元数据对齐而保留,--attention-dropout: 0.0/--hidden-dropout: 0.0确保推理期无随机性。
检查点加载与推理管线参数
--load: ${CHECKPOINT_LOAD_PATH}/model/deepseek_16b_pyt/dcp/mcore-v1_bf16/checkpoints --tokenizer-model: ${CHECKPOINT_LOAD_PATH}/model/deepseek_16b_pyt/dcp/mcore-v1_bf16/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json --tokenizer-type: TikTokenizer --tiktoken-pattern: v2 --ckpt-format: torch_dist --dist-ckpt-optim-fully-reshardable: true --ckpt-fully-parallel-save: true --ckpt-fully-parallel-load: true --ckpt-assume-constant-structure: true --dist-ckpt-strictness: log_unexpected --use-checkpoint-args: true --no-use-tokenizer-model-from-checkpoint-args: true --no-load-optim: true --bf16: true --deterministic-mode: true- 检查点与词表通过
${CHECKPOINT_LOAD_PATH}环境变量注入,该变量由 CI 配方(recipe)统一赋值。词表为 500K 规模的 128K 词汇 GPT-4o 风格 TikTokenizer。 - 分布式检查点(
torch_dist格式)全面启用并行保存/加载与--ckpt-assume-constant-structure,加载时对结构变化的处理策略为log_unexpected(记录但不阻断)。 --use-checkpoint-args: true表示加载检查点时恢复其保存时的参数,而--no-use-tokenizer-model-from-checkpoint-args: true则明确禁止从检查点参数中取词表路径,二者配合保证"架构参数跟检查点走、词表路径跟配置走"。- 推理测试不加载优化器状态(
--no-load-optim),配合--deterministic-mode: true锁定整体确定性。
静态推理与采样相关参数同样关键:
--micro-batch-size: 1 --inference-max-requests: 1 --num-tokens-to-generate: 80 --inference-max-seq-length: 4096 --max-tokens-to-oom: 3600000 --temperature: 1.0 --top_k: 1 --return-log-probs: true --attention-backend: flash --flash-decode: true --no-create-attention-mask-in-dataloader: true --use-legacy-static-engine: true --prompt-file: ./tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_16b_multiprompt_tokensmatch/test_prompts.jsonl --incoming-requests-per-sec: -1 --inference-logging-step-interval: 1 --inference-dynamic-batching-buffer-size-gb: 20 --output-path: ${INFERENCE_OUTPUT_PATH}几个值得展开的细节:
--micro-batch-size: 1的注释道出了它的真实目的:"确保 batch size 为 1,以测试连续多次推理调用覆盖状态清理(state cleanup)"。也就是说,这个用例不仅验证输出正确性,还顺带验证静态引擎在连续请求间是否正确清理 KV cache / 中间状态。--inference-max-requests: 1与--incoming-requests-per-sec: -1("所有请求一次性到达")配合,让两条提示词以单请求并发上限的形态进入引擎。--temperature: 1.0+--top_k: 1是贪心解码(greedy),是输出可复现的采样侧保证;--return-log-probs: true让输出附带逐 token logprob,供更严格的数值比对。--use-legacy-static-engine: true选择传统静态推理引擎路径。在 gpt_static_inference.py 中可以看到,该开关决定StaticInferenceEngine是否携带buffer_size_gb(动态批处理缓冲)参数,legacy引擎则不使用该缓冲。- 每条提示词生成
--num-tokens-to-generate: 80个 token,与两份 golden values 中generated_tokens数组均为 80 个元素的事实吻合。
评测指标声明
METRICS: - "generated_text"METRICS列表声明了本用例实际参与比对的指标。此处仅声明generated_text,意味着即使推理脚本输出了 logprobs、tpot、latency 等字段,比对阶段也只校验生成文本——比对逻辑见下文"pytest 对比"小节。LAUNCHER: ft_launcher则指定了功能测试专用的启动器。
从配置文件到回归测试的完整执行链路
测试配方:recipe YAML 如何组织 CI 产品
测试由 tests/test_utils/recipes/h100/gpt-static-inference.yaml 编排。该配方文件将多个静态推理用例聚合为一个"产品族":
- 顶层
spec声明model: gpt、nodes: 1、gpus: 1、n_repeat: 1,单节点单卡运行; script_setup负责在 CI 环境中检出指定 commit 的仓库(含 mcore 主仓与 backwards-ref 旧版目录);script定义ARGUMENTS数组,其中关键映射包括:TRAINING_SCRIPT_PATH=examples/inference/advanced/gpt_static_inference.py:指定推理入口脚本;TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml:把model_config.yaml注入测试;GOLDEN_VALUES_PATH=.../golden_values_{environment}_{platforms}.json:按环境+平台拼接黄金值路径;INFERENCE_OUTPUT_PATH={assets_dir}/golden_values_{environment}_{platforms}.json:推理实际输出也写入同名风格的文件;- 随后统一调用
bash ./tests/functional_tests/shell_test_utils/run_ci_test.sh ${ARGUMENTS[@]}。
在products一节,本用例被登记为:
- test_case: [gpt_static_inference_tp1_pp1_16b_multiprompt_tokensmatch] products: - environment: [dev] scope: [mr, mr-github] platforms: [dgx_h100]即该用例在 MR 与 MR-GitHub 两个 CI 范围内、dev 环境、DGX H100 平台执行。目录中同时保留的golden_values_dev_dgx_a100.json表明 A100 黄金值也处于维护状态,README 中的 "a100/h100" 警告与此对应。
驱动脚本:run_ci_test.sh 如何处理 frozen-start 推理测试
run_ci_test.sh 是整个功能测试的通用驱动,对本用例而言关键路径如下:
- 校验
TRAINING_SCRIPT_PATH、TRAINING_PARAMS_PATH、GOLDEN_VALUES_PATH、CHECKPOINT_LOAD_PATH、INFERENCE_OUTPUT_PATH等必需环境变量(脚本第 42-58 行); - 从
model_config.yaml用yq解析TEST_TYPE、MODE等字段(第 63-78 行); - 由于
TEST_TYPE=frozen-start,将CHECKPOINT_LOAD_PATH指向外部冻结检查点(第 264-268 行); - 通过
run_training_phase调用_run_training.sh执行推理脚本gpt_static_inference.py(第 228 行); - 在
MODE=inference && TEST_TYPE=frozen-start分支下,运行 pytest 比对(第 476-483 行):
uv run --no-sync pytest -s -o log_cli=true --log-cli-level=info \ $ROOT_DIR/tests/functional_tests/python_test_utils/test_inference_regular_pipeline.py \ --golden-values-path $GOLDEN_VALUES_PATH \ --test-values-path $INFERENCE_OUTPUT_PATH \ --model-config-path ${TRAINING_PARAMS_PATH} \ $ALLOW_NONDETERMINISTIC_ALGO_ARG注意--allow-nondeterministic-algo参数由环境变量NVTE_ALLOW_NONDETERMINISTIC_ALGO决定(第 436-439 行):本用例将其设为 0,因此不会放宽确定性要求。
推理入口:gpt_static_inference.py 的引擎装配与结果落盘
examples/inference/advanced/gpt_static_inference.py 是该用例的推理入口,其核心装配逻辑(第 56-82 行)为:
- 构建
StaticInferenceContext(args.inference_max_requests, args.inference_max_seq_length); - 用
GPTInferenceWrapper包装模型,配合TextGenerationController与 tokenizer; - 依据
args.use_legacy_static_engine选择传统静态引擎或带buffer_size_gb的新引擎; - 采样参数由
temperature / top_k / top_p / return_log_probs / num_tokens_to_generate组成(第 140-147 行)。
结果落盘逻辑在 rank 0 上执行(第 172-190 行):每个请求输出input_prompt、generated_text、generated_tokens(转为 list)、tpot、latency,并在return_log_probs时附上logprobs(提示词与生成 token 的对数概率拼接),最后写入args.output_path——即 recipe 中的INFERENCE_OUTPUT_PATH,供 pytest 读取。脚本还通过torch.cuda.memory_stats()打印峰值显存,并输出一条含请求数、batch、显存与延迟的汇总日志(第 214-236 行)。
对比验证:pytest 中的 generated_text 前缀匹配与 logprobs 容差
test_inference_regular_pipeline.py 是推理类测试的统一比对器,对本用例生效的逻辑包括:
- 请求 ID 一致性:黄金值中的请求 ID 必须覆盖实际输出(第 77-88 行);
generated_text前缀比对(第 179-192 行):对每个请求,取黄金值与当前输出生成文本的公共最小长度min_len,断言前min_len个字符完全一致;任一为空即失败;generated_tokens逐 token 相等(第 156-163 行):仅当指标列表包含generated_tokens时启用;logprobs长度一致 + 绝对容差比对(第 165-177 行):逐位置断言math.isclose(lp1, lp2, abs_tol=0.001),仅当指标列表包含logprobs时启用;routing_indices排序比对(第 194-206 行):针对 MoE 模型可选的专家路由索引校验。
由于本用例METRICS仅声明["generated_text"],实际生效的是前缀比对路径。这种"配置驱动比对项"的设计使得同一个比对器可复用于不同粒度的用例(logitsmatch、tokensmatch、cudagraphs 等)。
Golden Values 文件解读:A100 与 H100 的结构差异
两份黄金值文件结构不同,体现了录制与演进过程:
- golden_values_dev_dgx_a100.json 采用精简结构,仅按请求 ID 记录
generated_text(附id字段),两段补全分别停在 CC 许可证第一段与 GPL v3 Preamble 第一段; - golden_values_dev_dgx_h100.json 则记录了完整字段:
input_prompt、generated_text、generated_tokens(80 个 token id)、tpot(每 token 生成耗时,稳定在 0.075-0.078 秒量级)、latency(总延迟约 14 秒)以及逐 token 的logprobs。H100 版本中请求 0 的补全在 CC 许可证第一段后继续穿过了\n\n段落边界,多输出了 "License Elements" 定义段的部分内容。
两份文件的内容差异恰好印证了 README 的核心论断:虽然目标是"已知补全",但不同设备/库版本下模型可能在不同位置越过\n\n边界,因此黄金值需要按设备分别维护,且任何改动都必须以"结果仍然正确"为前提。同时由于METRICS只声明generated_text,即使 H100 文件携带了 tpot、latency、logprobs 等附加数据,比对阶段也只校验生成文本的前缀一致性。
维护与演进注意事项
README 的维护铁律
回归测试的黄金值维护必须遵循 README 的两条原则:
- 改值前先确认"正确性":若新输出与黄金值不一致,先判断新结果是否是许可证文本的精确续写;若是推理回归导致偏离标准文本,则应修复代码而非更新黄金值。
- 黄金值是"标准答案"而非"设备快照":补全内容由训练语料中过表征的固定文本决定,任何称职的模型都应能精确复现,因此这类黄金值天然具备跨实现的普适性。
相关性能测试注释揭示的选型约束
仓库中与本案共享同一模型配置的还有性能测试参数文件 tests/performance_tests/server/model_args/gpt_16b.args,其头部注释(标注为"reviewer asked for the full kitchen-sink list"后记录的选型纪要)从维护者视角说明了为什么本用例停留在transformer_engine+ chunked prefill 的安全子集:
--transformer-impl inference_optimized会拒绝--swiglu,而本模型使用 SwiGLU,因此无法启用;- vllm 融合 MoE、
nvlsdispatcher、shared-expert overlap 等参数均依赖inference_optimized,因传递闭包一并被排除; - 在
full_iteration_inference作用域启用 CUDA graph 会在捕获阶段崩溃,原因是 alltoall MoE dispatcher 在其 forward 中执行了d2h event.synchronize()(对应token_dispatcher.py中的_maybe_dtoh_and_synchronize),这在捕获流(capturing stream)中是非法的。
这段注释虽然来自性能测试目录,却直接解释了功能测试用例配置的边界来源,可作为后续在该 MoE 推理路径上扩展优化能力(如图形友好的 MoE dispatcher)时的背景依据。
结语
gpt_static_inference_tp1_pp1_16b_multiprompt_tokensmatch是理解 Megatron-LM 推理回归测试体系的绝佳样本:它用训练语料中天然确定的许可证文本作为提示词,以"已知补全"替代"设备相关补全",配合frozen-start的冻结检查点加载、确定性环境锁定、单卡 16B MoE 的轻量配置,以及 recipe → run_ci_test.sh → gpt_static_inference.py → pytest 的四级链路,构成了一套低成本、高信号、可跨硬件维护的静态推理回归保护。当你在 Megatron-LM 中修改推理引擎、MoE 调度或 CUDA graph 相关代码时,这个用例正是第一道防线。
【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考