企业官网建设流程全解析

Megatron-LM 梯度一致性测试实战：基于 Optimizer State 提取与跨并行配置校验

【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM

本指南以 Megatron-LM 仓库中tests/functional_tests/test_cases/gpt/gpt3_mcore_reruns_resume_check_grads测试用例为核心，系统讲解如何在单机单卡基准与 DP/TP/CP/PP 各种并行配置之间验证梯度一致性。读完本文，你将掌握“利用 Adam 动量缓存提取梯度”、“用旧版分布式 checkpoint 格式固化优化器状态”、“跨配置比对张量并处理行并行（Row Parallel）分片”的完整技术方案，并能据此扩展出属于你自己的梯度一致性测试。

测试的整体设计：为什么能“用 checkpoint 比对梯度”

该测试属于TEST_TYPE: checkpoint-consistency，其核心思路并非在训练过程中直接钩取梯度，而是利用优化器的动量状态作为梯度的“快照存储”：

• 测试只训练 1 个 step（--train-iters: 1），以收集这一轮的梯度。
• 将 Adam 的beta1、beta2均设为0.0，此时动量更新公式变为exp_avg = grad，即上一轮的梯度会直接覆盖写入Adam 的第一个动量状态exp_avg。
• 训练完成后，从 checkpoint 的 optimizer state 中把exp_avg张量取出来，得到的就是这一轮训练的真实梯度。
• 不同运行（基准配置 vs 各种并行配置）只要共享同一份数据与随机种子，其梯度在数学上应当一致，因此通过比对各自 checkpoint 中的exp_avg即可验证梯度一致性。

这一设计的巧妙之处在于：它完全绕开了训练框架内部的梯度收集接口，只依赖 checkpoint 文件这一对外产物，因此可以跨 DP/TP/CP/PP 各种配置公平比对，甚至可以推广到“任何应当产生相同梯度的配置组合”之间的验证。

从源码结构看，这一测试还与本仓库的 rerun 状态机（--rerun-mode: validate_results，实现见 megatron/core/rerun_state_machine.py）结合：训练以结果校验模式运行，叠加确定性模式与固定随机种子，确保重跑/断点续训后产生的梯度与首次运行严格一致，从而让“resume 后梯度不变”这一断言也得到覆盖。

运行前的关键约束与配置前提

原文档明确列出了运行该测试必须满足的几项前提，缺一不可：

1. 必须使用旧版（pre-mcore-0.14）分布式 checkpoint 格式：比对脚本当前依赖较老的 optimizer checkpoint 结构，因此配置中必须加上--dist-ckpt-save-pre-mcore-014: true。对应参数在源码中的定义为dist_ckpt_save_pre_mcore_014: bool = False（见 megatron/training/config/training_config.py），即默认关闭、测试场景需显式开启。
2. 必须关闭随机化：dropout 等随机性在单一 global batch 或模型/数据并行分片下会产生不同的模式，破坏梯度一致性，因此配置中--attention-dropout: 0.0、--hidden-dropout: 0.0，并开启--deterministic-mode: true。
3. GPU 数量与并行度要匹配：基准配置为单卡（DP=1, TP=1, CP=1, PP=1），各并行配置在同一批卡上以相同 global batch 切分 micro-batch。
4. 奇数 ROPE base 让梯度远离 0：配置中--rotary-base: 500（注释明确说明使用奇数值是为了让梯度数值显著非零，避免比对时出现 0 vs 0 的平凡通过）。

完整配置解析：model_config.yaml

测试用例的完整参数位于 tests/functional_tests/test_cases/gpt/gpt3_mcore_reruns_resume_check_grads/model_config.yaml，通过 YAML anchor（&BASE_MODEL_ARGS）复用公共参数，各并行配置仅覆盖并行度与 batch 相关项。

通用训练设置（BASE_MODEL_ARGS）

模型结构设置

该用例配置的是一个约 4B 规模的 GPT 模型：

• --num-layers: 32、--hidden-size: 3072、--ffn-hidden-size: 8192
• --num-attention-heads: 32、--num-query-groups: 8、--kv-channels: 128，并开启--group-query-attention: true
• --seq-length: 512、--max-position-embeddings: 4096
• --position-embedding-type: rope、--rotary-percent: 1.0
• --normalization: RMSNorm、--swiglu: true
• --transformer-impl: transformer_engine、--disable-bias-linear: true
• --untie-embeddings-and-output-weights: true（词嵌入与输出层权重解绑）

