NVIDIA Isaac Lab是一個(gè)適用于機(jī)器人學(xué)習(xí)的開源統(tǒng)一框架，基于 NVIDIA Isaac Sim 開發(fā)，其模塊化高保真仿真適用于各種訓(xùn)練環(huán)境，可提供各種物理 AI 功能和由 GPU 驅(qū)動(dòng)的物理仿真，縮小仿真與現(xiàn)實(shí)世界之間的差距。

Isaac Lab 主要通過 NVIDIA GPU 加速高度逼真的物理模擬、實(shí)時(shí)渲染、深度強(qiáng)化學(xué)習(xí)（RL）模仿學(xué)習(xí)（IL）集成等先進(jìn)技術(shù)，然而對于復(fù)雜的強(qiáng)化學(xué)習(xí)環(huán)境，可能需要在多個(gè) GPU 上擴(kuò)大訓(xùn)練規(guī)模。Isaac Lab 支持多 GPU 和多節(jié)點(diǎn)功能，與在單個(gè) GPU 上進(jìn)行訓(xùn)練相比，可以更快地加速訓(xùn)練過程并達(dá)到更高的性能水平。

在之前的機(jī)器人仿真教程里，我們分別介紹了Isaac Lab 的安裝教程以及Isaac Lab 的可用環(huán)境與強(qiáng)化學(xué)習(xí)腳本使用指南，本篇教程將帶大家了解如何通過多 GPU 和多節(jié)點(diǎn)進(jìn)行擴(kuò)展訓(xùn)練。

核心要點(diǎn)

Isaac Lab 支持多 GPU 和多節(jié)點(diǎn)強(qiáng)化學(xué)習(xí)。但是此功能目前僅可用于 RL-Games、RSL-RL 和 skrl 庫。多 GPU 和多節(jié)點(diǎn)訓(xùn)練僅支持 Linux 系統(tǒng)，因 NCCL 庫限制，暫不支持 Windows 系統(tǒng)。

一、多 GPU

Isaac Lab 支持以下兩種多 GPU 訓(xùn)練框架。

1. 通過 PyTorch Torchrun 進(jìn)行分布式訓(xùn)練

Torchrun 通過以下方式管理分布式訓(xùn)練：

進(jìn)程管理：為每個(gè) GPU 創(chuàng)建一個(gè)獨(dú)立進(jìn)程，并將每個(gè)進(jìn)程分配至指定的 GPU。

腳本執(zhí)行：在每個(gè)進(jìn)程上運(yùn)行相同的訓(xùn)練腳本（例如 RL Games）。

環(huán)境實(shí)例：每個(gè)進(jìn)程都會創(chuàng)建獨(dú)立的 Isaac Lab 環(huán)境實(shí)例。

梯度同步：收集并同步所有進(jìn)程的梯度，在每個(gè)訓(xùn)練步驟結(jié)束后將更新后的梯度廣播回各進(jìn)程。

此流程中的關(guān)鍵組件是：

Torchrun：處理進(jìn)程生成、通信和梯度同步。

RL 庫：運(yùn)行實(shí)際訓(xùn)練算法的強(qiáng)化學(xué)習(xí)庫。

Isaac Lab：提供每個(gè)進(jìn)程獨(dú)立實(shí)例化的模擬環(huán)境。

Torchrun 在底層使用 DistributedDataParallel 模塊來管理分布式訓(xùn)練。當(dāng)使用多個(gè) GPU 進(jìn)行訓(xùn)練時(shí)，會發(fā)生以下情況：

每個(gè) GPU 運(yùn)行一個(gè)獨(dú)立的進(jìn)程

每個(gè)進(jìn)程執(zhí)行完整的訓(xùn)練腳本

每個(gè)進(jìn)程都獨(dú)立擁有：

① Isaac Lab 環(huán)境實(shí)例（含 n 個(gè)并行環(huán)境）

② 策略網(wǎng)絡(luò)副本

③ 用于收集訓(xùn)練經(jīng)驗(yàn)的緩沖區(qū)

所有進(jìn)程僅針對梯度更新進(jìn)行同步

2. 通過 JAX 進(jìn)行分布式訓(xùn)練

在使用 JAX 時(shí)，我們借助skrl.utils.distributed.jax模塊來實(shí)現(xiàn)分布式訓(xùn)練。由于 JAX 等機(jī)器學(xué)習(xí)框架通常不會在單個(gè)程序調(diào)用中自動(dòng)啟動(dòng)多個(gè)進(jìn)程，因此 skrl 庫提供了該模塊，用于負(fù)責(zé)進(jìn)程的啟動(dòng)與管理。

*注：僅 skrl 庫支持 JAX。

二、多節(jié)點(diǎn)

除了將訓(xùn)練規(guī)模擴(kuò)展到單臺機(jī)器上的多個(gè) GPU 之外，還可以跨多個(gè)節(jié)點(diǎn)進(jìn)行訓(xùn)練。要跨多個(gè)節(jié)點(diǎn)/機(jī)器進(jìn)行訓(xùn)練，需要在每個(gè)節(jié)點(diǎn)上啟動(dòng)單獨(dú)的進(jìn)程。

