在大語言模型的強(qiáng)化學(xué)習(xí)訓(xùn)練過程中，GPU 性能優(yōu)化至關(guān)重要。隨著模型規(guī)模不斷擴(kuò)大，如何高效地分析和優(yōu)化 GPU 性能成為開發(fā)者面臨的主要挑戰(zhàn)之一。

NVIDIA Nsight? Systems是 NVIDIA 開發(fā)的一款強(qiáng)大的 GPU 性能分析工具。Ray 作為開源的分布式計(jì)算框架，在深度學(xué)習(xí)社區(qū)廣受歡迎。然而，Nsight Systems 用戶在分析 Ray 的任務(wù)負(fù)載時，會遇到與以往不同的挑戰(zhàn)。本文將介紹如何在 Ray 分布式計(jì)算框架下集成 Nsight Systems 進(jìn)行 GPU 性能分析，幫助開發(fā)者更好地理解和優(yōu)化他們的強(qiáng)化學(xué)習(xí)訓(xùn)練流程。

大語言模型強(qiáng)化學(xué)習(xí)與 Ray

近年來，大語言模型 (LLM) 的強(qiáng)化學(xué)習(xí) (RL) 取得了顯著進(jìn)展，從早期的監(jiān)督式微調(diào)發(fā)展到現(xiàn)在的 PPO、DPO 等先進(jìn)算法。這些算法通常需要復(fù)雜的分布式訓(xùn)練架構(gòu)，而 Ray 憑借其靈活的 Actor 模型和分布式任務(wù)調(diào)度能力，成為許多開源框架 (如 OpenRLHF、verl、AReal、ROLL、NeMo RL) 的首選基礎(chǔ)設(shè)施。

Ray 的核心優(yōu)勢包括：使用 Task 和 Actor 來抽象遠(yuǎn)端的計(jì)算任務(wù)，能夠動態(tài)調(diào)度任務(wù)實(shí)現(xiàn)負(fù)載均衡，提供了方便高效的跨節(jié)點(diǎn)通信機(jī)制，能夠?qū)?GPU 和 CPU 等資源進(jìn)行統(tǒng)籌配置和管理。這些特性使得 Ray 特別適合處理 LLM 強(qiáng)化學(xué)習(xí)中復(fù)雜的計(jì)算圖和工作流。在 Ray 的框架下，其他深度學(xué)習(xí)的訓(xùn)練和推理引擎可以作為一個子模塊，完成具體任務(wù)的實(shí)現(xiàn)，Ray 則像操作系統(tǒng)一樣把不同的任務(wù)調(diào)度到一起，配合完成強(qiáng)化學(xué)習(xí)的訓(xùn)練任務(wù)。

NVIDIA Nsight Systems 工具概述

NVIDIA Nsight Systems 是一款系統(tǒng)級性能分析工具，旨在實(shí)現(xiàn)應(yīng)用算法的可視化，找出程序中亟需優(yōu)化的“瓶頸”并進(jìn)行調(diào)整，以跨任意數(shù)量或大小的 CPU 和 GPU —— 從大型服務(wù)器到最小的系統(tǒng)級芯片 (SoC) —— 進(jìn)行高效擴(kuò)展。

首先，Nsight Systems 不僅可以跟蹤 GPU 的活動，還能夠跟蹤 CPU 的事件，從而幫助程序員建立起對程序整體工作流程的認(rèn)識。它會鎖定目標(biāo)應(yīng)用，以便在時間軸中同時顯示 GPU 和 CPU 活動、事件、注釋、吞吐量和性能指標(biāo)。GPU 工作負(fù)載與應(yīng)用內(nèi)的 CPU 事件進(jìn)行關(guān)聯(lián)，因此可以輕松識別和修復(fù)性能障礙。

Nsight Systems 能夠跟蹤 GPU 的活動，即追蹤 GPU 工作負(fù)載。為進(jìn)一步探索 GPU，切換 GPU 指標(biāo)采樣將繪制低級輸入 / 輸出 (IO) 活動，例如 PCIe 吞吐量、NVIDIA NVLink? 以及動態(tài)隨機(jī)訪問內(nèi)存 (DRAM) 活動。GPU 指標(biāo)采樣還可公開 SM 利用率、Tensor Core 活動、指令吞吐量和線程束占用率。同時，對于計(jì)算任務(wù)，Nsight Systems 支持研究 CUDA API 和追蹤 CUDA庫，包括 cuBLAS、cuDNN 和NVIDIA TensorRT?。對于程序員來說，在代碼中主動放入可執(zhí)行的注釋對于分析程序的流程和性能瓶頸非常必要。NVIDlA Tools Extension Library (NVTX) 就是這樣的工具，并且可以和 Nsight Systems 聯(lián)合使用。它的核心功能包括：

• 標(biāo)記事件 (markers)：在程序執(zhí)行的特定點(diǎn)標(biāo)記信息；

• 注釋范圍 (ranges)：在程序執(zhí)行的兩個點(diǎn)之間標(biāo)注范圍；

• 跟蹤資源：為對象分配可顯示的名稱并跟蹤其生命周期。

所有這些標(biāo)記都可以附加類別、顏色等信息，非常便于在 Nsight Systems 中進(jìn)行觀察。

Ray 與 Nsight Systems 集成的挑戰(zhàn)

在 Nsight Systems 的典型應(yīng)用場景中，Nsight Systems 可以自動捕捉目標(biāo)程序，完成事件跟蹤和數(shù)據(jù)記錄。但是用戶在分析基于 Ray 實(shí)現(xiàn)的 RL 程序時發(fā)現(xiàn)，原來的方法失效了。

Nsight Systems 的典型應(yīng)用方法是這樣的，使用命令 “nsys” 來啟動目標(biāo)程序：

nsys[command_switch][optional command_switch_options][application] [optional application_options]

這里 “nsys” 命令可以通過不同的子命令來選擇具體功能，子命令帶有可選項(xiàng)。后面緊接著目標(biāo)應(yīng)用程序 (application) 及其可選的參數(shù)選項(xiàng)。通過這種調(diào)用方法，Nsight Systems 可以捕捉到目標(biāo)應(yīng)用程序的執(zhí)行過程，進(jìn)而跟蹤其執(zhí)行的各種事件。

但是這種方法在捕捉以 Ray 為基礎(chǔ)框架的程序時卻失效了。其根本原因在于，在 Ray 任務(wù)里，上面命令行里的 application 并非真正的計(jì)算程序，而是給 Ray 系統(tǒng)提交命令的另一個命令行。通常啟動 Ray 程序有兩種方法，一種是使用 ray job submit 直接向 Ray 系統(tǒng)提交任務(wù)；一種是像普通的 Python 程序一樣使用 python 來啟動，但在其程序內(nèi)部仍然是啟動一個 Ray 系統(tǒng)并提交任務(wù)，或者直接向已有的 Ray 系統(tǒng)提交任務(wù)。也就是說我們關(guān)心的計(jì)算任務(wù)是由 Ray 在遠(yuǎn)端節(jié)點(diǎn)啟動的。以上兩種方法作為 application 附在 “nsys” 命令后面的時候，Nsight Systems 只能捕捉到提交任務(wù)這個前導(dǎo)程序，而不是真正的計(jì)算任務(wù)。

集成方案的實(shí)現(xiàn)

理解 Ray 與 Nsight Systems 集成的挑戰(zhàn)和背后的原因，解決方案便較為明確。實(shí)際上，Ray 的使用手冊已經(jīng)關(guān)注到這個問題，并給出了相應(yīng)的解決方案。在使用手冊的小節(jié) “Run Nsight on Ray” 和 “Custom options” 中講到在定義 RayActor 的時候，應(yīng)該在其 ray.remote 修飾中加入 runtime_env 的環(huán)境變量，那么 Ray 在啟動這個 Actor 的時候，會把這些與 Nsight Systems 相關(guān)的環(huán)境變量轉(zhuǎn)換為 “nsys” 的調(diào)用參數(shù)，并且使用 “nsys” 來啟動這個 Actor 對應(yīng)的進(jìn)程（可以參考 Ray 的使用手冊獲得詳細(xì)信息）。

問題似乎得到了解決。但大語言模型強(qiáng)化學(xué)習(xí)程序的復(fù)雜性引入了新問題。Ray 使用手冊中的方法要求在 Actor 定義階段就預(yù)先寫好 Nsight 相關(guān)的環(huán)境變量，但是像 verl 這樣的復(fù)雜程序，Actor 被抽象了，它是動態(tài)的定義并動態(tài)的調(diào)用，需要找到動態(tài)的時機(jī)把這些參數(shù)傳遞進(jìn)去。實(shí)際上，除了在定義 RayActor class 的時候可以指定它的環(huán)境變量，在構(gòu)造它的實(shí)例 (instance) 的時候，也可以指定它的環(huán)境變量，以手冊中示例為例，可以通過以下方式構(gòu)造 RayActor 的實(shí)例：

actor = RayActor.options(runtime_env={"nsight": {…}}).remote()

這樣，我們就可以為每個 RayActor 的實(shí)例傳入不同環(huán)境變量了。當(dāng)然 verl 的實(shí)際應(yīng)用就更復(fù)雜，需要更多的設(shè)計(jì)，來管理跟蹤這些環(huán)境變量，并在合適的時機(jī)傳遞給 Actor 進(jìn)程。

在 verl 中集成 Nsight Systems 的設(shè)計(jì)考量

理解了 Ray 和 Nsight Systems 集成的挑戰(zhàn)和解決方案，理論上我們就可以在任何一個基于 Ray 的強(qiáng)化學(xué)習(xí)框架中使用 Nsight Systems 來分析應(yīng)用性能和瓶頸了。我們以 verl 為例為其集成了 Nsight Systems 的分析工具，并簡要梳理總結(jié)了設(shè)計(jì)中的關(guān)鍵考量，期望能對其他框架集成 Nsight Systems 提供參考。

一、區(qū)分 controller 和 worker 進(jìn)程

目前絕大部分的強(qiáng)化學(xué)習(xí)框架都接受單一控制流 (single controller) 的概念。一方面，強(qiáng)化學(xué)習(xí)算法非常復(fù)雜，演化速度又快，算法研究者希望自己推演的公式能夠像腳本一樣在工具中快速實(shí)現(xiàn)，這時候用戶的思考模式是單進(jìn)程的。另一方面，超大規(guī)模的計(jì)算任務(wù)一定是多進(jìn)程進(jìn)行的，大語言模型的訓(xùn)練和推理引擎都需要支持各種并行優(yōu)化策略。這中間是有很大差距的。幸運(yùn)的是，Ray 作為并行計(jì)算系統(tǒng)的基礎(chǔ)架構(gòu)，使得彌合這個差距成為可能。OpenRLHF 率先使用 Ray 作為調(diào)度的基礎(chǔ)框架，verl 則率先提出使用單一控制流來簡化算法用戶的嘗試流程。

這就需要我們在跟蹤強(qiáng)化學(xué)習(xí)任務(wù)的時候，分別跟蹤 controller 和 worker 進(jìn)程。

• Controller 進(jìn)程雖然沒有 GPU 計(jì)算資源的信息，但是它為用戶提供統(tǒng)一的全盤程序執(zhí)行時間線，使用戶能全面把握各個計(jì)算任務(wù)的整體耗時、分割及占比，避免過早陷入細(xì)節(jié)。在 verl 中，我們在啟動 controller 的時候?yàn)槠鋫魅肓诉\(yùn)行時環(huán)境變量，這些變量是由 Hydra 配置傳入，給用戶精細(xì)控制 Nsight 的調(diào)用參數(shù)提供了可能性。

• verl 對 worker 進(jìn)程做了優(yōu)雅的抽象，以便能夠支持靈活的資源分配、共享，和不同計(jì)算任務(wù)的合并。忽略 Hydra 配置的傳遞細(xì)節(jié)，我們是在不同的角色 (role) 合并之后，最終的 worker 啟動之前，把用戶配置信息傳遞進(jìn)入每個 worker 實(shí)例。

二、細(xì)粒度跟蹤 worker 進(jìn)程

完成所有配置并運(yùn)行 Nsight Systems 進(jìn)行目標(biāo)程序分析時，用戶會遇到一個現(xiàn)實(shí)的問題 —— 跟蹤數(shù)據(jù)庫太大。強(qiáng)化學(xué)習(xí)的復(fù)雜性必須解決規(guī)模的問題。以下從三個角度來說明如何對 worker 進(jìn)程進(jìn)行更精細(xì)的控制。

1. 目標(biāo)訓(xùn)練步的選擇

大語言模型的訓(xùn)練是一個耗時的過程，而且目標(biāo)性能瓶頸不一定在初始的若干步之內(nèi)出現(xiàn)。我們無法對整個過程進(jìn)行數(shù)據(jù)跟蹤，這就需要我們在特定訓(xùn)練步的時候才開啟 Nsight Systems 跟蹤，該步結(jié)束的時候關(guān)閉跟蹤。Nsight Systems 提供了 capture-range 的能力，我們在 verl 程序內(nèi)部通過 torch.cuda.profiler.start() 和 torch.cuda.profiler.stop() 對來控制開啟和關(guān)閉 Nsight Systems 的跟蹤動作。這樣每步都會形成一個單獨(dú)的數(shù)據(jù)庫，方便后續(xù)分析。

2. 目標(biāo) worker 成員 (rank) 的選擇

在一般的訓(xùn)練和推理引擎設(shè)計(jì)里，一個 GPU 對應(yīng)一個進(jìn)程，一個強(qiáng)化學(xué)習(xí)任務(wù)可能使用成百上千的 GPU，甚至到達(dá)萬卡規(guī)模。因此，我們也可以使用 capture-range 機(jī)制，僅僅對感興趣的成員打開 Nsight Systems 的跟蹤動作。

3. 分立任務(wù)的選擇

強(qiáng)化學(xué)習(xí)的一個訓(xùn)練步驟包括多個子任務(wù)（序列生成、優(yōu)勢計(jì)算、模型更新），即使是一個訓(xùn)練步驟，數(shù)據(jù)的規(guī)模仍可能很大。我們可以使用 capture-range 機(jī)制，對每一個分立任務(wù)獨(dú)立啟停 Nsight Systems 跟蹤。

三、使用 NVTX 描繪整體運(yùn)行流程

如前面介紹，NVTX 可以在程序內(nèi)放置標(biāo)記和范圍，并在 Nsight Systems 時間線上顯示。我們對 verl 中每一個關(guān)鍵步驟都進(jìn)行了范圍標(biāo)記，使用戶清晰地看到每個計(jì)算子任務(wù)之間是怎樣的接力關(guān)系、時間占比如何。具體來看，我們給 step, gen, reward, old_log_prob, ref, values, adv, update_critic, update_actor, testing 等計(jì)算步驟做了標(biāo)記。在 verl 的設(shè)計(jì)中，我們把這些標(biāo)記與 verl 提供的計(jì)時函數(shù)進(jìn)行了融合，使用 marked_timer 進(jìn)行計(jì)時的任務(wù)都自動標(biāo)記了相應(yīng)的 NVTX 范圍。

四、為其他硬件平臺預(yù)留設(shè)計(jì)空間

