UFO 训练为何能把 GPU utilization 和功耗推到接近上限：链路分析

[TATP-233]

背景

我研究了 Roboparty/UFO 的 MJLab 训练链路，想记录一下它为什么在训练时能把 GPU utilization 打到接近 99%，功耗也接近电源上限。结论先说：它没有通过 nvidia-smi 锁频、修改 power limit 或显式“拉功耗”，而是由 GPU-native batched physics、大模型、大 batch、高 update-to-data ratio 和 CUDA Graph/compile 共同造成的自然结果。

1. 每张 GPU 跑 1024 个并行仿真环境

UFO 的推荐命令是多卡训练：

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
./run_train.sh \
  --agent fb \
  --gpu-ids all \
  --num-envs 1024 \
  --num-env-steps 192000000 \
  --buffer-size 5120000

其中 --num-envs 是 per GPU。--gpu-ids all 会通过 torchrunx 启动一卡一进程，每个 rank 使用 cuda:LOCAL_RANK 并初始化 NCCL。

因此 8 卡时相当于每张 GPU 独立运行：

• 1024 个 humanoid physics worlds
• 一份完整模型 / optimizer / replay buffer
• 独立 CUDA / Warp kernel 流
• NCCL 梯度同步

2. 物理仿真在 GPU 上批量执行

UFO 的环境适配层是 MJLab / MuJoCo-Warp。MJLab 的 Simulation 会调用：

mjwarp.put_data(..., nworld=self.num_envs)

也就是把 num_envs 个 world 放进 MuJoCo Warp 中批量仿真。

在 CUDA driver 和 memory pool 条件满足时，MJLab 还会自动捕获 CUDA Graph：

with wp.ScopedCapture() as capture:
    mjwarp.step(self.wp_model, self.wp_data)

之后每次 sim.step() replay 这个 graph：

wp.capture_launch(self.step_graph)

UFO 的 MuJoCo 配置是：

fps: 200
control_decimation: 4

所以每个 control step 会执行 4 次 GPU physics step。以 1024 envs/GPU 计算：

每张 GPU 每个 rollout step:
1024 worlds × 4 physics steps = 4096 batched physics world-steps

这本身已经是持续的大规模 GPU workload。

3. 真正的大头：极高的 optimizer update 密度

UFO 的训练循环近似为：

for rollout_step in ...:
    action = agent.act(obs)        # batched policy inference
    train_env.step(action)         # batched GPU simulation
    replay_buffer.extend(data)

    if past_seed and update_trigger:
        for _ in range(num_agent_updates):
            agent.update(replay_buffer, step)

关键默认值：

online_parallel_envs = 1024
update_agent_every  = 1024
batch_size          = 1024

FB:
    num_agent_updates = 16

TeCH:
    num_agent_updates = 128

因为一次 vector env step 会新增 1024 个 environment transitions，而 update_agent_every 也是 1024，所以过了 seed 阶段后：

每收集 1 个 batched rollout step:
FB   -> 16 轮完整 agent.update
TeCH -> 128 轮完整 agent.update

每一轮 agent.update() 会同时 sample expert batch 和 replay batch，各 1024，然后更新多个模块，例如 FB 中包含：

• discriminator
• forward/backward map
• critic
• auxiliary critic
• actor
• target network soft updates
• 多卡 NCCL gradient all-reduce

这是 GPU 长时间满载和功耗接近上限的主要原因。

4. 网络规模很大，而且默认不是 AMP

FB preset 中多个主要网络是 2048 宽、6 层残差结构：

forward map:    hidden_dim=2048, hidden_layers=6, num_parallel=2
actor:          hidden_dim=2048, hidden_layers=6
critic:         hidden_dim=2048, hidden_layers=6, num_parallel=2
aux critic:     hidden_dim=2048, hidden_layers=6, num_parallel=2
discriminator:  hidden_dim=1024, hidden_layers=3

按当前默认 G1 观测维度粗算：

FB:
    包含 target networks 总参数约 846M
    非 target 主网络约 441M

TeCH:
    包含 target networks 总参数约 574M
    非 target 主网络约 305M

同时模型配置里是：

amp=False

并且 workspace 里设置了：

torch.set_float32_matmul_precision("high")

也就是默认走 FP32/TF32 风格的大矩阵计算，而不是轻量混合精度训练。

num_parallel=2 还会把 forward/critic/aux critic 的有效 batch 从 1024 扩展到约 2048：

obs = obs.expand(num_parallel, -1, -1)
z = z.expand(num_parallel, -1, -1)
action = action.expand(num_parallel, -1, -1)

FB 里还有 pairwise 矩阵：

target_Ms = torch.matmul(target_Fs, target_B.T)  # [2, 1024, 1024]
Ms = torch.matmul(Fs, B.T)                      # [2, 1024, 1024]

再加上 discriminator gradient penalty 的 double backward，计算密度非常高。

5. replay/expert buffer 也在 GPU 上

UFO 训练配置中：

buffer_device = "cuda" if device.startswith("cuda") else "cpu"

replay buffer 和 expert trajectory buffer 都使用这个 device。采样和 indexing 大多留在 GPU 上，减少了 PCIe 搬运，也让 GPU kernel 队列更连续。

Rollout wrapper 中仍有一些 observation/action 的 CPU↔GPU 转换和同步，但相对大量的 update 计算来说，update 阶段通常 dominates wall time。

6. compile 与 CUDA Graph

单卡时 agent update 会启用：

torch.compile(..., mode="reduce-overhead")

多卡时 agent 侧 torch.compile 会关闭以简化分布式训练，但 MJLab/MuJoCo Warp 的物理仿真 CUDA Graph 仍会自动启用。

run_train.sh 也配置了 Triton/Inductor/CUDA/Warp cache，减少重复编译开销。

7. 关于 GPU util 99% 的解释

nvidia-smi 的 utilization.gpu 不完全等同于所有 SM 的算术利用率。它更接近采样周期内是否有 kernel 在执行。因此连续不断的小 kernel 也可能显示 99%。

但在这个案例里，确实存在大量真实重计算：

GPU-native physics
+ policy inference
+ huge residual networks
+ FP32/TF32 matmul
+ gradient penalty / double backward
+ Adam
+ NCCL all-reduce
+ GPU replay sampling

所以功耗接近 power limit 是合理结果。

总结

UFO 把 GPU 打满的实现要点可以概括为：

1. 使用 MJLab / MuJoCo Warp 做 GPU-native batched physics。
2. 默认每张 GPU 1024 个并行环境。
3. 一 GPU 一进程，使用 NCCL 同步。
4. simulation、policy、replay buffer 都放在同一 CUDA device。
5. 使用 2048 宽、多层残差网络，且 amp=False。
6. num_parallel=2 将部分网络有效 batch 扩大到约 2048。
7. 过 seed 后每个 rollout step 执行 16 或 128 轮完整 update。
8. 单卡使用 torch.compile(mode="reduce-overhead")，物理仿真使用 CUDA Graph。
9. 没有发现显式修改 GPU power limit、锁频或调用 nvidia-smi 的代码。

如果想降低功耗/占用，优先调整：

--num-envs
--num-agent-updates
--buffer-size
--gpu-ids

或修改网络宽度、层数、amp 配置；而不是从系统层面强制拉高功耗。