多 GPU 和多节点训练#

通过 train_multigpu ,可将一个强化学习作业扩展到工作站的多个 GPU 或多个节点。Isaac Lab 为每个 GPU 启动一个训练进程,各进程拥有自己的仿真环境,并在进程间同步策略更新。

同一种启动方式也用于三种多 GPU 基准测试,分别测量启动、仿真和端到端训练性能。

注意

多 GPU 和多节点训练需要 Linux 及 NVIDIA NCCL,不支持 Windows。

警告

train_multigpu 属于实验性功能,后续版本可能变化。

启动多 GPU 训练#

先验证任务能否在单个 GPU 上训练,以区分任务/配置问题与分布式启动问题:

uv run isaaclab train --task Isaac-Cartpole

然后在所有可见 GPU 上运行同一任务:

uv run isaaclab train_multigpu --task Isaac-Cartpole

这就完成了从单 GPU 运行的切换。 train_multigpu 自动添加 --distributed ,并选择分布式启动器。其他参数与 train 接受的参数相同:

uv run isaaclab train_multigpu \
   --task Isaac-Reorient-KukaAllegro \
   --num_envs 4096 \
   --max_iterations 100

--num_envs 是 每个 GPU 上的 环境数量。使用四个 GPU 和 --num_envs 4096 时,作业总共从 16,384 个环境采集经验。

所有 rank 写入同一个运行目录。各 rank 在自己的 GPU 上训练,种子为启动种子加自身 rank。rank 0 将设置( params/env.yaml 和 params/agent.yaml )及视频写入与单 GPU 运行相同的位置;其余 rank 将自己的 params/ 和视频写入 rank_<rank>/ 。

小技巧

加入 --dry_run 可打印解析后的启动命令,而不开始训练,适合检查 GPU 数量、会合选项和转发的训练参数。

选择 GPU#

启动器默认使用所有可见 GPU。需要指定工作进程数量时,设置 --num_gpus :

uv run isaaclab train_multigpu --num_gpus 2 --task Isaac-Cartpole

使用 CUDA_VISIBLE_DEVICES 选择物理设备,启动器只能看到列表中的设备:

CUDA_VISIBLE_DEVICES=1,3 uv run isaaclab train_multigpu \
   --num_gpus 2 --task Isaac-Cartpole

启动前运行 nvidia-smi ,确认预期设备可用且有足够空闲显存。

选择强化学习库#

多 GPU 训练支持 RSL-RL、RL-Games 和 skrl。大多数核心任务默认使用 RSL-RL。

库

分布式后端

命令

RSL-RL

PyTorch

uv run isaaclab train_multigpu --rl_library rsl_rl ...

RL-Games

PyTorch

uv run --extra rl-games isaaclab train_multigpu --rl_library rl_games ...

skrl

PyTorch

uv run --extra skrl isaaclab train_multigpu --rl_library skrl ...

skrl

JAX

uv run --extra skrl-jax isaaclab train_multigpu --rl_library skrl --ml_framework jax ...

skrl 使用 JAX 时采用 skrl 的分布式启动器,而非 torchrun 。请为 --num_gpus 传入整数,并通过 --coordinator_address 配置协调器:

uv run --extra skrl-jax isaaclab train_multigpu \
   --rl_library skrl --ml_framework jax --num_gpus 4 \
   --coordinator_address localhost:5000 \
   --task Isaac-Cartpole

通过三种多 GPU 基准测试测量扩展性#

如果目标是测量性能,而非训练供后续使用的策略,请使用基准测试命令。Isaac Lab 提供三种多 GPU 工作流程:

基准测试

测量内容

可回答的问题

startup_multigpu

所有 GPU 同时启动相同负载时,rank 0 的启动时间。

满负载节点如何影响启动?

runtime_multigpu

各 GPU 独立运行时,rank 0 的仿真吞吐量。

主机资源竞争如何影响环境步进?

training_multigpu

所有 rank 全局同步的训练吞吐量。

端到端学习在多 GPU 上的扩展效果如何?

运行各基准测试时,使用与 train_multigpu 相同的启动选项:

uv run isaaclab benchmark startup_multigpu \
   --num_gpus 2 --task Isaac-Cartpole