梯度与优化器设置（测试核心）

--clip-grad: 1.0 --overlap-grad-reduce: true --overlap-param-gather: true --lr: 3e-4 --lr-warmup-samples: 0 # Set adam beta1 and beta2 to 0 so that we can easily check the gradients --adam-beta1: 0.0 --adam-beta2: 0.0 --adam-eps: 1e-8 --use-distributed-optimizer: true --split: 949,50,1 --no-gradient-accumulation-fusion: true

其中--adam-beta1: 0.0与--adam-beta2: 0.0是整条测试链路的核心，注释也直接点明其目的。--use-distributed-optimizer: true意味着优化器状态本身是分片存储的——这也正是比对脚本需要处理分片/去分片逻辑的原因。

Checkpoint 设置

--save-interval: 1 # 每 1 步保存一次 checkpoint --eval-interval: 1000 --ckpt-format: torch_dist --dist-ckpt-strictness: log_all --dist-ckpt-optim-fully-reshardable: true --save: ${CHECKPOINT_SAVE_PATH} --load: ${CHECKPOINT_LOAD_PATH}/model/mcore_gpt/gpt3_4b_pyt/25.03.05_bf16_rerun-enabled_v2 --async-save: true --use-persistent-ckpt-worker: true

并行配置矩阵

注意：所有并行配置的global-batch-size 恒为 2，与基准保持一致，只是 micro-batch 与并行维度不同。这样保证数学上各配置消费的是相同的数据切片（配合同一随机种子），梯度才有可比性。另外通过ENV_VARS设置了CUDA_DEVICE_MAX_CONNECTIONS: 1、NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0、NCCL_ALGO: Ring、CUBLAS_WORKSPACE_CONFIG: :4096:8等环境变量，进一步收敛非确定性来源。

比对脚本剖析：test_optimizer_grads_match.py

核心比对逻辑位于 tests/functional_tests/python_test_utils/test_optimizer_grads_match.py，它既可以作为库函数被调用，也可以作为独立 CLI 脚本运行：

python tests/functional_tests/python_test_utils/test_optimizer_grads_match.py <ckpt1> <ckpt2> [<ckpt3> ...]

其内部工作流分四步：

1. 无进程组加载分布式 checkpoint

load_dist_checkpoint_pt()使用torch.distributed.checkpoint.filesystem.FileSystemReader读取 checkpoint 元数据（read_metadata()），仅挑选 key 匹配正则optimizer的张量（默认 pattern 为r"optimizer"），构造 CPU 上的空占位张量，再以no_dist=True调用load()流式加载——整个过程不需要初始化任何分布式进程组，一个普通 Python 进程即可完成多 checkpoint 的比对。这大大降低了测试门槛。

2. 过滤出优化器动量张量

_filter_optimizer_tensors()只保留 key 以optimizer.开头且包含.exp_avg.的张量。由于 Adam 的beta1=beta2=0，这些exp_avg张量就是我们要比对的梯度。

3. 处理分片形状差异

_assert_optimizer_tensors_equal()是比对的核心。当左右两侧张量形状不一致（例如非并行 checkpoint 与 TP 分片 checkpoint）时，它需要先“对齐形状”：

• 对于row parallel linear 层（mlp.linear_fc2.weight与self_attention.linear_proj.weight），其分片方式与列并行不同，需要专门的unshard_row_parallel_state()函数还原：先按[..., tp, O, I/tp]视图展开，再 permute 与 reshape 成完整的[..., O, I]。
• 对embedding.word_embeddings.weight与.output_layer.weight等涉及 padding 变化的张量，则按最小公共维度截断比较。
• 其余张量尝试简单 reshape 对齐。

源码注释中明确标注TODO come up with a better way to determine row parallel linear layers——即当前脚本对 GPT 模型有硬编码假设，文档也特别指出：扩展到其他模型时，最关键的是重新实现“哪些层是 row parallel linear”的判断逻辑及其对应的去分片 reshape 函数。

4. 数值容差判定：基于理论误差界而非固定阈值

比对并非简单使用torch.allclose，而是实现了 arxiv 2506.09280 Theorem 5.3 的误差界模型：

