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

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

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

核心要點

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

一、多 GPU

Isaac Lab 支持以下兩種多 GPU 訓練框架。

1. 通過 PyTorch Torchrun 進行分布式訓練

Torchrun 通過以下方式管理分布式訓練：

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

腳本執(zhí)行：在每個進程上運行相同的訓練腳本（例如 RL Games）。

環(huán)境實例：每個進程都會創(chuàng)建獨立的 Isaac Lab 環(huán)境實例。

梯度同步：收集并同步所有進程的梯度，在每個訓練步驟結束后將更新后的梯度廣播回各進程。

此流程中的關鍵組件是：

Torchrun：處理進程生成、通信和梯度同步。

RL 庫：運行實際訓練算法的強化學習庫。

Isaac Lab：提供每個進程獨立實例化的模擬環(huán)境。

Torchrun 在底層使用 DistributedDataParallel 模塊來管理分布式訓練。當使用多個 GPU 進行訓練時，會發(fā)生以下情況：

每個 GPU 運行一個獨立的進程

每個進程執(zhí)行完整的訓練腳本

每個進程都獨立擁有：

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

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

③ 用于收集訓練經驗的緩沖區(qū)

所有進程僅針對梯度更新進行同步

2. 通過 JAX 進行分布式訓練

在使用 JAX 時，我們借助skrl.utils.distributed.jax模塊來實現(xiàn)分布式訓練。由于 JAX 等機器學習框架通常不會在單個程序調用中自動啟動多個進程，因此 skrl 庫提供了該模塊，用于負責進程的啟動與管理。

*注：僅 skrl 庫支持 JAX。

二、多節(jié)點

除了將訓練規(guī)模擴展到單臺機器上的多個 GPU 之外，還可以跨多個節(jié)點進行訓練。要跨多個節(jié)點/機器進行訓練，需要在每個節(jié)點上啟動單獨的進程。

運行訓練

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

一、訓練環(huán)境

1. CUDA 安裝：

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

圖1

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

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

1.3運行以下命令：

sudosh cuda_12.8.0_570.86.10_linux.run

2. Pytorch安裝

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

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

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

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

二、運行結果（多節(jié)點訓練）

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

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

2. 節(jié)點運行

在節(jié)點 1 運行：

${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é)點 2 運行：

${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

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

--nproc_per_node

每個節(jié)點（機器）上啟動的進程數(shù)，通常設置為該節(jié)點的 GPU 數(shù)量。例如，若單機有 8 塊 GPU，--nproc_per_node=8。

--nnodes

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

--rdzv_endpoint

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

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

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

推薦硬件配置

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

LEADTEK WS3008

LEADTEK WS3008 產品特性

支持選配 2 張 NVIDIA RTX 5880 Ada GPU

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

Intel W790 芯片組

支持最高至 350W CPU TDP

支持最多 DDR5-4800MHz x8 內存

支持 4 PCIe 5.0 x16 插槽

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

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

支持 1 個 IPMI 管理口

支持 2 個 USB-A

支持 1 個 VGA 口，1 個 COM 口

支持 1 個 1200W/1300W 鉑金電源

機箱體積：400.0×278.0×167.6 mm