uv run isaaclab benchmark runtime_multigpu \
   --num_gpus 2 --task Isaac-Cartpole --num_envs 4096

uv run isaaclab benchmark training_multigpu \
   --rl_library rsl_rl --num_gpus 2 \
   --task Isaac-Cartpole --num_envs 4096 --max_iterations 100

training_multigpu 支持基于 PyTorch 的 RSL-RL、RL-Games 和 skrl,不支持基于 JAX 的 skrl 或 SB3。它也不接受 --video 和 --capture_env_sensors ,因为这些选项不能产生有意义的跨 rank 汇总结果。RSL-RL 和 RL-Games 支持 --check_success :成功指标在全部 rank 的环境中取平均,使所有 rank 在同一迭代停止。

多节点基准测试中,在每个节点传入 跨节点训练 所述的 --nnodes 、 --node_rank 和会合选项。

读取多 GPU 基准测试结果#

只有全局 rank 0 写出结果包。 extra 字段记录 rank 布局,并说明结果包覆盖的测量范围:

extra 字段

含义

world_size, local_world_size, num_nodes

作业的 rank 布局。

num_envs_per_rank

每个 rank 承载的环境数量。

workload_scope

training_multigpu 为 global ,因为各 rank 同步训练。 startup_multigpu 和 runtime_multigpu 为 rank0 ,因为各 rank 独立运行,且只测量 rank 0。

measurement_scope

rank0_process :时间、学习曲线、CPU 和 RAM 数据仅来自 rank 0。

gpu_measurement_scope

rank0_node : resources.devices 报告 rank 0 可见的全部 GPU。GPU 利用率和显存值仍仅对应 rank 0 的设备。

比较单 GPU 和多 GPU 时,保持 每个 rank 的 --num_envs 不变。在相同逐 rank 设置下,N 个 GPU 的作业处理的环境数为单 GPU 作业的 N 倍。应将 training_multigpu 的全局吞吐量与单 GPU 吞吐量的 N 倍比较。不要将 startup_multigpu 或 runtime_multigpu 当作汇总速率;它们测量的是其他 rank 造成主机资源竞争时的 rank 0。

多 GPU 训练的工作原理#

在 PyTorch 工作流程中, train_multigpu 包装 torchrun ,为每个 GPU 启动一个进程。各进程拥有:

  • 一个 Isaac Lab 应用及其向量化环境;

  • 一份策略副本;

  • 一个轨迹采样缓冲区;以及

  • 由本地 rank 选择的一个 GPU。

各进程独立采集经验,在策略更新期间通过 DistributedDataParallel 同步梯度。仿真不会在 GPU 之间迁移,因此 GPU 数量增加时,主机 CPU、RAM 和 I/O 可能成为限制因素。

每个 GPU 对应一个训练进程和仿真负载 每个 GPU 对应一个训练进程和仿真负载

读取分布式运行日志#

各 rank 产生类似的启动消息、警告和模型摘要。启动器默认显示本地 rank 0,使控制台易于阅读。训练指标本就来自全局 rank 0;即使隐藏的 rank 崩溃,仍会报告出错 rank 及其异常堆栈。

进程表现似乎不一致时,显示所有 rank 的输出:

uv run isaaclab train_multigpu \
   --task Isaac-Cartpole --log_all_ranks

要保持控制台简洁,并在磁盘上保存完整的逐 rank 日志,使用 torchrun 日志重定向:

uv run isaaclab train_multigpu \
   --task Isaac-Cartpole --tee 3 --log_dir /tmp/isaaclab-rank-logs

日志过滤选项适用于 PyTorch 工作流程。使用 JAX 的 skrl 会将所有 rank 输出到控制台。

跨节点训练#

各节点必须使用相同的 Isaac Lab 代码、依赖和任务配置,并能访问训练资产。节点间还必须能够通过会合端口通信。

选择一个节点作为会合主机。对于每节点四个 GPU 的双节点 PyTorch 作业,在第一个节点运行:

uv run isaaclab train_multigpu \
   --nnodes 2 --node_rank 0 --num_gpus 4 \
   --master_addr 10.0.0.10 --master_port 29500 \
   --task Isaac-Cartpole

在第二个节点使用其自身 rank 运行相同命令:

