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

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

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

近年來(lái)，大語(yǔ)言模型 (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)度能力，成為許多開(kāi)源框架 (如 OpenRLHF、verl、AReal、ROLL、NeMo RL) 的首選基礎(chǔ)設(shè)施。

Ray 的核心優(yōu)勢(shì)包括：使用 Task 和 Actor 來(lái)抽象遠(yuǎn)端的計(jì)算任務(wù)，能夠動(dòng)態(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)練和推理引擎可以作為一個(gè)子模塊，完成具體任務(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)級(jí)性能分析工具，旨在實(shí)現(xiàn)應(yīng)用算法的可視化，找出程序中亟需優(yōu)化的“瓶頸”并進(jìn)行調(diào)整，以跨任意數(shù)量或大小的 CPU 和 GPU —— 從大型服務(wù)器到最小的系統(tǒng)級(jí)芯片 (SoC) —— 進(jìn)行高效擴(kuò)展。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

如前面介紹，NVTX 可以在程序內(nèi)放置標(biāo)記和范圍，并在 Nsight Systems 時(shí)間線上顯示。我們對(duì) verl 中每一個(gè)關(guān)鍵步驟都進(jìn)行了范圍標(biāo)記，使用戶清晰地看到每個(gè)計(jì)算子任務(wù)之間是怎樣的接力關(guān)系、時(shí)間占比如何。具體來(lái)看，我們給 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í)函數(shù)進(jìn)行了融合，使用 marked_timer 進(jìn)行計(jì)時(shí)的任務(wù)都自動(dòng)標(biāo)記了相應(yīng)的 NVTX 范圍。

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

NVIDIA 一直致力于推動(dòng)人工智能的發(fā)展，推進(jìn)算法的極限，擴(kuò)展人類認(rèn)知領(lǐng)域的邊界。主流大語(yǔ)言模型強(qiáng)化學(xué)習(xí)框架除了采用 NVIDIA 的加速計(jì)算方案，也考慮兼容不同的硬件方案。因此在 verl 集成 Nsight Systems 的設(shè)計(jì)內(nèi)，我們對(duì) 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)抽象提供一個(gè) Nsight Systems 的具體實(shí)現(xiàn)，如 NsightSystemsProfiler 等。Verl 可以在未來(lái)實(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 的時(shí)間線示例

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

常見(jiàn)的反饋問(wèn)題

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

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

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

2. 文件格式如何解讀?

文件格式為 worker_process_.nsys-rep 或者 worker_process_..nsys-rep，PID 是進(jìn)程號(hào)，RID 是第幾個(gè) capture range。

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

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

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

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

作者

馬立偉

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