Megatron-LM 确定性训练开发者参考：位级可复现契约、确定性模式与内核测试门禁

Megatron-LM 确定性训练开发者参考位级可复现契约、确定性模式与内核测试门禁【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM本指南面向 Megatron-LMMCore的开发者与代码评审者系统讲解仓库中确定性训练deterministic training的完整技术体系从位级可复现的契约定义到--deterministic-mode的启动参数与环境变量校验再到操作目录op catalog中每个内核的确定性分支选择以及每个内核必须有 bit-exact 测试的注册表与 PR 门禁机制。阅读本文后你将理解确定性模式的底层原理、能够解释 15% 左右的性能代价从何而来并掌握为新增或修改的内核编写、注册 bit-exact 测试并本地运行门禁检查的完整方法。本文主体整理自仓库内的确定性开发参考文档即 docs/developer/determinism/README.md 及其四个子页status.md、op-catalog.md、testing.md、glossary.md面向普通使用者的配置说明见 确定性训练用户指南本文不再重复其内容。一、确定性契约什么是位级可复现位级确定性bit-exact determinism的定义是在配置、数据、随机种子、软件与硬件完全一致的前提下两次独立运行产生完全相同的结果。更严格地说文档将其表述为bitwise determinism在输入数据及其顺序、模型配方与配置、并行布局TP/PP/DP/EP/CP/VPP、容器镜像与库版本Megatron-Core、CUDA、cuDNN、NCCL、Transformer Engine、PyTorch、驱动、NCCL 设置以及硬件类型与拓扑全部固定的情况下每次运行都产生逐位相同的数值结果。需要特别强调的是这份契约约束的是数值结果而不是计算图本身。内核的调度与执行顺序可以变化只要所有浮点归约reduction都以相同的顺序发生即可。这正是整个确定性工程的核心难点浮点加法不满足结合律归约顺序不同会带来舍入差异而几乎所有的训练不确定性问题都源于归约顺序不可控。1.1 验证方式把 loss 当作模型状态的校验和验证确定性的标准做法是在两次相互独立的运行中追踪训练指标loss、梯度范数、参数范数、num-zeros要求其曲线逐位一致。文档给出的直觉解释是——loss 是对采样 logits 的归约等价于模型状态的校验和checksum任何一处的数值不一致都会在几步之内传播为 loss 与梯度曲线上的逐位差异。一个重要的实操细节控制台日志只打印有限精度严格比较时必须使用全精度序列化指标例如 TensorBoard event 文件而不是打印出来的数值。1.2 契约相关的关键术语术语含义Deterministic mode以--deterministic-mode运行由 megatron/training/determinism.py 完成参数校验与环境变量默认值填充并调用torch.use_deterministic_algorithms(True)库代码通过config.deterministic_mode或torch.are_deterministic_algorithms_enabled()选择确定性分支Default mode不带--deterministic-mode的运行Bit-exact / bitwise identical被比较的张量或序列化指标字节级一致见上方验证Reproducible在相同条件下同一分配、缓存、环境得到相同结果保证弱于契约Cross-allocation在独立分配的资源上比较运行新的调度分配通常对应不同的物理节点与网络环这是确定性测试要达到的目标与门槛Collision-free唯一索引写入索引互不重复的写入scatter、index_put_(accumulateFalse)不发生浮点累加顺序不影响结果天然确定Fail closed确定性模式遇到无法保证的功能时在校验阶段直接拒绝该配置而不是静默运行二、确定性模式--deterministic-mode的三层动作--deterministic-mode的入口实现位于 megatron/training/determinism.py。从源码结构看该模块对外暴露两个函数apply_determinism_env与apply_determinism_to_args后者由validate_args调用。apply_determinism_to_args是幂等的按顺序完成三件事校验配置参数断言cross_entropy_loss_fusion与tp_comm_overlap必须为False见源码中的ARG_VALUES_REQUIRED_FOR_DETERMINISM。这是纯校验绝不静默改回——不兼容选项会以显式错误被拒绝用户必须自行去掉这些 flag以确保实际运行与提交的配置一致。校验环境变量并填充规范默认值调用apply_determinism_env处理os.environ。开启 torch 确定性算法torch.use_deterministic_algorithms(True)置于所有断言通过之后。源码 docstring 特别强调这些环境变量被各自的库在首次使用时捕获NCCL 在通信器初始化、cuBLAS 在句柄创建、Transformer Engine 在首次 attention 前向时因此apply_determinism_to_args必须在validate_args的早期执行先于任何 cuBLAS / Transformer Engine 内核调用。2.1 环境变量校验规则与默认值每个变量可以由启动器预先设置也可以保持未设置。已设置时其值必须是通过校验的确定性取值否则断言失败fail hard未设置时apply_determinism_env会填充规范默认值MAMBA_DETERMINISTIC与CAUSAL_CONV1D_DETERMINISTIC除外——这两个由各自内核从torch.are_deterministic_algorithms_enabled()自动检测。下表来自 确定性训练用户指南 并对应 determinism.py 中的常量变量接受值或未设置未设置时填充的默认值原因NCCL_ALGO{Ring, CollnetDirect, CollnetChain, ^NVLS}的子集Ring保守默认值——Ring的归约顺序由拓扑固定在所有受支持的 NCCL 版本上跨运行 bit-exactNVTE_ALLOW_NONDETERMINISTIC_ALGO00强制 Transformer Engine 只使用确定性算法CUBLAS_WORKSPACE_CONFIG:4096:8或:16:8:4096:8确定性 cuBLAS workspaceNVIDIA 文档标注两种尺寸均可复现:4096:8更快:16:8更省显存MAMBA_DETERMINISTIC任意以1开头的字符串无——SSM 自动检测未设置时 Mamba SSM 自动跟随torch.are_deterministic_algorithms_enabled()只有显式的非确定性覆盖会被拒绝CAUSAL_CONV1D_DETERMINISTIC任意以1开头的字符串无——内核自动检测causal_conv1d ≥ 1.6.0 在未设置时自动跟随确定性开关通过 workspace 而非atomicAdd归约卷积权重/偏置梯度Mamba 与 GDP mixer 在缺少它时会拒绝确定性运行2.2 为什么排除Tree算法NCCL_ALGO的接受集合在源码中是一个显式白名单ACCEPTED_NCCL_ALGO_TOKENS frozenset({Ring, CollnetDirect, CollnetChain, ^NVLS})Tree被有意排除。源码注释解释了原因Tree 的节点内链式归约顺序不受用户控制而跨节点的树形拓扑在没有固定拓扑文件的情况下可能随运行变化因此无法担保跨运行跨栈的 bit-exact。^NVLS被接受因为禁用 NVLS在暴露该特性的硬件上是合法的用户选择此时由用户负责确保 NCCL 回退到的算法在自己的环境中是确定性的。2.3 配置要求fail closed 的旗标Flag--deterministic-mode下的行为--cross-entropy-loss-fusion必须关闭——被断言拒绝融合 CE 内核不确定需自行去掉该 flag--tp-comm-overlap必须关闭——被断言拒绝重叠路径不 bit-exact需自行去掉该 flagtorch.use_deterministic_algorithms被置为True需要指出的是--use-flash-attn不会被拒绝Transformer Engine 的 flash-attention 后端在NVTE_ALLOW_NONDETERMINISTIC_ALGO0时是确定性的包括确定性的 FlashAttention 反向并且由 bit-exact 正确性测试套件覆盖源码注释中对此有明确说明。三、操作目录Op Catalog哪些操作有确定性路径大多数训练步骤中的操作并不需要特殊处理它们天然确定逐元素操作elementwise ops在固定的 cuBLAS workspace 下的 GEMM按 rank 索引的集合通信all-gather、all-to-all、broadcast稳定排序stable sorts唯一索引写入unique-index writes操作目录op-catalog.md只列出那些需要做出选择的操作分为两类有确定性代码路径的操作与尚无确定性支持的操作。项目目标是缩小第二类、并让第一类更快。3.1 有确定性路径的操作下表完整整理自 op-catalog.md。确定性分支由torch.are_deterministic_algorithms_enabled()或config.deterministic_mode选择默认模式路径保留在另一分支操作位置确定性路径默认路径MoE token 反置换combinemegatron/core/transformer/moe/moe_utils.pyindex_add_——在 torch 确定性算法下确定且 CUDA-graph 安全scatter_add_原子累加MoE 路由映射与概率megatron/core/transformer/moe/moe_utils.pyindex_put_(accumulateFalse)按行写入就地out-of-placescatterVocab 并行 embeddingmegatron/core/tensor_parallel/layers.py直接索引weight[idx]确定性反向F.embedding非确定性原子反向Gated-delta-net 内核megatron/core/ssm/gated_delta_net.pytorchchunk_gated_delta_ruleFLA 融合内核Gated-delta-net 因果卷积megatron/core/ssm/gated_delta_net/F.conv1d外加转置FLAcausal_conv1dMamba/SSM Triton 操作megatron/core/ssm/ops/common/determinism.py固定一份 autotune 配置 零初始化的分块 workspace用有序sum归约基于计时的 autotune、未初始化 workspaceMamba/GDP 因果卷积megatron/core/ssm/causal_conv1d.pycausal_conv1d ≥ 1.6.0——权重与偏置梯度使用 per-block workspace以有序sum归约atomicAdd累加每次 launch 顺序可变Transformer Engine 注意力megatron/core/extensions/transformer_engine.py要求NVTE_ALLOW_NONDETERMINISTIC_ALGO0此时 TE 只选择支持确定性执行的 backend含确定性 FlashAttention 反向TE 自由选择包括原子累加注意力的反向推理 DP 调度与 RL rollout 顺序megatron/core/inference/engines/dynamic_engine.py、megatron/rl/rl_utils.py按稳定 key 排序完成顺序这里有两处容易混淆两行 conv 是不同内核。Mamba 与 GDP 调用 Dao-AILab 的causal_conv1d自带确定性归约而 gated-delta-net 绑定的是 FLA 的版本不带确定性归约因此回退到F.conv1d。作为源码佐证MoE token 反置换的确定性分支可以在 moe_utils.py#L604-L617 中直接看到torch.are_deterministic_algorithms_enabled()为真时使用output_tokens.index_add_(0, sorted_indices, permuted_tokens)并注释说明index_add在确定性算法开启时是确定的且兼容 CUDA graph不同于scatter_add默认分支才是scatter_add_。3.2 尚无确定性支持的操作操作或特性强制位置或观察位置状态融合交叉熵损失--cross-entropy-loss-fusion被--deterministic-mode拒绝megatron/training/determinism.py融合内核不确定是否存在确定性变体仍是开放问题。在此之前框架使用原生 vocab-parallel 路径TP 通信重叠--tp-comm-overlap被--deterministic-mode拒绝重叠后的集合通信顺序不可复现gated-delta-net 中的 packed sequencethdmegatron/core/ssm/gated_delta_net.py 中的断言尚无确定性的 packed-sequence SSM 路径跨分配cross-allocation的浮点集合通信TP all-reduce、DP 梯度 reduce-scatter已知开放缺口NCCL_ALGORing固定了算法但未固定某个分配allocation实际获得的物理环。仅靠该环境变量无法保证这些归约在不同分配之间 bit-exact同一分配内的重复运行、或拓扑完全一致的分配间运行仍然 bit-exact3.3 性能代价约 15% 的步时间与默认模式相比确定性路径大约增加15%的步时间因模型与精度而异重度依赖 MoE 的模型代价更高。已测案例范围从大型稠密模型约4%到混合 MoE 模型约17%。文档点名的性能热点包括确定性的 MoE scatter 与 unpermute 路径排序后的 router top-k注意力反向attention backwardGrouped-GEMM 的权重梯度wgrad项目的目标是让代价降到10% 以下冲刺目标接近5%从而可以在生产运行中长期开启确定性模式。热点清单与优化进展记录在 op catalog 与仓库跟踪的 issue #5785 中。任何改动上面表格任一行都需要 bit-exact 测试注册在tests/unit_tests/determinism/kernels/manifest.py以及确定性与默认性能的对比见下节。四、内核确定性测试bit-exact 测试套件测试页testing.md提出的要求非常明确Megatron 分发的每一个 GPU 内核都必须有 bit-exact 确定性测试且每个新增或修改内核的 PR 都必须同时新增或更新该测试。4.1 什么算一个内核任何 Megatron 自身启动或选择 GPU 内核、且其数值结果可能取决于调度顺序的代码路径Tritontriton.jit内核以及 TileLang / cuTile 内核jit_fuser/torch.compile融合函数C/CUDA 扩展load_inline、CUDAExtension、.cu源文件Transformer Engine 与外部库的算法选择分发grouped GEMM、注意力后端、causal_conv1d、mamba_ssm、FLA、DeepEP、cuTile、FlashInfer、apex multi-tensor 内核具有非确定性累加的 torch 操作scatter_add_、index_add_、index_put_(accumulateTrue)、bincount、embedding 反向判定规则在 manifest.py 中以代码形式固化KERNEL_DIRECTORIESmegatron/core/fusions与megatron/core/ssm/ops下的每个.py文件都必须注册和KERNEL_CONTENT_PATTERNS正则模式匹配triton.jit、torch.compile、CUDAExtension(、.scatter_add_(、F.embedding(、tex.rmsnorm_fwd(、causal_conv1d_fn(、buffer.dispatch(等。需要注意外部库模式匹配的是调用点call site从不匹配裸 import——只导入类或做isinstance检查的模块不算内核载体。而像 Mamba mixer、gated delta product、RoPE 分发、FP8 master-weight 转换这类选择或调用外部内核但不自己定义内核的模块以kinddispatch注册它们由所调用内核的回放测试覆盖在条目notes中指名并在存在模块级测试时由模块级测试覆盖。4.2 测试的存放位置层级位置断言内容内核级tests/unit_tests/determinism/kernels/test_*.py每个内核族一个文件。同一输入上多次运行内核所有输出张量与梯度必须字节一致harness.assert_replays_bit_exact、assert_module_replays_bit_exactharness.bytes_equal比较比特模式因此带符号零与 NaN payload 也必须一致模块/模型级tests/unit_tests/determinism/correctness/GPT、TransformerBlock、HybridModel 以及 FP8/FP4 recipeBitExactRunner覆盖张量并行、专家并行、全分片数据并行、流水线并行与虚拟流水线并行等并行单元FP8/FP4tensorwise、delayed、mxfp8、nvfp4见 test_fp8_determinism.pyBlackwell 专属 recipe 在 Hopper 上按能力跳过端到端带--deterministic-mode的功能测试loss 与num-zeros按记录精度与 golden 值比较旧版 golden 使用五位小数4.3 注册表与 PR 门禁三层强制机制仓库不变量test_manifest.py纯 CPU跑在 unit-test 桶megatron/下每个内核载体文件都已注册、每个注册路径真实存在、每个条目都有测试或豁免理由。往新文件里加一个 Triton 内核而不注册单元测试直接失败。PR 门禁tools/check_kernel_determinism_coverage.py由.github/workflows/cicd-main.yml中的lintingjob 在 PR push 时运行被修改的内核载体文件必须已注册不可覆盖被修改的已注册内核源文件必须同时修改其至少一个确定性测试。当改动不影响数值纯注释、重构时可用determinism-exemptPR 标签豁免此时检查仅记录豁免日志而不失败。评审PR 模板复选框与/claude review提示都会显式要求测试。本地对照main运行门禁python3 tools/check_kernel_determinism_coverage.py --base-ref origin/main内核测试桶运行在 H100 单元测试 recipetests/test_utils/recipes/h100/unit-tests.yaml中。由于硬件相关的调度正是这些测试要捕获的目标在 GB200/GB300 级 runner 上也运行该桶GB200 单元测试由launch_on_gb200marker 选择是已跟踪的后续工作在此之前需要在 Blackwell 硬件上手工复现发现方法见 status.md。4.4 如何新增一个内核测试五步流程放置测试把测试放入与内核族对应的tests/unit_tests/determinism/kernels/test_*.py模块或新建包__init__在 import 时钉住确定性环境。使用 harnessfrom tests.unit_tests.determinism.kernels.harness import ( assert_replays_bit_exact, assert_module_replays_bit_exact, deterministic_algorithms, seeded, ) seeded() x torch.randn(16384, 2048, devicecuda, dtypetorch.bfloat16, requires_gradTrue) assert_replays_bit_exact(my_kernel, (x,), replays3, contentionTrue, whatmy_kernel)contentionTrue会在旁路 GEMM 压力下执行回放让顺序相关的归约暴露出来。输入尺寸比回放次数更重要两个块的归约可能碰巧确定要把尺寸开到许多 CTA 竞争同一批输出harness 中的CONTENTION_TOKENS 4096就是为此设计。根据torch.are_deterministic_algorithms_enabled()选择分支的内核应使用deterministic_algorithms(...)在两个分支上都测试。消费 RNG 的内核dropout使用restore_rngTrue。添加负对照当默认路径已知存在竞态原子操作时用count_differing_replays(...) 0添加一个充分大的负对照确保严格断言确实敏感参考test_moe_kernels.py。注册内核在manifest.py中登记源文件、测试文件、内核类型triton/torch.compile/cuda-ext/te-wrapper/torch-op/external-lib/tilelang/dispatch以及关于非确定性机制与确定性分支如何被选择的备注。每个KernelEntry要么有非空tests元组要么有显式exempt_reason例如需要 NVLink peer group 的多 rank NVLS 集合通信、或没有计算内核的分配器。豁免是可见的覆盖债务而不是沉默。在 GPU 节点上运行uv run python -m torch.distributed.run --nproc-per-node 8 -m pytest -q \ tests/unit_tests/determinism/kernels/test_my_family.py \ tests/unit_tests/determinism/kernels/test_manifest.py4.5 测试 harness 的字节级严格性harness.py 的实现细节值得注意bytes_equal比torch.equal更严格——0.0与-0.0视为不同NaN 只有在 payload 相同时才匹配比较的是逻辑内容按逻辑顺序读取的uint8视图因此返回不同内存布局但内容相同的张量不算失败。_clone_preserving_layout还会保留expand产生的 stride-0 视图避免克隆后内核走入与生产不同的特化路径。五、模块级与端到端验证比内核回放更广的覆盖模块级 bit-exact 套件tests/unit_tests/determinism/correctness/在恢复 RNG 状态的前提下两次运行模型或模块块断言输出与梯度 bit 一致。覆盖范围包括 GPTModel、TransformerBlock、HybridModel张量并行、专家并行、全分片数据并行、流水线并行与虚拟流水线并行FP8 与 FP4 recipe以及用于暴露潜在顺序竞态的调度压力测试scheduling stressors。性能门禁tests/performance_tests/shell_test_utils/determinism/在 Nsight Systems 下分别以确定性与默认模式运行小型 recipe输出各 NVTX 区间range的排行榜当确定性步时间超过文档化阈值时失败。实际脚本为 run_nsys_breakdown.sh可包装任意训练入口如pretrain_hybrid.py --profile与 print_nsys_leaderboard.py将两份 CSV 拼成并排表格CI 调用见 tests/test_utils/recipes/h100/determinism-perf.yaml。端到端验证跨两次独立运行比较全精度训练指标曲线见术语表中Verification说明。将已入库的覆盖扩展到生产规模架构是路线图上的待办项。六、常用并行与模型缩写速查开发者参考中大量使用缩写glossary.md 给出了权威定义摘录如下并行与基础设施MCore本仓库的模型并行训练库、DP数据并行、TP张量并行、PP流水线并行、VPP虚拟流水线并行、EP专家并行、CP上下文并行、A2Aall-to-all 集合通信MoE token 分发与合并是按 rank 索引的置换而非浮点归约、TETransformer Engine、wgrad/dgrad线性层反向中的权重梯度/输入梯度。模型相关MoE专家混合、MLA多潜在注意力DeepSeek 家族的低秩潜在 qkv 投影、DSV3/DSV4DeepSeek-V3-/V4 风格配置DSV3 组合 MLA 与细粒度 MoEDSV4 额外使用 DSA 稀疏注意力、DSADeepSeek 稀疏注意力、MTP多 token 预测、SSM状态空间模型层Mamba 家族、GDNGated delta net一种 SSM 变体见 megatron/core/ssm/gated_delta_net.py。七、路线图与总结确定性的路线图在仓库跟踪的 issue #5785 中动态维护当前工作流集中在两条线上缩小尚无确定性支持的操作集合例如为融合交叉熵与 packed-sequence GDN 寻找确定性变体、研究跨分配集合通信的确定性方案以及降低确定性路径的性能代价目标从当前约 15% 降到 10% 以下、冲刺 5%。总结这份开发者参考的核心脉络契约bit-exact determinism 要求两次相同运行逐位一致本质是保证所有浮点归约顺序一致模式--deterministic-mode通过 determinism.py 校验参数与环境变量fail closed、填充规范默认值并开启 torch 确定性算法目录op-catalog.md 逐条记录了每个需要选择的操作的确定性分支与不支持项门禁testing.md 与 manifest.py 通过内核必须注册 改内核必改测试的机制把 bit-exact 验证固化为仓库不变量与 PR 门禁check_kernel_determinism_coverage.py 可在本地一键复现。对于想要深入代码的读者推荐按以下顺序研读先看 megatron/training/determinism.py 理解三层动作再对照 op-catalog.md 在 moe_utils.py、layers.py、transformer_engine.py 中逐一印证确定性分支最后通读 manifest.py 与 harness.py即可完整掌握为内核编写 bit-exact 测试的工程方法。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考