• machine_epsilon_for_dtype()根据张量 dtype 返回机器精度（FP8 场景按论文建议使用 BF16 的 epsilon）。
• expected_rel_bound(l, L, C, dtype, k)给出随网络深度指数增长的相对误差上界k * (C^(L+1-l)) * eps_mch（L=32对应本用例层数，C=1.03为接近 1 的增长因子，k=4.0吸收大 O 常数）。
• check_gradient()计算相对差||g_hat - g_ref||_F / ||g_ref||_F并判定是否落在理论界内。
• 若失败，再尝试一次对张量做torch.roll移位后的对比，并辅以torch.testing.assert_close输出详细诊断信息，帮助区分“相对范数超界”与“张量内容确实不等”两类失败。

比对还包含两个兜底断言：优化器 tensor 的 key 集合必须完全一致；所有被比较张量中至少要存在一个非零张量（避免 0 vs 0 的平凡通过，some_non_zero检查）。

测试注册与 CI 集成：recipes 与 golden values

该用例通过 tests/test_utils/recipes/h100/gpt-grads.yaml 注册到 CI（README 中提到的tests/functional_tests/test_utils/recipes/gpt-grads.yaml在仓库当前结构中对应位于tests/test_utils/recipes/h100/下的同名文件）。recipe 的关键信息：

• spec.model: gpt、platforms: dgx_h100、单节点 8 卡（nodes: 1、gpus: 8）；
• script段将TESTING_PARAMS_PATH指向./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml，将GOLDEN_VALUES_PATH指向同目录下的 golden values JSON，再调用bash ./tests/functional_tests/shell_test_utils/run_ci_test.sh执行；
• products段将test_case: [gpt3_mcore_reruns_resume_check_grads]、environment: [dev]、scope: [mr]、platforms: [dgx_h100]绑定，即该用例在 dev 环境的 MR 级别 CI 中运行，且注释强调该测试很昂贵，硬编码N_REPEAT=1。

同目录的 golden_values_dev_dgx_h100.json 记录了基准运行在第 1 步的期望观测值，包括：

• lm loss：29.29548
• num-zeros：152955200.0（梯度零值计数，配合--log-num-zeros-in-grad输出）
• mem-allocated-bytes/mem-max-allocated-bytes：约 67.28 GB（与 32 层、hidden 3072 的 4B 级模型在 H100 上的显存占用相符）

这些 golden values 用于验证基准配置本身没有回归，而跨并行配置的梯度一致性则由上述比对脚本负责。

如何扩展：添加新的梯度一致性测试

原文档明确给出了扩展指引，结合源码可总结为以下步骤：

1. 复制测试目录：拷贝tests/functional_tests/test_cases/gpt/gpt3_mcore_reruns_resume_check_grads/并改名，如gpt3_mcore_check_grads_<新场景>。
2. 修改 model_config.yaml：在BASE_MODEL_ARGS之上，通过 YAML anchor 继承公共参数，再为每个新配置新增MODEL_ARGS_*块，覆盖并行度、micro-batch 等差异项；若测试新模型，还需同步调整--num-layers等结构参数，使expected_rel_bound的L假设与模型深度一致。
3. 适配比对脚本的层类型判断：如果是非 GPT 模型，必须修改_assert_optimizer_tensors_equal()中判定 row parallel linear 层的逻辑（当前硬编码为mlp.linear_fc2.weight与self_attention.linear_proj.weight），并实现对应的unshard_row_parallel_state变体。
4. 注册测试：参照gpt-grads.yaml的写法在 recipes 目录下新增或扩展 recipe，绑定test_case、environment、platforms与scope。
5. 更新 golden values：先跑通基准配置，将输出的 loss、num-zeros、显存等观测值写入新目录的 golden values JSON。

需要强调的是，这套方法不限于模型并行配置之间的比对——任何“在数学上应当产生相同梯度”的配置组合（例如不同的重计算策略、不同的梯度归并顺序、关闭/开启某些融合算子等）都可以套用同样的套路：beta1=beta2=0固化梯度 → checkpoint 提取 → 去分片对齐 → 理论误差界判定。

总结

gpt3_mcore_reruns_resume_check_grads测试用例为 Megatron-LM 提供了一条“以 checkpoint 为媒介、以优化器状态为存储”的梯度一致性验证通路。它借助 Adam 动量清零这一巧妙手段规避了梯度接口差异，借助 torch.distributed.checkpoint 的无进程组加载降低了比对门槛，并借助基于误差传播理论的自适应容差取代了拍脑袋的固定阈值。无论是用于并行实现回归验证，还是用于探索新的训练配置组合，这一模式都值得直接复用与扩展。

【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM