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

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

大語言模型強化學習與 Ray

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

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

NVIDIA Nsight Systems 工具概述

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

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

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

標記事件(markers)：在程序執(zhí)行的特定點標記信息；

注釋范圍 (ranges)：在程序執(zhí)行的兩個點之間標注范圍；

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

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

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

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

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

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

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

但是這種方法在捕捉以 Ray 為基礎(chǔ)框架的程序時卻失效了。其根本原因在于，在 Ray 任務(wù)里，上面命令行里的 application 并非真正的計算程序，而是給 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)心的計算任務(wù)是由 Ray 在遠端節(jié)點啟動的。以上兩種方法作為 application 附在 “nsys” 命令后面的時候，Nsight Systems 只能捕捉到提交任務(wù)這個前導程序，而不是真正的計算任務(wù)。

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

理解 Ray 與 Nsight Systems 集成的挑戰(zhàn)和背后的原因，解決方案便較為明確。實際上，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)的進程（可以參考 Ray 的使用手冊獲得詳細信息）。

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

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

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

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

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

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

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

這就需要我們在跟蹤強化學習任務(wù)的時候，分別跟蹤 controller 和 worker 進程。

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

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

二、細粒度跟蹤 worker 進程

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

1. 目標訓練步的選擇

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

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

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

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

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

三、使用 NVTX 描繪整體運行流程

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

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

NVIDIA 一直致力于推動人工智能的發(fā)展，推進算法的極限，擴展人類認知領(lǐng)域的邊界。主流大語言模型強化學習框架除了采用 NVIDIA 的加速計算方案，也考慮兼容不同的硬件方案。因此在 verl 集成 Nsight Systems 的設(shè)計內(nèi)，我們對 Nsight Systems 的使用方式和 NVTX 的主要函數(shù)接口進行了標準抽象，如 ProfilerConfig, DistProfiler, DistProfilerExtension, mark_range_start, mark_range_end, mark_annotate 等，并提供了空白實現(xiàn)。我們?yōu)闃藴食橄筇峁┮粋€ Nsight Systems 的具體實現(xiàn)，如 NsightSystemsProfiler 等。Verl 可以在未來實現(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 進程顯示訓練一共持續(xù)了 6 步，選取的兩個 worker 進程詳細地跟蹤了其中的第 1、2、5 步。有了詳細的 GPU 活動記錄，對齊到整體的控制流圖，我們就可以有針對性的分析程序性能了。

常見的反饋問題

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

1. 去哪里找跟蹤到的數(shù)據(jù)結(jié)果，這個目錄可以指定嗎?

按照 Ray 的使用手冊，跟蹤到的數(shù)據(jù)庫保存在每個服務(wù)器節(jié)點的 /tmp/ray/session_*/logs/{profiler_name}，并且跟蹤目錄是不能修改的。這是當前版本 Ray 的一個局限。

2. 文件格式如何解讀?

文件格式為 worker_process_.nsys-rep 或者 worker_process_..nsys-rep，PID 是進程號，RID 是第幾個 capture range。

3. 去哪里找 controller 的跟蹤數(shù)據(jù)?

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

4. 怎么查看這么多數(shù)據(jù)?

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

作者

馬立偉

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