🖥️ 多 GPU 模拟#

Genesis 支持多 GPU 执行以扩展模拟。

单 GPU 配置#

import genesis as gs

# 自动 GPU 选择
gs.init(backend=gs.gpu)

# 强制特定后端
gs.init(backend=gs.cuda)   # NVIDIA CUDA
gs.init(backend=gs.metal)  # Apple Metal
gs.init(backend=gs.cpu)    # CPU 回退

并行环境（单 GPU）#

通过在单个 GPU 上批处理环境来扩展：

scene.build(n_envs=2048, env_spacing=(1.0, 1.0))
# 所有环境在同一 GPU 上并行运行

使用多进程的多 GPU#

每 GPU 运行单独的进程：

import os
import multiprocessing

def run_simulation(gpu_id):
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    os.environ["TI_VISIBLE_DEVICE"] = str(gpu_id)
    os.environ["EGL_DEVICE_ID"] = str(gpu_id)

    import genesis as gs
    gs.init(backend=gs.gpu)
    # ... 模拟代码 ...

if __name__ == "__main__":
    for i in range(2):  # 2 个 GPU
        p = multiprocessing.Process(target=run_simulation, args=(i,))
        p.start()

分布式训练 (DDP)#

使用 PyTorch 分布式数据并行：

torchrun --standalone --nnodes=1 --nproc_per_node=2 train.py

import os
import torch
import torch.distributed as dist
import genesis as gs

local_rank = int(os.environ.get("LOCAL_RANK", 0))
os.environ["CUDA_VISIBLE_DEVICES"] = str(local_rank)
os.environ["TI_VISIBLE_DEVICE"] = str(local_rank)

gs.init(backend=gs.gpu, seed=local_rank)
scene.build(n_envs=2048)

torch.cuda.set_device(0)
dist.init_process_group(backend="nccl", init_method="env://")
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[0])

# 带梯度同步的训练循环
for step in range(steps):
    scene.step()
    loss.backward()  # DDP 处理 all-reduce
    optimizer.step()

dist.barrier()
dist.destroy_process_group()

环境变量#

变量

目的

CUDA_VISIBLE_DEVICES

PyTorch/CUDA GPU 选择

TI_VISIBLE_DEVICE

Taichi GPU 选择

EGL_DEVICE_ID

渲染 GPU (OpenGL/EGL)

对于多 GPU 设置，始终同时设置所有三个。

GPU 选择模式#

模式

方法

GPU

复杂度

单 GPU

gs.init(backend=gs.gpu)

1

低

批处理环境

scene.build(n_envs=N)

1

低

多进程

多进程 + 环境变量

N

中

分布式

torchrun + DDP

N

高

最佳实践#

  1. 优先批处理：在扩展到多 GPU 之前，先在单个 GPU 上使用大的 n_envs

  2. 设置所有环境变量：始终同时设置 CUDA、Taichi 和 EGL 设备

  3. 同步 DDP：在销毁进程组之前调用 dist.barrier()

  4. 无头渲染：在服务器上设置 pyglet.options["headless"] = True

  5. 监控内存：在批处理模拟期间使用 nvidia-smi

设备访问#

初始化后：

gs.device    # PyTorch 设备 (例如 "cuda:0", "mps:0")
gs.backend   # 后端类型 (gs.cuda, gs.metal, gs.cpu)