Ray Compiled Graph 性能剖析指南：PyTorch Profiler、Nsight 与编译图可视化

Ray Compiled Graph 性能剖析指南PyTorch Profiler、Nsight 与编译图可视化【免费下载链接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.项目地址: https://gitcode.com/gh_mirrors/ra/ray导读Ray Compiled Graph编译图CGraph是 Ray Core 面向 GPU 加速工作负载提供的高性能执行路径它将一组跨 Actor 的 DAG 预先编译成确定性的调度与通信方案从而降低每次执行时的任务级与系统级开销。本文以 doc/source/ray-core/compiled-graph/profiling.rst 为核心完整讲解编译图的两种性能剖析方案——基于 PyTorch Profiler 的 Torch 追踪以及基于 Nsight Systems NVTX 的系统级剖析并附带编译图结构的可视化方法。读完本文你将能够通过一行环境变量开启 Torch profiling 并产出每个 Actor 独立的 trace 文件通过runtime_env为参与编译图的 Actor 挂载 Nsight 剖析并解读/tmp/ray/session_*/logs下的剖析结果使用 NVTX 对执行循环中的方法调用做细粒度标注最后用visualize()将编译后的图结构导出为 PNG 以便直观排查调度问题。为什么需要剖析 Compiled Graph编译图将 DAG 中每个 Actor 的方法调用与它们之间的数据传输如 NCCL 张量传输通道预编译为一份执行计划schedule由每个 Actor 内的执行循环按计划顺序执行。性能瓶颈往往来自两类开销任务级开销单个 task 从准备、反序列化到实际方法执行的耗时以及多次执行之间的波动系统级开销调度、通信如 NCCL 通道建立、张量传输、内存拷贝等不属于用户计算代码的部分。官方文档即本主题来源文档指出Compiled Graph 提供基于 PyTorch 和基于 Nsight 两套剖析能力目的是更好地理解单个任务、系统开销与性能瓶颈开发者可以按偏好任选其一。从源码看这两套剖析均由环境变量开关控制并在 Actor 的执行循环入口统一挂载见 python/ray/dag/compiled_dag_node.py 中的do_exec_tasks实现。PyTorch Profiler一行环境变量开启 Torch 追踪开启方式PyTorch 剖析是成本最低的切入方式运行脚本前设置环境变量RAY_CGRAPH_ENABLE_TORCH_PROFILING1即可。例如对于编译图脚本example.pyRAY_CGRAPH_ENABLE_TORCH_PROFILING1 python3 example.py无需修改任何业务代码。从 python/ray/dag/constants.py 可以看到该开关的定义# Feature flag to turn on torch profiling. # This cannot be used together with RAY_CGRAPH_ENABLE_NVTX_PROFILING. RAY_CGRAPH_ENABLE_TORCH_PROFILING ( os.environ.get(RAY_CGRAPH_ENABLE_TORCH_PROFILING, 0) 1 )底层实现执行循环中挂载 torch.profiler开启后每个 Actor 的do_exec_tasks执行循环会在进入无限调度循环前启动torch.profiler.profile并同时采集 CPU 与 CUDA 活动、记录调用栈通过 TensorBoard trace handler 落盘见 python/ray/dag/compiled_dag_node.pyif RAY_CGRAPH_ENABLE_TORCH_PROFILING: assert ( not RAY_CGRAPH_ENABLE_NVTX_PROFILING ), NVTX and torch profiling cannot be enabled at the same time. import torch torch_profile torch.profiler.profile( activities[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], with_stackTrue, on_trace_readytorch.profiler.tensorboard_trace_handler( compiled_graph_torch_profiles ), ) torch_profile.start()要点每个参与编译图的 Actor 都会各自启动一个 profile 实例因此每个 Actor 生成一份独立的 trace 文件结果输出到当前工作目录下的compiled_graph_torch_profiles目录with_stackTrue会记录 Python 调用栈方便定位到用户方法同时采集 CPU 与 CUDA 活动可观察 GPU 内核执行与 CPU 侧调度的时间关系。查看与可视化 trace运行结束后用浏览器打开 https://ui.perfetto.dev/Perfetto UI将compiled_graph_torch_profiles目录下的 trace 文件拖入即可查看时间线。通过对比各 Actor 的时间线可以直观发现某个 Actor 的执行循环是否长时间处于空闲等待上游数据/通信通道方法调用与 NCCL 传输在时间轴上的重叠程度单个 task 在 CPU 与 GPU 上的耗时构成。Nsight Systems通过 runtime_env 启用系统级剖析前置为 Actor 配置 nsight runtime_envCompiled Graph 建立在 Ray 既有 profiling 能力之上。要开启 Nsight 剖析不需要改脚本执行方式而是为涉及的 Actor 指定runtime_env{nsight: ...}具体配置方式参考 Ray 的 Nsight 使用说明。nsight配置项可以是字符串default使用默认配置也可以是 Nsight Systems 选项的字典参见 doc/source/ray-core/handling-dependencies.rst 对 runtime_envnsight键的说明。完整示例构建编译图并执行以下代码来自仓库示例 doc/source/ray-core/doc_code/cgraph_profiling.py先创建带 Nsight runtime_env 的 GPU Actorimport ray import torch from ray.dag import InputNode ray.remote(num_gpus1, runtime_env{nsight: default}) class RayActor: def send(self, shape, dtype, value: int): return torch.ones(shape, dtypedtype, devicecuda) * value def recv(self, tensor): return (tensor[0].item(), tensor.shape, tensor.dtype) sender RayActor.remote() receiver RayActor.remote()然后按常规方式构建并编译 DAG注意这里通过with_tensor_transport(transportnccl)指定张量走 NCCL 通道传输shape (10,) dtype torch.float16 # Test normal execution. with InputNode() as inp: dag sender.send.bind(inp.shape, inp.dtype, inp[0]) dag dag.with_tensor_transport(transportnccl) dag receiver.recv.bind(dag) compiled_dag dag.experimental_compile() for i in range(3): shape (10 * (i 1),) ref compiled_dag.execute(i, shapeshape, dtypedtype) assert ray.get(ref) (i, shape, dtype)最后按常规方式运行脚本python3 example.py执行结束后Compiled Graph 会将剖析结果输出到/tmp/ray/session_*/logs/{profiler_name}目录下session_*为本次 Ray 会话目录{profiler_name}为 profiler 名称。NVTX细粒度方法级标注如果希望对方法调用与系统开销做更细粒度的分析可额外设置环境变量RAY_CGRAPH_ENABLE_NVTX_PROFILING1 python3 example.py该开关在 python/ray/dag/constants.py 中定义。开启后Compiled Graph 在底层利用 NVTXNVIDIA Tools Extension Library自动为编译图各 Actor 执行循环中调用的所有方法添加标注使 Nsight Systems 时间线上能清晰地区分每个方法调用的起止见 python/ray/dag/compiled_dag_node.pyif RAY_CGRAPH_ENABLE_NVTX_PROFILING: assert ( not RAY_CGRAPH_ENABLE_TORCH_PROFILING ), NVTX and torch profiling cannot be enabled at the same time. try: import nvtx except ImportError: raise ImportError( Please install nvtx to enable nsight profiling. You can install it by running pip install nvtx. ) nvtx_profile nvtx.Profile() nvtx_profile.enable()需要注意的是使用 NVTX 剖析前需要安装nvtx包pip install nvtxNVTX 与 Torch profiling 二者互斥不能同时开启源码中通过assert强制校验剖析结果的查看方式与 Ray 常规的 Nsight 剖析结果相同即打开 Nsight Systems 分析/tmp/ray/session_*/logs/{profiler_name}下的结果文件。三种 profiling 开关小结环境变量作用输出位置依赖备注RAY_CGRAPH_ENABLE_TORCH_PROFILING1启动 torch.profiler采集 CPU/CUDA 活动当前目录compiled_graph_torch_profiles/每个 Actor 一份 tracetorch与 NVTX 互斥runtime_envnsight: default启动 Nsight Systems 系统级剖析/tmp/ray/session_*/logs/{profiler_name}Nsight Systems通过 runtime_env 配置RAY_CGRAPH_ENABLE_NVTX_PROFILING1NVTX 自动标注执行循环中的方法调用随 Nsight 结果一起pip install nvtx与 Torch profiling 互斥可视化编译图结构基本用法剖析着眼于时间而理解结构则需要可视化。在调用experimental_compile()编译图之后调用CompiledDAG.visualize()即可将图结构导出。来自 doc/source/ray-core/doc_code/cgraph_visualize.py 的完整示例import ray from ray.dag import InputNode, MultiOutputNode ray.remote class Worker: def inc(self, x): return x 1 def double(self, x): return x * 2 def echo(self, x): return x sender1 Worker.remote() sender2 Worker.remote() receiver Worker.remote() with InputNode() as inp: w1 sender1.inc.bind(inp) w1 receiver.echo.bind(w1) w2 sender2.double.bind(inp) w2 receiver.echo.bind(w2) dag MultiOutputNode([w1, w2]) compiled_dag dag.experimental_compile() compiled_dag.visualize()默认情况下Ray 会在当前工作目录生成名为compiled_graph.png的 PNG 图片。注意这需要安装graphvizpip install graphviz否则会抛出 ImportError见 python/ray/dag/compiled_dag_node.py。接口签名与参数visualize()的完整签名来自 python/ray/dag/compiled_dag_node.pydef visualize( self, filename: str compiled_graph, format: str png, view: bool False, channel_details: bool False, ) - str:参数默认值说明filenamecompiled_graph输出文件名不含扩展名ASCII 格式下该参数被忽略formatpng输出格式如png、pdf、jpegascii则直接打印到控制台viewFalse非 ASCII 格式下是否用默认查看器打开ASCII 格式下是否打印并返回channel_detailsFalse为True时在边上附加通道类型与细节与ascii格式不兼容返回值对 Graphviz 格式png/pdf/jpeg 等返回图的 DOT 字符串表示对 ASCII 格式返回 ASCII 字符串。读图节点与边的含义下面这张图展示了上述示例代码的可视化结果。同一 Actor 的任务使用相同颜色可以据此快速识别任务所属的 Actor 以及任务间的依赖关系。结合 python/ray/dag/compiled_dag_node.py 的绘制逻辑节点标注规则如下InputNode蓝色矩形与InputAttributeNode蓝色矩形图的输入ClassMethodNode椭圆按 Actor 着色标注为Actor: 类名、ID: Actor ID 前 6 位...、Method: 方法名是图中最核心的节点MultiOutputNode黄色矩形图的汇合输出点同色椭圆即属于同一 Actor 的任务边表示数据流依赖方向若设置channel_detailsTrue边上还会标注通道类型如 NCCL与传输细节。这张图可以帮助你在剖析之前确认数据流是否如预期地在各 Actor 之间传递、是否存在意外的串行依赖、多个输入分支是否真正并行。实战建议如何系统性定位编译图瓶颈综合以上三套工具推荐按以下步骤定位 Compiled Graph 的性能瓶颈先用visualize()检查结构确认 DAG 编译结果符合预期Actor 归属、任务依赖、并行分支排除结构性问题开启 PyTorch profiling 观察任务级耗时RAY_CGRAPH_ENABLE_TORCH_PROFILING1跑一次在 Perfetto 中对比各 Actor trace定位单个 task 的 CPU/GPU 耗时与空闲等待开启 Nsight NVTX 深入系统开销为 Actor 配置runtime_env{nsight: default}配合RAY_CGRAPH_ENABLE_NVTX_PROFILING1在 Nsight Systems 中观察调度、通信NCCL与用户方法调用的时间线细节交叉验证注意 NVTX 与 Torch profiling 互斥两套剖析需要分开运行剖析本身会引入一定开销建议用多次运行的平均趋势而非单次结果做结论。参考路径速查官方文档源文件doc/source/ray-core/compiled-graph/profiling.rst剖析开关定义python/ray/dag/constants.py剖析挂载与执行循环实现python/ray/dag/compiled_dag_node.pyNsight 剖析示例doc/source/ray-core/doc_code/cgraph_profiling.py可视化示例doc/source/ray-core/doc_code/cgraph_visualize.py可视化输出示例图doc/source/images/compiled_graph_viz.pngruntime_envnsight配置说明doc/source/ray-core/handling-dependencies.rst【免费下载链接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.项目地址: https://gitcode.com/gh_mirrors/ra/ray创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考