NVIDIA 一直致力于推動人工智能的發(fā)展，推進(jìn)算法的極限，擴(kuò)展人類認(rèn)知領(lǐng)域的邊界。主流大語言模型強(qiáng)化學(xué)習(xí)框架除了采用 NVIDIA 的加速計(jì)算方案，也考慮兼容不同的硬件方案。因此在 verl 集成 Nsight Systems 的設(shè)計(jì)內(nèi)，我們對 Nsight Systems 的使用方式和 NVTX 的主要函數(shù)接口進(jìn)行了標(biāo)準(zhǔn)抽象，如 ProfilerConfig, DistProfiler, DistProfilerExtension, mark_range_start, mark_range_end, mark_annotate 等，并提供了空白實(shí)現(xiàn)。我們?yōu)闃?biāo)準(zhǔn)抽象提供一個 Nsight Systems 的具體實(shí)現(xiàn)，如 NsightSystemsProfiler 等。Verl 可以在未來實(shí)現(xiàn)其他工具的集成。具體代碼可以參考https://github.com/volcengine/verl/blob/main/verl/utils/profiler/profile.py, performance.py, empty_annotations.py, nvtx_profile.py.

五、Nsight Systems 的時間線示例

下圖是使用 verl v0.4.1 版本的例子 examples/ppo_trainer/run_qwen2-7b_rm_seq_balance_nsys.sh 得到 Nsight Systems 時間線。從圖中可以看到，controller 進(jìn)程顯示訓(xùn)練一共持續(xù)了 6 步，選取的兩個 worker 進(jìn)程詳細(xì)地跟蹤了其中的第 1、2、5 步。有了詳細(xì)的 GPU 活動記錄，對齊到整體的控制流圖，我們就可以有針對性的分析程序性能了。

常見的反饋問題

在 Nsight Systems 相關(guān)代碼合入 verl 的主分支后，很多用戶積極試用，在各自的強(qiáng)化學(xué)習(xí)訓(xùn)練任務(wù)的分析和優(yōu)化中取得了很好的效果。這里匯總一些比較常見的反饋問題，供參考。

1. 去哪里找跟蹤到的數(shù)據(jù)結(jié)果，這個目錄可以指定嗎？
按照 Ray 的使用手冊，跟蹤到的數(shù)據(jù)庫保存在每個服務(wù)器節(jié)點(diǎn)的 /tmp/ray/session_*/logs/{profiler_name}，并且跟蹤目錄是不能修改的。這是當(dāng)前版本 Ray 的一個局限。

2. 文件格式如何解讀？
文件格式為 worker_process_.nsys-rep 或者 worker_process_..nsys-rep，PID 是進(jìn)程號，RID 是第幾個 capture range。

3. 去哪里找 controller 的跟蹤數(shù)據(jù)？
因?yàn)?Ray 是在遠(yuǎn)端服務(wù)器節(jié)點(diǎn)調(diào)度計(jì)算任務(wù)的，所以它可能會把 controller 任務(wù)調(diào)度到任何一個節(jié)點(diǎn)，這樣找起來就不方便。我們在 verl 的程序內(nèi)讓這個 controller 自己打印了自己的位置信息，包括節(jié)點(diǎn)的 hostname 和進(jìn)程號，這樣用戶就可以在對應(yīng)的位置找到相應(yīng)的數(shù)據(jù)。

4. 怎么查看這么多數(shù)據(jù)？
Nsight 提供 multi-report view 的模式，用戶可以把 controller 和 workers 的數(shù)據(jù)一次性全部導(dǎo)入，工具會自動按照時間順序把數(shù)據(jù)排布在視圖之中，方便對齊事件的時間點(diǎn)。

NVIDIA Nsight Systems 使用手冊：

https://docs.nvidia.com/nsight-systems/UserGuide/index.html

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

ROLL:

https://github.com/alibaba/ROLL

NeMo RL:

https://github.com/NVIDIA-NeMo/RL

作者

馬立偉

NVIDIA 加速計(jì)算專家。清華大學(xué)電子科學(xué)與技術(shù)工學(xué)博士，上海交通大學(xué)信息與控制工程工學(xué)學(xué)士，研究方向包括計(jì)算機(jī)體系架構(gòu)、人工智能算法和大語言模型強(qiáng)化學(xué)習(xí)框架，持有和申請相關(guān)專利 20 余項(xiàng)。