uv run isaaclab train_multigpu \
   --nnodes 2 --node_rank 1 --num_gpus 4 \
   --master_addr 10.0.0.10 --master_port 29500 \
   --task Isaac-Cartpole

总进程数(world size)为 nnodes * num_gpus ,此例为八个 rank。也可以通过 --rdzv_backend 、 --rdzv_endpoint 和 --rdzv_id 配置弹性 torchrun 会合。先加入 --dry_run ,验证各节点命令。

各节点按自身启动时间命名运行目录。节点共享存储时,应在每个节点的 train_multigpu 命令中传入相同的 --run_timestamp (格式为 YYYY-MM-DD_HH-MM-SS ),使所有 rank 写入同一个目录。

使用 JAX 的 skrl 时,在各节点传入 --nnodes 、 --node_rank 、整数 --num_gpus 及相同的 --coordinator_address 。不要向 JAX 启动命令传入 PyTorch 会合选项。

多节点扩展性很大程度上取决于节点间网络。如果梯度同步主导训练迭代,多节点作业可能比单节点更慢。

排查分布式训练问题#

从最小而有效的诊断开始:

  1. 确认相同任务、后端和训练参数能够通过 isaaclab train 在单 GPU 上运行。

  2. 加入 --dry_run ,检查所选 GPU 和节点数量。

  3. 分别以总进程数 2、3、4 重试。只有某个进程数失败时,通常应检查通信传输层,而非任务。

  4. 通过 nvidia-smi topo -m 检查 GPU 布局和互连。

  5. 设置 NCCL_DEBUG=INFO ,查看选用的 NCCL 传输方式。

  6. 每次只应用一种解决办法,并验证移除后故障是否再次出现。

NCCL 挂起和错误#

运行停止、没有异常堆栈,而所有参与的 GPU 利用率仍为 100% 时,通常是卡在 NCCL 集合通信中。以下办法针对已知的特定系统传输问题:

现象

可尝试的方法

没有 NVLink 的 PCIe 系统在总进程数为 2 时挂起。

NCCL_P2P_DISABLE=1 或 NCCL_P2P_LEVEL=LOC

ProcessGroupNCCL 中出现 illegal memory access 。

NCCL_SHM_DISABLE=1

CUDA 与渲染器枚举 GPU 的顺序不同,导致启用渲染的作业失败。

CUDA_DEVICE_ORDER=PCI_BUS_ID

启用渲染的作业跨 NUMA 节点执行 BROADCAST 或 ALLREDUCE 时超时。

NCCL_CUMEM_HOST_ENABLE=0 ;必要时尝试 NCCL_CUMEM_ENABLE=0 。

通信器初始化或传输故障持续存在。

NCCL_IB_DISABLE=1 或 NCCL_ALGO=Ring 。

例如,不修改共享配置,仅测试第一种办法:

NCCL_P2P_DISABLE=1 uv run isaaclab train_multigpu \
   --num_gpus 2 --task Isaac-Cartpole

这些变量可能降低通信性能,应仅在受影响机器上设置。 NCCL_P2P_DISABLE=1 与 NCCL_P2P_LEVEL=LOC 二选一,不要同时设置。两者都会阻止直接 P2P 通信,可能迫使 NCCL 选择其他传输方式而降低带宽。除非所有受支持系统都需要某项设置,否则不要将它提交到任务或启动器中。设置 CUDA_DEVICE_ORDER=PCI_BUS_ID 前,通过 nvidia-smi --query-gpu=name,pci.bus_id 查看 GPU 总线 ID。

脱离 Isaac Lab 定位挂起问题

以发生挂起的总进程数运行最小 NCCL 集合通信测试,保存为 nccl_probe.py :

import os

import torch
import torch.distributed as dist

local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
dist.init_process_group("nccl")
tensor = torch.ones(1024, device=f"cuda:{local_rank}")
dist.all_reduce(tensor)
torch.cuda.synchronize()
print(f"rank {dist.get_rank()} ok", flush=True)
dist.destroy_process_group()

使用相同 rank 数启动测试:

uv run python -m torch.distributed.run --nproc_per_node 2 nccl_probe.py

如果此测试也挂起,问题位于 NCCL 或系统拓扑,而非 Isaac Lab、任务或强化学习库。