多 GPU 仿真与训练#

Genesis World 沿两个相互独立的维度扩展。在单张 GPU 内，一个场景以批处理 并行环境 的方式同时运行同一个世界的许多副本。跨 GPU 时，每个设备一个进程，各自运行绑定到单张 GPU 的场景。本页讨论第二个维度：如何把工作分摊到多张 GPU，以及如何与数据并行训练结合。

建议只在用满第一张 GPU 之后再加第二张。单张现代 GPU 就能并行运行数千个环境，而把它们批处理起来比跨进程边界既更简单也更快。当你需要比单张 GPU 显存能容纳的更多环境，或者数据并行训练需要每个设备一个 worker 时，才需要横向扩展。

有两个可运行示例是下面这些模式的权威来源：

每 GPU 一个进程的模型#

Genesis World 不会把单个场景拆分到多张 GPU。每个进程初始化自己的 Genesis 运行时，构建自己的场景，并且只在一张设备上运行。可以运行多个这样的进程，让每个进程绑定到不同的 GPU。

用以下环境变量固定每个进程，并且要在 gs.init() 运行之前设置，因为它们决定了 Genesis、其编译器以及渲染器在初始化时绑定到哪张设备：

  • CUDA_VISIBLE_DEVICES： 限制 CUDA 运行时与 PyTorch 能看到哪些物理 GPU。将其设为单个索引，使进程只看到一张设备，之后该设备在进程内以 cuda:0 访问。

  • QD_VISIBLE_DEVICE： 为 Quadrants 选择 GPU，Quadrants 是生成并运行 Genesis World 内核的编译器。

  • EGL_DEVICE_ID： 选择用于离屏（EGL）渲染的 GPU，仅在进程需要在 GPU 上渲染图像时才有意义。

把 CUDA_VISIBLE_DEVICES 和 QD_VISIBLE_DEVICE 一起设为同一个索引，这样仿真和所有 PyTorch 张量都会落在同一张设备上。每个进程只能看到那一张 GPU，因此在进程内部始终以索引 0 引用它。

每张 GPU 运行一个进程#

多进程示例会为每张 GPU 派生一个 worker。每个 worker 设置自己的设备变量，然后运行一个普通的单 GPU Genesis 程序：

def run(gpu_id, func):
    # Pin this process to one physical GPU before Genesis initializes.
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    os.environ["QD_VISIBLE_DEVICE"] = str(gpu_id)
    os.environ["EGL_DEVICE_ID"] = str(gpu_id)
    func()

num_gpus = 2
processes = []
for i in range(num_gpus):
    p = multiprocessing.Process(target=run, args=(i, main))
    processes.append(p)
    p.start()

main() 的主体是一个普通仿真（gs.init(backend=gs.gpu)、构建场景、步进），其中没有任何与 GPU 索引相关的内容。隔离完全由父进程为每个子进程设置的环境变量实现。

这种模式适合高度并行的任务：彼此独立的 rollout、参数扫描或数据生成，这些进程从不需要交换梯度。当需要交换梯度时，请改用 DDP。

使用 PyTorch DDP 进行数据并行训练#

examples/rigid/ddp_multi_gpu.py 使用 PyTorch DistributedDataParallel 跨 GPU 训练策略。每个 rank 拥有一个完整的 Genesis 场景及其自己的环境批次；DDP 在每个 rank 上保留一份模型副本，并在每次反向传播时跨 rank 平均梯度。有效批大小为单 GPU 的 n_envs 乘以 GPU 数量，因此增加 GPU 会降低梯度噪声，而不会改变任何单个场景。

用 torchrun 启动它；torchrun 会为每张 GPU 启动一个进程，并设置 DDP 读取的会合变量：

torchrun --standalone --nnodes=1 --nproc_per_node=2 examples/rigid/ddp_multi_gpu.py

每个 worker 读取自己的 rank，把自身绑定到对应的 GPU，并按 rank 为 Genesis 设置种子，使各 GPU 上的环境彼此去相关，而不是完全相同：

local_rank = int(os.environ.get("LOCAL_RANK", 0))

os.environ["CUDA_VISIBLE_DEVICES"] = str(local_rank)
os.environ["QD_VISIBLE_DEVICE"] = str(local_rank)
gs.init(backend=gs.gpu, seed=local_rank)  # distinct seed per rank

之后构建场景并初始化进程组。由于 CUDA_VISIBLE_DEVICES 已把该进程限制到一张 GPU，这里的设备就是 cuda:0，DDP 包装器会绑定到它：

scene.build(n_envs=args.num_envs)

gpu_id = 0
torch.cuda.set_device(gpu_id)
dist.init_process_group(backend="nccl", init_method="env://")
device = torch.device("cuda", gpu_id)

model = TinyMLP(obs_dim, act_dim).to(device)
model = DDP(model, device_ids=[gpu_id])

训练循环会步进仿真，把状态读入张量，并让 DDP 同步梯度。观测直接来自刚体求解器的广义坐标，因此没有数据离开 GPU：

for step in range(args.steps):
    scene.step()
    qpos = rigid.get_qpos()  # shape (n_envs, n_qs), on the GPU

    obs = qpos + torch.randn_like(qpos)
    logits = model(obs)
    target = qpos.sum(dim=1, keepdim=True)
    loss = torch.nn.functional.mse_loss(logits, target)

    optim.zero_grad(set_to_none=True)
    loss.backward()  # DDP averages gradients across ranks
    optim.step()

最后要干净地关闭。这个 barrier 让每个 rank 在 NCCL 销毁之前到达同一点，从而避免某个 rank 提前退出时发生挂起：

dist.barrier()
dist.destroy_process_group()
gs.destroy()

把仿真状态读入张量#

在 GPU 后端上，实体状态已经驻留在设备上，因此可以直接送入模型，而无需往返主机。DDP 示例从刚体求解器读取广义位置：

rigid = scene.sim.rigid_solver
qpos = rigid.get_qpos()  # shape ([n_envs,] n_qs)

当场景以并行环境构建时，返回的张量带有前置的 n_envs 维度；单环境场景则会去掉该维度。这遵循 Genesis World 中通用的 形状约定。在 gs.init() 之后，gs.device 与 gs.backend 会报告解析得到的 PyTorch 设备与后端，便于把自己的张量放到同一张 GPU 上。

注意事项与常见坑#

Warning

在 gs.init() 之前设置设备环境变量。它们只在初始化时读取一次；之后再修改对正在运行的进程没有影响。

Note

用 EGL_DEVICE_ID 固定渲染 GPU 并非在每台机器上都可靠。正因如此，examples/rigid/ddp_multi_gpu.py 没有设置它。无头仿真或训练并不要求在 GPU 上渲染，因此除非你确实要按 rank 渲染图像并已确认硬件支持，否则应省略该变量。

Tip

在添加第二张 GPU 之前，先把单张 GPU 能容纳的环境尽量批量化。提高 n_envs 时用 nvidia-smi 观察显存，只有当单张设备已满，或者训练配方需要每个 GPU 一个 rank 时，才转向多 GPU。

另见#

  • 并行仿真：在单张 GPU 上批处理环境，以及 ([n_envs,] ...) 形状约定。

  • Hello, Genesis World：这些模式中每个进程所运行的单 GPU 程序。

  • Scene API：build、step，以及 n_envs 和 env_spacing 参数。