NVFuser 融合代码生成器：PyTorch TorchScript 的 GPU 算子融合后端

NVFuser 融合代码生成器PyTorch TorchScript 的 GPU 算子融合后端【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorchNVFuser 是 PyTorch 中面向 NVIDIA GPU 的融合代码生成器它作为 TorchScript Profiling Graph Executor 的后端将add、relu等相邻算子融合为单个自定义 CUDA kernel以降低 kernel 启动开销与中间结果的内存往返。本文围绕该组件的官方指南系统讲解融合行为的调节开关single node / horizontal / guard / skip node、基于 TorchScript 与PYTORCH_JIT_LOG_LEVEL的融合调试方法、PYTORCH_NVFUSER_DISABLE/PYTORCH_NVFUSER_ENABLE环境变量族以及禁用 nvfuser 的三种方式帮助你快速定位“没融合”“融合了但回退”“结果不正确”三类问题。读完本文你将掌握一套可复制的 nvfuser 调参与故障排查流程。说明本文内容以 torch/csrc/jit/codegen/cuda/README.md 为主体骨架展开并补充当前仓库源码中的实现细节作佐证。需特别提醒在当前仓库版本中nvfuser 在 TorchScript 侧已进入弃用deprecated状态torch/csrc/jit/codegen/cuda/interface.cpp 中相关 API 均返回false并打印弃用警告下文的所有开关、环境变量与调试技巧均以该 README 所描述的 nvfuser 集成时代为准适用于包含该集成的历史 PyTorch 版本。什么是 NVFuser从 TorchScript 图到单个 CUDA Kernel在传统执行方式中x 1.0、relu会被逐个 dispatch 到独立的 CUDA kernel每次 kernel 启动都要付出 launch 开销并且中间结果需要写回显存再读入。NVFuser 的思路是把一批兼容的算子合并成一个prim::CudaFusionGroup由代码生成器产出单个融合 kernel 一次执行完整个子图。在 TorchScript 的 Profiling Graph Executor 中这个流程分三步Profiling前几次运行收集张量的 shape、dtype、device 等元信息prim::profile节点Fusion pass基于 profile 信息把相邻且兼容的算子划入融合组。当前仓库中通用的图融合 pass 位于 torch/csrc/jit/passes/graph_fuser.cpp其中isSimpleMapgraph_fuser.cpp定义了“simple mappable operator”的判定条件单一 tensor 输出、输出与所有 tensor 输入 shape 一致、scalar type 一致、位于同一 device 且产出 dense 无重叠输出Codegen 与执行融合组交给 nvfuser 的 codegen 运行时系统生成并缓存 kernel。正因为 nvfuser 是 ProfilingExecutor 的后端前几次运行只是做 profiling真正的融合 kernel 需要等优化包括 fusion pass生效之后才会出现——这是理解后续一切调试技巧的前提。调节融合行为的四个简单开关README 给出了四个 Python 层开关全部通过torch._C暴露1. 单节点融合_jit_set_nvfuser_single_node_modetorch._C._jit_set_nvfuser_single_node_mode(True)默认情况下只有当两个及以上兼容算子聚合在一起时才会创建融合组。开启该模式后融合 pass 允许仅包含单个节点的融合组。其价值在于测试便利可以逐个算子验证 codegen 的正确性性能探索某些场景下单个节点生成的专用 kernel 可能优于框架原生 CUDA kernel例如包含融合优化如 FMA、索引提升时。2. 水平融合_jit_set_nvfuser_horizontal_modetorch._C._jit_set_nvfuser_horizontal_mode(True)默认融合 pass 采用“生产者→消费者”的垂直融合producer to consumer把数据依赖链上的算子串成一个 kernel。水平模式则允许共享同一 tensor 输入的兄弟节点合并例如x / \ a b - 兄弟节点共享输入 x水平模式可将 a、b 融合这能节省输入的内存带宽——同一份输入数据只被读取一次。特别适合“一个张量同时喂给多个 elementwise 分支”的模型结构。3. 关闭融合守卫_jit_set_nvfuser_guard_modetorch._C._jit_set_nvfuser_guard_mode(False)融合 pass 会为每个融合组生成一个运行时检查即 IR 中的prim::CudaFusionGuard用于校验融合组的前置假设tensor 元信息shape/strides/dtype/device、常量输入、profiled 常量是否与编译时一致。关闭 guard 意味着跳过这些运行时校验。该开关“really is only used for testing”——用于强制确保生成的 kernel 确实被测试到。训练脚本中应避免关闭 guard否则一旦实际输入的 shape 与编译时不符kernel 将基于错误的假设执行可能产生错误结果。4. 按节点类型跳过融合_jit_set_nvfuser_skip_node_kindtorch._C._jit_set_nvfuser_skip_node_kind(aten::add, True)第一个参数是节点类型node kind如aten::add第二个参数控制该节点在融合中是开启还是关闭。用于在保持其他算子继续融合的前提下定向禁用某一类算子的融合——例如怀疑某个算子的融合 kernel 有 bug 时可以先把它从融合中摘除来二分定位。在底层这些开关经由 torch/csrc/jit/codegen/cuda/interface.cpp 中的setSingletonFusion、setHorizontalFusion、getCudaFusionGuardMode、skipNode等接口与融合运行时交互当前版本这些接口已标记弃用并返回false见下节。重要前提当前仓库中 nvfuser 的弃用状态在阅读和运行本节及后续示例前必须先了解一个事实在当前仓库的源码中nvfuser 的 TorchScript 集成已经被移除并标记为弃用。torch/csrc/jit/python/init.cpp 中_jit_set_nvfuser_single_node_mode、_jit_set_nvfuser_enabled、_jit_set_nvfuser_horizontal_mode、_jit_set_nvfuser_guard_mode、_jit_set_nvfuser_skip_node_kind等 API 全部实现为“打印弃用警告并空操作”TORCH_WARN( nvfuser is no longer supported in torch script, use _jit_set_nvfuser_skip_node_kind is deprecated and a no-op);torch/csrc/jit/codegen/cuda/interface.cpp 中isEnabled()、setEnabled()、canBeEnabled()等一律返回falsesetEnabled(true)会直接TORCH_INTERNAL_ASSERT失败并提示 nvfuser support in torchscript is removed and cannot be enabled!。因此本文后面给出的所有开关、环境变量与日志调试方法均针对 README 所描述的 nvfuser 集成时代历史版本。若你在当前版本上执行这些命令只会看到弃用警告如需实际体验 nvfuser 调试流程应使用仍包含该集成的 PyTorch 版本。当前仓库中实际承担 GPU 融合职责的是 torch/csrc/jit/passes/graph_fuser.cpp 所代表的通用图融合 pass 及其它后端。融合调试TorchScript 侧的两条主线README 以一个最小脚本贯穿调试讲解。脚本内容如下import torch def forward(x): o x 1.0 o o.relu() return o shape (2, 32, 128, 512) input torch.rand(*shape).cuda() t torch.jit.script(forward) with torch.jit.fuser(fuser2): for k in range(4): o t(input)注意脚本中torch.jit.fuser(fuser2)显式选择 nvfuserfuser2 即 nvfuser 的代号作为融合器并循环运行 4 次让 ProfilingExecutor 完成 profiling 并触发优化。1. TorchScript IR 图确认融合是否发生有两种方式查看融合后的图方式一在 Python 脚本中打印优化后图print(t.graph_for(input))多次运行让优化生效之后打印直接呈现最终的优化 IR。方式二开启 profiling executor 的图转储PYTORCH_JIT_LOG_LEVELprofiling_graph_executor_impl python your pytorch scriptPYTORCH_JIT_LOG_LEVEL是 TorchScript 的日志等级机制profiling_graph_executor_impl会 dump ProfilingExecutor 内部执行的各个图。如何阅读输出在图打印中关键是寻找prim::CudaFusionGroup_X节点——它代表一个将被 nvfuser 编译执行的融合 kernel。以 README 的示例输出为例Optimized Graph: graph(%x.1 : Tensor): %12 : bool prim::CudaFusionGuard[types[Float(2, 32, 128, 512, strides[20965536, 512, 1], requires_grad0, devicecuda:0)]](%x.1) %11 : Tensor prim::If(%12) block0(): %o.8 : Tensor prim::CudaFusionGroup_0cache_id0 - (%o.8) block1(): %18 : Function prim::Constant[namefallback_function, fallback1]() %19 : (Float(2, 32, 128, 512, strides[20965536, 512, 1], requires_grad0, devicecuda:0)) prim::CallFunction(%18, %x.1) %20 : Float(2, 32, 128, 512, strides[20965536, 512, 1], requires_grad0, devicecuda:0) prim::TupleUnpack(%19) - (%20) return (%11) with prim::CudaFusionGroup_0 graph(%2 : Float(2, 32, 128, 512, strides[20965536, 512, 1], requires_grad0, devicecuda:0)): %4 : int prim::Constant[value1]() %3 : float prim::Constant[value1.]() # test.py:6:12 %o.1 : Float(2, 32, 128, 512, strides[20965536, 512, 1], requires_grad0, devicecuda:0) aten::add(%2, %3, %4) # test.py:6:8 %o.5 : Float(2, 32, 128, 512, strides[20965536, 512, 1], requires_grad0, devicecuda:0) aten::relu(%o.1) # test.py:7:8 return (%o.5)这个输出揭示了融合后的完整结构prim::CudaFusionGuard携带编译时的 tensor 类型shape、strides、requires_grad、device即运行时校验prim::If的两个分支构成guard 成功走融合 kernel、guard 失败走 fallback的双路径prim::CudaFusionGroup_0[cache_id0]是融合 kernel 的调用点其子图内是aten::addaten::relu——一次 kernel 启动完成两个算子cache_id0表明 nvfuser 按输入类型缓存已编译的 kernelGraphCache 机制。训练场景的关键提醒融合 pass 只会在prim::DifferentiableGraph内部运行。如果脚本化模型有需要梯度的输入第一件事就是确认目标算子确实位于 differentiable 子图内。由于 profiling executor 会朴素地 dump 所有执行过的图、且 differentiable 图由嵌套的图执行器运行输出可能相当混乱——每个原始图对应一段分段的Optimized Graph其中一段对应一个 differentiable 节点。相关实现见 torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp。2. CUDA 融合图检查融合 pass 的输入输出PYTORCH_JIT_LOG_LEVELgraph_fuser python your pytorch scriptgraph_fuser级别会 dump 融合 pass 的输入图和输出图是检查融合 pass 逻辑为什么某些节点没被融合的最佳入口。日志中需要关注两张图Before Fusion融合 pass 运行的子图。注意此时算子间还夹着prim::profile节点profiled_type...这是 profiling 阶段的产物Before Compilation交给 codegen 后端的图其中每个CudaFusionGroup都会触发 codegen 运行时系统生成 kernel 来执行子图同时带有CudaFusionGuardprim::IfFallbackGraph的回退结构。README 给出的示例输出完整展示了这两个阶段此处节选核心部分Before Fusion: graph(%x.1 : Tensor): %2 : float prim::Constant[value1.]() %1 : int prim::Constant[value1]() %3 : Tensor prim::profile[profiled_typeFloat(2, 32, 128, 512, strides[20965536, 512, 1], requires_grad0, devicecuda:0)](%x.1) %o.10 : Tensor aten::add(%3, %2, %1) # test.py:6:8 %5 : Tensor prim::profile[profiled_typeFloat(2, 32, 128, 512, strides[20965536, 512, 1], requires_grad0, devicecuda:0)](%o.10) %o.7 : Tensor aten::relu(%5) # test.py:7:8 ... return (%7, %8) Before Compilation: graph(%x.1 : Tensor): %13 : bool prim::CudaFusionGuard[types[...]](%x.1) %12 : Tensor prim::If(%13) block0(): %o.11 : Tensor prim::CudaFusionGroup_0(%x.1) - (%o.11) block1(): %o.7 : Tensor prim::FallbackGraph_1(%x.1) - (%o.7) return (%12, %12)排查“没发生融合”debug no-fusion目前有多个消费者通过“把计算 lower 成 TorchScript、再用 ProfilingExecutor 执行”的方式使用 nvfuser。在不深入集成细节的前提下README 给出了三条排查“未融合”的经验法则多运行几次 TorchScript 模块5 次是个幸运数字以触发融合。ProfilingExecutor 用前几次运行做 profiling后续的优化包括启用 nvfuser 的融合 pass依赖 profiling 信息才能运行所以最初的几次运行不会出现融合 kernel。注意 profiling 的运行次数与模型有关不同模型可能需要不同次数。融合 kernel 应以prim::CudaFusionGroup出现在 TorchScript IR 中。用jit_model.graph_for(*inputs)查看优化后的图确认融合是否发生。如果脚本化模型的输入需要梯度融合只发生在prim::DifferentiableGraph内部的图。图无法 autodiff 的原因很多可查看 torch/csrc/jit/runtime/symbolic_scripts.cpp——它列出了所有可 autodiff 的算子注意这与 autograd 支持的算子列表不是同一个列表。此外还存在一个阈值过小的 autodiff 图会被内联/回退可通过torch._C._debug_set_autodiff_subgraph_inlining(False)禁用。排查“融合 kernel 出问题”debug nvfuser mal-functioning假设 ProfilingExecutor 已正常工作——你看到一个本应融合的区域却没有进入融合 kernelREADME 给出了四个递进的排查手段1. Dump 融合 pass 结果PYTORCH_JIT_LOG_LEVELgraph_fuser python your_script.py log查找日志中Before Fusion与Before Compilation两张图前者展示融合 pass 运行的图区间后者展示融合结果CudaFusionGroup。2. 查看哪些算子未融合及其大致原因PYTORCH_JIT_LOG_LEVELpartition:graph_fuser python your_script.py logpartition语法提升partition模块的日志级别。启用来自 partition.cpp 的 GRAPH_UPDATE 后每当某节点被融合 pass 拒绝时都会 dump 一条日志直接告诉你拒绝原因如isSimpleMap不满足、device 不一致、shape 不匹配等判定条件。3. 禁用 FALLBACK 路径如果执行模型时出现“走了 FALLBACK 路径”的警告说明 codegen 或融合 pass 出现了非预期失败。虽然功能上仍然正确fallback 保证结果对但这极可能带来性能回退。建议禁用 fallback 路径让错误被如实上报以便提交信息完整的 issuePYTORCH_NVFUSER_DISABLEfallback python your_script.py log4. 精确定位引发错误的 kernel / 融合模式对于包含多个融合模式的大模型要定位究竟是哪个融合触发了 FALLBACK 并构造最小复现并不容易。README 建议的组合拳PYTORCH_NVFUSER_DISABLEfallback \ PYTORCH_JIT_LOG_LEVELpartition:graph_fuser:kernel_cache \ python your_script.py logpartition:graph_fuser提升融合 pass 分区日志kernel_cache进一步提升 kernel_cache 的日志记录所有由 nvfuser 解析到 codegen IR 的 TorchScript IR以及 nvfuser 生成并执行的 kernel由于已禁用 fallback日志中最后一条往往就是失败的融合点。阅读提示寻找最后一个Before Compilation:表示解析失败或running GraphCache: xxxxx表示 JIT 编译/执行失败。对于后者还可以反查该 GraphCache 地址——前面应该已经 dump 过对应的 TorchScript IR。查询 nvfuser codegen kernelPYTORCH_NVFUSER_DUMP系列README 指出PYTORCH_NVFUSER_DUMP相关的调试开关定义在[pytorch_source_path]/torch/csrc/jit/codegen/cuda/utils.cpp该目录在当前仓库中仅保留 README 与 interface 文件utils.cpp 属于 nvfuser 集成时代的源码。几个常用的 dump 开关开关值作用dump_eff_bandwidth打印每个生成 kernel 的有效带宽。实现上为 kernel 执行时间除以 I/O buffer 大小是衡量带宽受限 kernel 性能的简单直观指标cuda_kernel打印生成的 CUDA kernel 源码launch_param打印生成 kernel 的 launch 配置网格/块维度等kernel_args打印所有已执行 codegen kernel 的输入/输出/buffer tensor。注意 buffer 会标注是否零初始化——这暗示可能存在一个额外的 kernel 用于在 codegen kernel 之前填充 tensor使用方式PYTORCH_NVFUSER_DUMPcuda_kernel,launch_param python your_script.py常见问题FAQQ1开启 nvfuser 后出现性能回退第一步确认融合 kernel 是否真正在跑用export PYTORCH_NVFUSER_DISABLEfallback禁用 fallback 后运行模型看是否命中错误导致回退。如果开启 nvfuser 后输出异常可通过PYTORCH_NVFUSER_DISABLE关闭部分可选特性来定位例如fma禁用 FMA融合乘加指令index_hoist禁用公共索引表达式提升优化predicate_elimination禁用冗余谓词消除优化unroll_with_rng使用 RNG 时禁用展开unroll。export PYTORCH_NVFUSER_DISABLEfma,index_hoist以上逗号分隔的写法会同时禁用 FMA 与索引提升。Q2nvfuser 没有带来任何加速检查脚本模型里是否有融合发生。运行PYTORCH_JIT_LOG_LEVELgraph_fuser如果日志中没有任何融合 pass 的前后图 dump说明 TorchScript 侧出了问题、融合 pass 根本没执行。此时回到上文“排查没发生融合”章节逐条核对运行次数、DifferentiableGraph、可 autodiff 算子列表等。Q3遇到 codegen 问题如何禁用 nvfuserREADME 给出了三种方式按优先级从高到低排列强制使用 NNC 而非 nvfuser 做 GPU 融合环境变量export PYTORCH_JIT_USE_NNC_NOT_NVFUSER1用 torch API 禁用torch._C._jit_set_nvfuser_enabled(False)用环境变量禁用export PYTORCH_JIT_ENABLE_NVFUSER0Q4还有更多调节 nvfuser 融合的旋钮吗除前述开关外两类可选特性通过环境变量暴露PYTORCH_NVFUSER_DISABLEopt-out 特性例如export PYTORCH_NVFUSER_DISABLEfallback,fmafallback编译失败时禁用 aten fallbackfma禁用融合乘加。README 特别提醒禁用 fma 通常会带来性能回退强烈建议不要禁用。PYTORCH_NVFUSER_ENABLEopt-in 特性complex启用 nvfuser 对复数浮点类型的支持。当前为实验特性默认关闭以避免功能回归linear_decomposition启用 linear 层中 bias 加法的分解同理conv_decomposition启用 conv 层中 bias 加法的分解。在部分小型 benchmark 模型中这类分解带来的编译开销超过了快速 kernel 的收益因此被改为 opt-in。示例export PYTORCH_NVFUSER_ENABLEcomplex总结一套完整的 nvfuser 调参与排障流程把 README 的知识点串成一条可执行的排障路径确认融合多运行几次约 5 次脚本用t.graph_for(input)或PYTORCH_JIT_LOG_LEVELprofiling_graph_executor_impl找prim::CudaFusionGroup排查未融合用PYTORCH_JIT_LOG_LEVELgraph_fuser看Before Fusion/Before Compilation用PYTORCH_JIT_LOG_LEVELpartition:graph_fuser看节点被拒原因训练场景核对是否在prim::DifferentiableGraph内排查融合异常用PYTORCH_NVFUSER_DISABLEfallback暴露真实错误必要时叠加PYTORCH_JIT_LOG_LEVELpartition:graph_fuser:kernel_cache精确定位失败融合点并关注Before Compilation:与running GraphCache:两类日志标记细粒度调参单节点/水平融合/guard/按节点跳过四个开关按需开启PYTORCH_NVFUSER_DISABLE关闭fma、index_hoist等可疑优化PYTORCH_NVFUSER_ENABLE开启complex、linear_decomposition等实验特性回退方案按优先级使用PYTORCH_JIT_USE_NNC_NOT_NVFUSER1、torch._C._jit_set_nvfuser_enabled(False)、PYTORCH_JIT_ENABLE_NVFUSER0禁用 nvfuser。最后再次强调当前仓库中 nvfuser 的 TorchScript 集成已弃用torch/csrc/jit/codegen/cuda/interface.cpp 与 torch/csrc/jit/python/init.cpp 中的相关 API 均为空操作并提示弃用上述全部技巧适用于仍集成 nvfuser 的历史 PyTorch 版本。如需研究当前版本的实际 GPU 融合实现可继续阅读 torch/csrc/jit/passes/graph_fuser.cpp 与 torch/csrc/jit/OVERVIEW.md。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考