Megatron-LM CUDA Graph 完整指南：local / transformer_engine / full_iteration 三种实现的配置、原理与迁移

Megatron-LM CUDA Graph 完整指南local / transformer_engine / full_iteration 三种实现的配置、原理与迁移【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LMCUDA Graphs 通过一次性录制 GPU 算子执行序列、后续迭代直接回放replay大幅削减 kernel 启动开销。本篇技术指南以 Megatron-LM 的 CUDA Graph 特性文档 为主线系统讲解该仓库中由--cuda-graph-impl统一控制的三种实现local、transformer_engine、full_iteration的适用场景、配置方法、底层机制与旧参数迁移路径并结合megatron/training/arguments.py、megatron/core/transformer/cuda_graph_config.py、megatron/core/full_cuda_graph.py等源码给出可复制、可运行的实战示例帮助你在稠密模型、MoE 模型、推理/RL 等不同场景下正确启用 CUDA Graphs。概述三个正交的控制开关Megatron-LM 用三个相互独立的命令行参数共同决定 CUDA graph 行为它们各司其职Flag取值作用--cuda-graph-implnone/local/transformer_engine/full_iteration选择捕获后端或整体策略--cuda-graph-modulesattn/mlp/moe/moe_router/moe_preprocess/mamba训练期按层级的捕获覆盖范围多值参数仅对local与transformer_engine有意义--inference-cuda-graph-scopenone/layer/block推理期 CUDA graph 的粒度仅local支持非none取值三个开关可以组合出如下受支持的搭配--cuda-graph-impl后端训练期捕获推理期捕获none—关闭关闭localMCoreCudaGraphManager按层捕获由--cuda-graph-modules控制layer默认或block由--inference-cuda-graph-scope控制transformer_engineTEmake_graphed_callables()按层捕获由--cuda-graph-modules控制不支持只能nonefull_iterationMCoreFullCudaGraphWrapper每个训练迭代一张图--cuda-graph-modules必须为空不支持只能none从源码层面看这三个开关在 arguments.py 中完成解析与合法性校验--cuda-graph-modules被定义为nargs的多值参数默认[]并在validate_args阶段通过normalize_cuda_graph_modules规范化为CudaGraphModule枚举--inference-cuda-graph-scope则通过normalize_inference_cuda_graph_scope规范化且当未显式给出时按--cuda-graph-impl推导默认值local默认layer其余默认none。非法组合如full_iteration搭配非空--cuda-graph-modules、或推理粒度超出ALLOWED_INFERENCE_SCOPES白名单会在启动阶段直接报错而不是在训练中途失败。CUDA Graph 实现一本地实现--cuda-graph-impl locallocal实现使用 MCore 内置的CudaGraphManager核心代码位于 megatron/core/transformer/cuda_graphs.py。在训练期它工作在**按层per-layer**模式不设置--cuda-graph-modules时捕获整个 Transformer 层指定模块时只捕获所选子区域。在推理期local支持把图挂在层边界layer或外层 block 边界block由--inference-cuda-graph-scope控制。对应的枚举定义见 megatron/core/transformer/enums.py 中的InferenceCudaGraphScopenone表示推理走 eager 模式layer表示图归模块/层所有如TransformerLayer或MambaLayerblock表示图归外层 block 所有如TransformerBlock或HybridBlock。在运维层面这条路径与 MCore 训练/推理流程深度集成使用者只需通过配置开关选择模式无需额外处理静态输入缓冲区可图化的模块会自动创建并持有各自的CudaGraphManager实例现有训练调度schedule自动驱动 warmup / 捕获 / 回放没有单独的自定义训练循环接入 API也无需手动处理静态输入缓冲。从源码可以看到local模式的接入点遍布训练与推理链路例如 megatron/core/pipeline_parallel/schedules.py 的多个调度器在config.cuda_graph_impl local时执行图化逻辑dynamic_engine.py 与 text_generation_controller.py 也仅在cuda_graph_impl local时才启用推理 CUDA graph。这正对应文档中推理 CUDA graph 目前仅local支持的说明。用法--cuda-graph-impl local--cuda-graph-modules可选项模块取值捕获内容空 / 不设置整个 Transformer 层默认attnTransformerLayer._forward_attention()mlp稠密层的TransformerLayer._forward_mlp()moeMoE 层的TransformerLayer._forward_mlp()仅 drop-and-pad 模式moe_routerMoE router 共享专家若未与 EP 通信重叠moe_preprocessMoELayer.preprocess()—— 必须与moe_router搭配使用mambaMamba SSM 层示例 —— MoE 模型仅捕获注意力与 router--cuda-graph-impl local \ # 可选限制被捕获的模块默认捕获整个层但对 MoE 动态形状无效 --cuda-graph-modules attn moe_router moe_preprocess上述每个取值的语义在 arguments.py 的 help 文本中有逐条定义例如moe_router捕获TransformerLayer._forward_mlp()直到MoELayer.router()的运算含未被 EP 通信重叠的共享专家moe_preprocess捕获MoELayer.preprocess()且强制要求与moe_router同时出现。此外mamba_layer.py 中可见 Mamba 层对cuda_graph_impl local有显式断言assert self.config.cuda_graph_impl local印证了 Mamba SSM 层目前只走local路径。CUDA Graph 实现二Transformer Engine 实现--cuda-graph-impl transformer_engine该实现走 Transformer Engine 的make_graphed_callables()路径。在 Megatron-LM 的 CLI 中它的训练粒度与local完全相同不设置--cuda-graph-modules捕获整个 Transformer 层指定模块则只捕获所选子区域。与local的主要差异在于后端实现与功能兼容性与local不同此路径不支持推理期 CUDA graph相比localTE 路径通过make_graphed_callables()暴露了更通用、更自包含的 API让用户在自定义训练循环中拥有更大的灵活性和控制力。灵活性提升的代价是手动设置更多训练循环必须实例化TECudaGraphHelper训练循环必须在正确的时点调用create_cudagraphs()与cuda_graph_set_manual_hooks()等辅助方法。Megatron-LM 自带的标准训练循环已在 megatron/training/training.py 中接入这些调用TECudaGraphHelper的实例化与create_cudagraphs()、cuda_graph_set_manual_hooks()的调用点位于训练主循环内但自定义训练脚本必须自行完成同样的接线工作。用法--cuda-graph-impl transformer_engine \ --cuda-graph-modules attn moe_router moe_preprocess训练期--cuda-graph-modules的可选值与local完全一致省略该 flag 时默认同样是捕获整个训练层。CUDA Graph 实现三全迭代训练 CUDA Graph--cuda-graph-impl full_iterationfull_iteration将整个训练迭代排除优化器捕获为一张 CUDA graph。同一个包装器megatron/core/full_cuda_graph.py 中的FullCudaGraphWrapper也被用于训练循环的验证/评估阶段forward-only 模式可提供最大的训练/验证延迟削减。FullCudaGraphWrapper内部维护每个阶段training/validation的 CUDA graph 对象与迭代计数运行--cuda-graph-warmup-steps次 warmup 后在curr_iteration cuda_graph_warmup_steps时创建torch.cuda.CUDAGraph并录制 forward/backward 结果后续迭代直接调用graph.replay()同时支持注册生成器状态register_generator_state并在阶段切换时清理已录制图详见 full_cuda_graph.py。该实现不创建推理 CUDA graph。推理场景请改用--cuda-graph-impl local --inference-cuda-graph-scope layer|block硬性要求必须设置--no-check-for-nan-in-loss-and-gradNaN 检查涉及 CPU-GPU 同步无法在 CUDA graph 内部执行。arguments.py 中对此有显式断言--no-check-for-nan-in-loss-and-grad should be set with --cuda-graph-implfull_iteration for training。必须省略或留空--cuda-graph-modules当整个迭代被捕获为一张图时按模块选择没有意义validate_args会强制--cuda-graph-modules为空--cuda-graph-modules must be empty when --cuda-graph-implfull_iteration。示例--cuda-graph-impl full_iteration \ --no-check-for-nan-in-loss-and-grad常见配置示例稠密模型训练三种实现都适用于稠密模型# 按层local --cuda-graph-impl local # 等价写法--cuda-graph-impl local --cuda-graph-modules attn mlp # 按层TE --cuda-graph-impl transformer_engine # 等价写法--cuda-graph-impl transformer_engine --cuda-graph-modules attn mlp # 全迭代 --cuda-graph-impl full_iteration \ --no-check-for-nan-in-loss-and-gradMoE 模型训练MoE 专家调度涉及动态形状dynamic shapes无法被捕获。此时用--cuda-graph-modules只捕获静态部分注意力、router、preprocess让专家计算保持 eager 模式。以下用transformer_engine举例local用法相同--cuda-graph-impl transformer_engine \ --cuda-graph-modules attn moe_router moe_preprocessMegatron-LM 的 MoE 文档也印证了这一约束megatron/core/transformer/moe/README.md 指出当--moe-expert-capacity-factor与--moe-pad-expert-input-to-capacity均未设置时MoE 层动态形状导致不可捕获此时可仅对注意力层启用 CUDA graph--cuda-graph-modulesattnMoE 层保持原样。注意--cuda-graph-modules moe_router moe_preprocess的组合存在使用前提transformer_layer.py 中对相关限制有明确提示配置时请留意模型是否符合条件。使用 paged stash 解锁 MoE 全迭代图paged stash当前仅在dev分支可用参见 docs/user-guide/features/paged_stash.md可以让专家调度的形状变为静态通过--moe-expert-rank-capacity-factor预分配大小从而使 MoE 模型也能使用全迭代 CUDA graph--cuda-graph-impl full_iteration \ --no-check-for-nan-in-loss-and-grad \ --moe-flex-dispatcher-backend hybridep \ --use-transformer-engine-op-fuser \ --moe-expert-rank-capacity-factor float \ --moe-paged-stash补充说明--cuda-graph-warmup-steps默认 3控制捕获 CUDA graph 前运行的 warmup 步数。不建议设为 0部分算子依赖前几次迭代完成惰性初始化或自动调优autotuning过早捕获可能生成错误或次优的图。从源码看该默认值对应FullCudaGraphWrapper的cuda_graph_warmup_steps参数full_cuda_graph.py。推理 CUDA graphserving 或 RL rollout目前必须使用--cuda-graph-impl local。搭配local使用--inference-cuda-graph-scope layer|block其余所有实现必须将--inference-cuda-graph-scope设为none即推理走 eager 模式。这一约束同时被 cuda_graph_config.py 中的ALLOWED_INFERENCE_SCOPES白名单强制校验。在 RL强化学习场景中相关参数还有联动约束例如--rl-persist-cuda-graphs与--rl-training-cuda-graphs都要求--cuda-graph-impl非none--rl-kv-cache-management-modeoffload又要求--rl-persist-cuda-graphs开启arguments.py。配置 RL 工作负载时需一并核对。迁移指南从旧参数到新参数旧式配置包括--enable-cuda-graph、--external-cuda-graph、已更名的--cuda-graph-scope现为--cuda-graph-modules以及full_iteration、full_iteration_inference等废弃模块值仍然被接受并在运行时自动迁移但官方鼓励将配置更新为新形式。迁移逻辑集中在 cuda_graph_config.py 中normalize_cuda_graph_modules将废弃字符串如full_iteration、full_iteration_inference识别为需迁移项validate_deprecated_cuda_graph_modules_migration_inputs负责拒绝新旧混用的歧义输入get_deprecated_cuda_graph_modules_migration返回实际迁移目标旧布尔开关的迁移则在 arguments.py 中完成--enable-cuda-graph→cuda_graph_impllocal--external-cuda-graph→cuda_graph_impltransformer_engine且二者与--cuda-graph-impl并用会直接报错。对照迁移表旧命令新命令--enable-cuda-graph--cuda-graph-impl local--external-cuda-graph--cuda-graph-impl transformer_engine--cuda-graph-scope modules--cuda-graph-modules modules--cuda-graph-impl local --cuda-graph-scope full_iteration--cuda-graph-impl full_iteration--cuda-graph-impl local --cuda-graph-scope full_iteration_inference--cuda-graph-impl local --inference-cuda-graph-scope block--cuda-graph-impl local --cuda-graph-scope attn moe_router moe_preprocess full_iteration_inference--cuda-graph-impl local --cuda-graph-modules attn moe_router moe_preprocess --inference-cuda-graph-scope block迁移提示--cuda-graph-scope参数在 arguments.py 中仍保留解析入口但已被argparse.SUPPRESS隐藏并重定向到cuda_graph_scope_deprecated字段任何新代码都不应再依赖它CudaGraphScope枚举也被标记为废弃enums.py仅保留用于旧 checkpoint 的反序列化兼容。另外历史写法--cuda-graph-modules full同样被接受并等价转换为空列表 捕获整个层迁移日志会在 rank 0 打印提示。选型建议速览你的场景推荐配置稠密模型训练追求最大延迟削减--cuda-graph-impl full_iteration --no-check-for-nan-in-loss-and-grad稠密模型训练需要精细控制捕获范围--cuda-graph-impl local或transformer_engine 按需--cuda-graph-modules attn mlpMoE 模型训练动态形状--cuda-graph-impl transformer_engine --cuda-graph-modules attn moe_router moe_preprocess或local等价写法MoE 模型 paged stashdev分支--cuda-graph-impl full_iteration--moe-paged-stash等参数组合推理 / RL rollout--cuda-graph-impl local --inference-cuda-graph-scope layer|block自定义训练循环、需要最大 API 灵活性--cuda-graph-impl transformer_engine 手动调用TECudaGraphHelper最后提醒本仓库文档引用 NVIDIA 官方《Transformer Engine and Megatron-LM CUDA Graph Support》一文作为概念背景参考但该文中的部分示例仍使用--enable-cuda-graph或--cuda-graph-scope full_iteration等旧 flag在当前仓库中请一律使用上文所述的--cuda-graph-impl local|transformer_engine|full_iteration新式配置。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考