運(yùn)行訓(xùn)練

接下來，我們來演示下多 GPU 多節(jié)點(diǎn)訓(xùn)練操作步驟。本期教程以 NVIDIA RTX 5880 Ada GPU 為底層硬件支持。

一、訓(xùn)練環(huán)境

1. CUDA 安裝：

1.1使用 nvidia-smi 查看當(dāng)前驅(qū)動(dòng)支持的最高 cuda 版本。如下圖所示，可以看到當(dāng)前顯示最高版本為 CUDA Version：12.8。

圖1

1.2訪問官網(wǎng)（https://developer.nvidia.com/cuda-toolkit-archive），下載并安裝 CUDA 12.8 及以下版本。

根據(jù)提示，選擇對應(yīng)系統(tǒng)版本，獲取 CUDA 工具包安裝程序下載鏈接以及安裝方式。

1.3運(yùn)行以下命令：

sudosh cuda_12.8.0_570.86.10_linux.run

2. Pytorch安裝

2.1訪問https://pytorch.org/，根據(jù)提示選擇對應(yīng)系統(tǒng)、CUDA版本等下載安裝即可。

2.2如果提示沒有 pip，按照要求 apt 安裝。

2.3執(zhí)行pip3 install torch torchvision。

2.4安裝完成后，查看 pytorch 版本。

二、運(yùn)行結(jié)果（多節(jié)點(diǎn)訓(xùn)練）

1. 在窗口中再次確認(rèn)環(huán)境變量

exportISAACSIM_PATH="${HOME}/isaacsim"
exportISAACSIM_PYTHON_EXE="${ISAACSIM_PATH}/python.sh"

2. 節(jié)點(diǎn)運(yùn)行

在節(jié)點(diǎn) 1 運(yùn)行：

${ISAACSIM_PYTHON_EXE}-m torch.distributed.run --nproc_per_node=1 --nnodes=2 --node_rank=0 --rdzv_id=123 --rdzv_backend=c10d --rdzv_endpoint=192.168.150.161:5555 scripts/reinforcement_learning/rl_games/train.py --task=Isaac-Cartpole-v0 --headless –distributed

在節(jié)點(diǎn) 2 運(yùn)行：

${ISAACSIM_PYTHON_EXE}-m torch.distributed.run --nproc_per_node=1 --nnodes=2 --node_rank=1 --rdzv_id=123 --rdzv_backend=c10d --rdzv_endpoint=192.168.150.161:5555 scripts/reinforcement_learning/rl_games/train.py --task=Isaac-Cartpole-v0 --headless –distributed

注意，需要調(diào)整的參數(shù)主要是：

--nproc_per_node

每個(gè)節(jié)點(diǎn)（機(jī)器）上啟動(dòng)的進(jìn)程數(shù)，通常設(shè)置為該節(jié)點(diǎn)的 GPU 數(shù)量。例如，若單機(jī)有 8 塊 GPU，--nproc_per_node=8。

--nnodes

參與訓(xùn)練的物理節(jié)點(diǎn)總數(shù)。例如，--nnodes=4表示使用 4 臺機(jī)器組成集群進(jìn)行訓(xùn)練。

--rdzv_endpoint

主節(jié)點(diǎn)的 IP 地址和端口號，格式為host:port。所有節(jié)點(diǎn)通過此端點(diǎn)進(jìn)行通信協(xié)調(diào)，例如--rdzv_endpoint=192.168.1.100:29500。

*如需了解詳細(xì)步驟，可參考官方文檔：

https://isaac-sim.github.io/IsaacLab/main/source/features/multi_gpu.html

推薦硬件配置

以下是麗臺針對個(gè)人開發(fā)者/研究人員的機(jī)器人仿真及訓(xùn)練環(huán)境推薦配置，能流暢運(yùn)行 Isaac Lab 以及機(jī)器人訓(xùn)練和仿真所需的算力，同時(shí)可再與企業(yè)級多卡集群方案形成互補(bǔ)，能夠完全滿足中小規(guī)模的實(shí)驗(yàn)需求。

LEADTEK WS3008

LEADTEK WS3008 產(chǎn)品特性

支持選配 2 張 NVIDIA RTX 5880 Ada GPU

支持單 Intel Xeon W-3400、W-2400 系列處理器

Intel W790 芯片組

支持最高至 350W CPU TDP

支持最多 DDR5-4800MHz x8 內(nèi)存

支持 4 PCIe 5.0 x16 插槽

支持 1 M.2 NVMe PCI-E 4.0 x4

支持 2 個(gè) 10GbE BaseT、1 個(gè) 2.5GbE BaseT 和 1 個(gè) 1GbE BaseT 網(wǎng)口

支持 1 個(gè) IPMI 管理口

支持 2 個(gè) USB-A

支持 1 個(gè) VGA 口，1 個(gè) COM 口

支持 1 個(gè) 1200W/1300W 鉑金電源

機(jī)箱體積：400.0×278.0×167.6 mm