在大语言模型的强化学习训练过程中,GPU 性能优化至关重要。随着模型规模不断扩大,如何高效地分析和优化 GPU 性能成为开发者面临的主要挑战之一。
NVIDIA Nsight™ Systems 是 NVIDIA 开发的一款强大的 GPU 性能分析工具。Ray 作为开源的分布式计算框架,在深度学习社区广受欢迎。然而,Nsight Systems 用户在分析 Ray 的任务负载时,会遇到与以往不同的挑战。本文将介绍如何在 Ray 分布式计算框架下集成 Nsight Systems 进行 GPU 性能分析,帮助开发者更好地理解和优化他们的强化学习训练流程。
大语言模型强化学习与 Ray
近年来,大语言模型 (LLM) 的强化学习(RL)取得了显著进展,从早期的监督式微调发展到现在的 PPO、DPO 等先进算法。这些算法通常需要复杂的分布式训练架构,而 Ray 凭借其灵活的 Actor 模型和分布式任务调度能力,成为许多开源框架 (如 OpenRLHF、verl、AReal、ROLL、NeMo RL) 的首选基础设施。
Ray 的核心优势包括:使用 Task 和 Actor 来抽象远端的计算任务,能够动态调度任务实现负载均衡,提供了方便高效的跨节点通信机制,能够对 GPU 和 CPU 等资源进行统筹配置和管理。这些特性使得 Ray 特别适合处理 LLM 强化学习中复杂的计算图和工作流。在 Ray 的框架下,其他深度学习的训练和推理引擎可以作为一个子模块,完成具体任务的实现,Ray 则像操作系统一样把不同的任务调度到一起,配合完成强化学习的训练任务。
NVIDIA Nsight Systems 工具概述
NVIDIA Nsight Systems 是一款系统级性能分析工具,旨在实现应用算法的可视化,找出程序中亟需优化的”瓶颈”并进行调整,以跨任意数量或大小的 CPU 和 GPU——从大型服务器到最小的系统级芯片 (SoC)——进行高效扩展。
首先, Nsight Systems 不仅可以跟踪 GPU 的活动,还能够跟踪 CPU 的事件,从而帮助程序员建立起对程序整体工作流程的认识。它会锁定目标应用,以便在时间轴中同时显示 GPU 和 CPU 活动、事件、注释、吞吐量和性能指标。GPU 工作负载与应用内的 CPU 事件进行关联,因此可以轻松识别和修复性能障碍。
Nsight Systems 能够跟踪 GPU 的活动,即追踪 GPU 工作负载。为进一步探索 GPU,切换 GPU 指标采样将绘制低级输入/ 输出 (IO) 活动,例如 PCIe 吞吐量、NVIDIA NVLink™ 以及动态随机访问内存 (DRAM) 活动。GPU 指标采样还可公开 SM 利用率、Tensor Core 活动、指令吞吐量和线程束占用率。同时,对于计算任务,Nsight Systems 支持研究 CUDA API 和追踪 CUDA 库,包括 cuBLAS、cuDNN 和 NVIDIA TensorRT™。对于程序员来说,在代码中主动放入可执行的注释对于分析程序的流程和性能瓶颈非常必要。NVIDIA Tools Extension Library (NVTX) 就是这样的工具,并且可以和 Nsight Systems 联合使用。它的核心功能包括:
- 标记事件(markers):在程序执行的特定点标记信息;
- 注释范围(ranges):在程序执行的两个点之间标注范围;
- 跟踪资源:为对象分配可显示的名称并跟踪其生命周期。
所有这些标记都可以附加类别、颜色等信息,非常便于在 Nsight Systems 中进行观察。
Ray 与 Nsight Systems 集成的挑战
在 Nsight Systems 的典型应用场景中, Nsight Systems 可以自动捕捉目标程序,完成事件跟踪和数据记录。但是用户在分析基于 Ray 实现的 RL 程序时发现,原来的方法失效了。
Nsight Systems 的典型应用方法是这样的,使用命令“nsys”来启动目标程序:
nsys [command_switch][optional command_switch_options][application] [optional application_options]
这里“nsys”命令可以通过不同的子命令来选择具体功能,子命令带有可选项。后面紧接着目标应用程序(application)及其可选的参数选项。通过这种调用方法,Nsight Systems 可以捕捉到目标应用程序的执行过程,进而跟踪其执行的各种事件。
但是这种方法在捕捉以 Ray 为基础框架的程序时却失效了。其根本原因在于,在 Ray 任务里,上面命令行里的 application 并非真正的计算程序,而是给 Ray 系统提交命令的另一个命令行。通常启动 Ray 程序有两种方法,一种是使用 ray job submit 直接向 Ray 系统提交任务;一种是像普通的 Python 程序一样使用 python 来启动,但在其程序内部仍然是启动一个 Ray 系统并提交任务,或者直接向已有的 Ray 系统提交任务。也就是说我们关心的计算任务是由 Ray 在远端节点启动的。以上两种方法作为 application 附在“nsys”命令后面的时候, Nsight Systems 只能捕捉到提交任务这个前导程序,而不是真正的计算任务。
集成方案的实现
理解 Ray与 Nsight Systems 集成的挑战和背后的原因,解决方案便较为明确。实际上,Ray 的使用手册已经关注到这个问题,并给出了相应的解决方案。在使用手册的小节“Run Nsight on Ray”和“Custom options”中讲到在定义 RayActor 的时候,应该在其 ray.remote 修饰中加入 runtime_env 的环境变量,那么 Ray 在启动这个 Actor 的时候,会把这些与 Nsight Systems 相关的环境变量转换为“nsys”的调用参数,并且使用“nsys”来启动这个 Actor 对应的进程(可以参考 Ray 的使用手册获得详细信息)。
问题似乎得到了解决。但大语言模型强化学习程序的复杂性引入了新问题。Ray 使用手册中的方法要求在 Actor 定义阶段就预先写好 Nsight 相关的环境变量,但是像 verl 这样的复杂程序,Actor 被抽象了,它是动态的定义并动态的调用,需要找到动态的时机把这些参数传递进去。实际上,除了在定义 RayActor class 的时候可以指定它的环境变量,在构造它的实例(instance)的时候,也可以指定它的环境变量,以手册中示例为例,可以通过以下方式构造 RayActor 的实例:
actor = RayActor.options(runtime_env={“nsight”: {…}}).remote()
这样,我们就可以为每个 RayActor 的实例传入不同环境变量了。当然 verl 的实际应用就更复杂,需要更多的设计,来管理跟踪这些环境变量,并在合适的时机传递给 Actor 进程。
在 verl 中集成 Nsight Systems 的设计考量
理解了 Ray 和 Nsight Systems 集成的挑战和解决方案,理论上我们就可以在任何一个基于 Ray 的强化学习框架中使用 Nsight Systems 来分析应用性能和瓶颈了。我们以 verl 为例为其集成了 Nsight Systems 的分析工具,并简要梳理总结了设计中的关键考量,期望能对其他框架集成 Nsight Systems 提供参考。
一、区分 controller 和 worker 进程
目前绝大部分的强化学习框架都接受单一控制流(single controller)的概念。一方面,强化学习算法非常复杂,演化速度又快,算法研究者希望自己推演的公式能够像脚本一样在工具中快速实现,这时候用户的思考模式是单进程的。另一方面,超大规模的计算任务一定是多进程进行的,大语言模型的训练和推理引擎都需要支持各种并行优化策略。这中间是有很大差距的。幸运的是,Ray 作为并行计算系统的基础架构,使得弥合这个差距成为可能。OpenRLHF 率先使用 Ray 作为调度的基础框架,verl 则率先提出使用单一控制流来简化算法用户的尝试流程。
这就需要我们在跟踪强化学习任务的时候,分别跟踪 controller 和 worker 进程。
- Controller 进程虽然没有 GPU 计算资源的信息,但是它为用户提供统一的全盘程序执行时间线,使用户能全面把握各个计算任务的整体耗时、分割及占比,避免过早陷入细节。在 verl 中,我们在启动 controller 的时候为其传入了运行时环境变量,这些变量是由 Hydra 配置传入,给用户精细控制 Nsight 的调用参数提供了可能性。
- Verl 对 worker 进程做了优雅的抽象,以便能够支持灵活的资源分配、共享,和不同计算任务的合并。忽略 Hydra 配置的传递细节,我们是在不同的角色(role)合并之后,最终的 worker 启动之前,把用户配置信息传递进入每个 worker 实例。
二、细粒度跟踪 worker 进程
完成所有配置并运行 Nsight Systems 进行目标程序分析时,用户会遇到一个现实的问题——跟踪数据库太大。强化学习的复杂性必须解决规模的问题。以下从三个角度来说明如何对 worker 进程进行更精细的控制。
- 目标训练步的选择。
大语言模型的训练是一个耗时的过程,而且目标性能瓶颈不一定在初始的若干步之内出现。我们无法对整个过程进行数据跟踪,这就需要我们在特定训练步的时候才开启 Nsight Systems 跟踪,该步结束的时候关闭跟踪。 Nsight Systems 提供了 capture-range 的能力,我们在verl程序内部通过 torch.cuda.profiler.start() 和 torch.cuda.profiler.stop() 对来控制开启和关闭 Nsight Systems 的跟踪动作。这样每步都会形成一个单独的数据库,方便后续分析。
- 目标 worker 成员(rank)的选择。
在一般的训练和推理引擎设计里,一个 GPU 对应一个进程,一个强化学习任务可能使用成百上千的 GPU,甚至到达万卡规模。因此,我们也可以使用 capture-range 机制,仅仅对感兴趣的成员打开 Nsight Systems 的跟踪动作。
- 分立任务的选择。
强化学习的一个训练步骤包括多个子任务(序列生成、优势计算、模型更新),即使是一个训练步骤,数据的规模仍可能很大。我们可以使用 capture-range 机制,对每一个分立任务独立启停 Nsight Systems 跟踪。
三、使用 NVTX 描绘整体运行流程
如前面介绍,NVTX 可以在程序内放置标记和范围,并在 Nsight Systems 时间线上显示。我们对 verl 中每一个关键步骤都进行了范围标记,使用户清晰地看到每个计算子任务之间是怎样的接力关系、时间占比如何。具体来看,我们给 step, gen, reward, old_log_prob, ref, values, adv, update_critic, update_actor, testing 等计算步骤做了标记。在 verl 的设计中,我们把这些标记与 verl 提供的计时函数进行了融合,使用 marked_timer 进行计时的任务都自动标记了相应的 NVTX 范围。
四、 Nsight Systems 的时间线示例
下图是使用 verl v0.4.1 版本的例子examples/ppo_trainer/run_qwen2-7b_rm_seq_balance_nsys.sh 得到 Nsight Systems 时间线。从图中可以看到,controller 进程显示训练一共持续了 6 步,选取的两个 worker 进程详细地跟踪了其中的第 1、2、5 步。有了详细的 GPU 活动记录,对齐到整体的控制流图,我们就可以有针对性的分析程序性能了。
常见的反馈问题
在 Nsight Systems 相关代码合入 verl 的主分支后,很多用户积极试用,在各自的强化学习训练任务的分析和优化中取得了很好的效果。这里汇总一些比较常见的反馈问题,供参考。
- 去哪里找跟踪到的数据结果,这个目录可以指定吗?
按照 Ray 的使用手册,跟踪到的数据库保存在每个服务器节点的/tmp/ray/session_*/logs/{profiler_name},并且跟踪目录是不能修改的。这是当前版本 Ray 的一个局限。 - 文件格式如何解读?
文件格式为 worker_process_<PID>.nsys-rep 或者 worker_process_<PID>.<RID>.nsys-rep,PID 是进程号,RID 是第几个 capture range。 - 去哪里找 controller 的跟踪数据?
因为 Ray 是在远端服务器节点调度计算任务的,所以它可能会把 controller 任务调度到任何一个节点,这样找起来就不方便。我们在 verl 的程序内让这个 controller 自己打印了自己的位置信息,包括节点的 hostname 和进程号,这样用户就可以在对应的位置找到相应的数据。 - 怎么查看这么多数据?
Nsight 提供 multi-report view 的模式,用户可以把 controller 和 workers 的数据一次性全部导入,工具会自动按照时间顺序把数据排布在视图之中,方便对齐事件的时间点。
参考链接
NVIDIA Nsight Systems 使用手册: User Guide — nsight-systems 2025.3 documentation
Ray 使用手册:https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html
OpenRLHF:https://github.com/OpenRLHF/OpenRLHF
verl:https://github.com/volcengine/verl
AReaL:https://github.com/inclusionAI/AReaL