调研报告：GR00T-WholeBodyControl / SONIC 多卡 PPO 的 DDP 与数据并行语义

[TATP-233]

摘要

直接回答多卡训练的核心问题：

• GR00T-WholeBodyControl 的 SONIC 训练入口使用 Hugging Face Accelerate。官方单机示例是 8 个进程，官方多机示例是 8 台机器、64 个进程；普通 MULTI_GPU 配置下是一张 GPU 一个进程，actor + critic 被合成一个 wrapper 后交给 PyTorch DDP。
• num_envs 是每个 rank 的环境数，不会按 world size 拆分。每个 rank 独立创建 Isaac Lab 环境，使用 base_seed + rank，并持有自己的 [T, N, ...] rollout。若每卡环境数为 N、rollout 长度为 T、rank 数为 W，每 iteration 的全局新 transition 数是 W * N * T。
• raw observation/action/transition 和 rollout storage 不跨 rank 交换。GAE 先在本地计算；但默认会 all-gather 全部 advantage 做全局 mean/std normalization，再让每个 rank 取回自己的 slice，因此不能笼统说“训练数据完全没有 collective”。
• 稳态同步的主体是梯度，不是周期参数平均：每个需要同步的 backward 由 DDP 对 gradient buckets 做平均 all-reduce，随后各 rank 本地执行相同 optimizer step。DDP 创建时会同步参数和 registered buffers；稳态没有“每 N 轮重新广播参数”。
• 默认 SONIC 配置是 E=5 个 PPO epoch、M=4 个 mini-batch、无 micro-batch、gradient_accumulation_steps=1，所以每 iteration 有 20 次 backward/梯度同步 phase 和 20 次实际 optimizer step。DDP 会按 bucket 发起 collective，不能把它误写成“20 次单一 flat-gradient all-reduce”；实际 NCCL collective 数取决于 bucket 划分。
• adaptive KL 在每个 micro-batch 都 all-gather 一个标量，并让所有 rank 独立执行同样的 LR 调整；默认又是 20 次/iteration，没有额外的 rank-0 LR broadcast。
• 默认 critic 使用 RunningMeanStd、actor 不使用。普通 DDP 默认 broadcast_buffers=True，会在 DDP forward 边界从 rank 0 广播这些 buffers；PPO update 后代码还会显式 gather 各 rank 的 mean/variance 并简单算术平均。后一个实现不合并 count、不做 count 加权，也没有 between-group variance，源码自己称为 “crude sync”。
• adaptive motion sampling 每个 rank 每轮本地更新；默认每 200 iteration 才 gather episode/failure counts 做一次跨 rank 平均，所以两次全局同步之间各 rank 的 sampling distribution 可以不同。
• 仓库内置的 trainer 只有 PPO 和“PPO + reconstruction/latent auxiliary loss”。没有 SAC、TD3、DQN、replay buffer 或 off-policy runner。代码中虽有 use_dagger、distill_only 和几个 distillation forward mode，但没有对应的独立 trainer/config，主 PPO update 也没有调用 teacher/reference model；它们不能视为已接通的另一种分布式 RL 算法。

这是一份源码审计，不是多卡性能 benchmark。主要审计基线是本地 NVlabs/GR00T-WholeBodyControl@1983e88。审计当天远端为 c374bae；本地落后的提交只涉及 README、live-camera 文档/媒体，以及 open3d / vector_quantize_pytorch 两项依赖声明，没有修改本文涉及的训练和 distributed 代码。

本文讨论官方普通 Accelerate MULTI_GPU 路径。若另外提供 DeepSpeed/FSDP 的 Accelerate 配置，accelerator.prepare() 的包装和同步语义会变化，不能直接套用本文的 DDP collective 结论。

1. 启动方式与进程拓扑

官方训练文档给出的启动方式是：

# 单机 8 卡
accelerate launch --num_processes=8 ...

# 8 机 64 卡
accelerate launch --multi_gpu \
  --num_machines=8 --num_processes=64 \
  --machine_rank=... --main_process_ip=... --main_process_port=... ...

官方同时建议使用 64+ GPU，8 GPU 也能训练但收敛明显更慢：training.md#L118-L145。

入口创建 Accelerator 时只显式修改两个 distributed 选项：

DistributedDataParallelKwargs(find_unused_parameters=False)
InitProcessGroupKwargs(timeout=6000s)
Accelerator(
    gradient_accumulation_steps=...,
    kwargs_handlers=[ddp_kwargs, init_pg_kwargs],
)

随后以 accelerator.process_index 设置 global rank，并把每个 rank 的 seed 改为 base_seed + process_index：train_agent_trl.py#L177-L198。

每个进程各自创建 Isaac Lab environment。代码用 /tmp/isaaclab_app_launcher.lock 串行化同一节点上的 AppLauncher 初始化，目的是避免多个进程同时创建 app；锁只约束初始化临界区，不会把后续 rollout/training 串行化：train_agent_trl.py#L252-L279。

actor、critic 和可选 discriminator 被放入一个 PolicyAndValueWrapper，然后 accelerator.prepare(model, optimizer, dataloader) 负责 DDP 包装：ppo_trainer.py#L579-L583、ppo_trainer.py#L632-L653。这里训练数据来自 simulator，train_dataset=None，所以主路径没有普通监督学习 DataLoader/DistributedSampler 的分片问题。

wrapper 的目的不只是组织代码。PPO update 把 policy 和 value 放在同一个 DDP forward 里，避免同一个 DDP module 在一次 backward 前被分别 forward 多次：ppo_trainer.py#L70-L76、ppo_trainer.py#L1256-L1304。

2. 数据并行：每个 rank 一套环境和本地 rollout

trainer 直接定义：

local_batch_size = env.config.num_envs = N
batch_size       = N * world_size    = N * W

这里的 batch_size 实际按“环境轨迹条数”计，不是 transition 条数：ppo_trainer.py#L519-L545。每个 rank 的 RolloutStorage 独立分配，底层 tensor 形状是 [T, N, ...]：ppo_trainer.py#L752-L790、data_utils.py#L29-L68。

rollout 时 trainer unwrap DDP model，切到 eval mode，在 torch.no_grad() 下让本地 policy 与本地 env 交互 T 步，随后用本地 critic 计算 value/GAE；这段路径没有 raw rollout collective：ppo_trainer.py#L885-L1012、ppo_trainer.py#L1717-L1729。

因此基本样本规模是：

local new transitions / iteration  = N * T
global new transitions / iteration = W * N * T

SONIC release 默认 N=4096, T=24：sonic_release.yaml#L20-L28、sonic_release.yaml#L76-L83。所以：

• 每卡每 iteration 新采 4096 * 24 = 98,304 个 transition；
• 官方 64 卡示例每 iteration 新采 6,291,456 个 transition；
• 加卡会线性放大全局环境数/全局采样量，并不是把固定的 4096 个环境拆到多卡。

Mini-batch 是按完整环境轨迹切分

rollout 在 update 前从 [T, N, ...] 转为 [N, T, ...]，随后 randperm(N) 并按环境索引切 mini-batch：ppo_trainer.py#L1123-L1205、ppo_trainer.py#L1734-L1760。

在 N 能被 M 整除时：

local trajectories / mini-batch = N / M
local transitions  / mini-batch = (N / M) * T
global effective transitions per synchronized update = W * (N / M) * T

默认 M=4，所以每卡每次 update 使用 1024 条完整的 24-step 轨迹，即 24,576 个 transition；64 卡时一次同步 update 的全局有效样本量是 1,572,864。默认 ppo_shuffle_every_epoch=True，每个 rank 在每个 PPO epoch 都独立重洗自己的环境轨迹：ppo_im_phc.yaml#L9-L30。

代码要求 N % M == 0 才能得到预期切分，但异常被 catch 后会把 local/global mini-batch size 都退化为 1，而不是继续抛错：ppo_trainer.py#L524-L538。这是配置校验缺口，不应依赖该 fallback 做正常训练。

GAE 本地算，advantage 默认全局标准化

return/GAE 的递推使用本 rank 的 rewards/dones/values；得到 advantage 后，默认 sync_advantage_normalization=True，会：

1. all-gather 所有 rank 的 advantage；
2. 在全局 W * T * N 样本上计算 mean/std；
3. reshape 后让每个 rank 取回自己的部分。

源码见 ppo_trainer.py#L2096-L2152。因此每 iteration 有一次较大的 advantage all-gather；它交换的是派生统计输入，不会把 observation/action/return storage 变成共享 rollout。

3. 梯度同步还是参数同步

准确答案是“初始化同步模型状态，稳态同步梯度”。

对象	机制	默认周期
actor + critic 参数	DDP 构造时从 authoritative rank 同步	setup 一次
actor + critic 梯度	DDP gradient bucket 平均 all-reduce	每次 backward，同步 phase 默认 20 次/iteration
optimizer moments	无 collective，各 rank 本地维护	每次实际 optimizer step
模型参数稳态	无周期 parameter averaging/broadcast	依靠相同平均梯度、LR 和 optimizer state 保持一致
advantage	accelerator.gather 后全局 normalization	1 次/iteration
adaptive KL	标量 accelerator.gather	每个 micro-batch，默认 20 次/iteration
registered buffers	DDP 默认从 rank 0 broadcast	默认每个需要 buffer sync 的 DDP forward
running mean/variance	额外显式 gather 后简单平均	默认 1 次/iteration
adaptive sampling counts	gather 后简单平均	默认每 200 iteration
raw rollout	不同步	无

真正触发梯度 collective 的位置是 accelerator.backward(loss)；之后才做 gradient clipping、optimizer.step() 和 zero_grad()：ppo_trainer.py#L1734-L1811。

这不是像 Holosoma FastSAC 或 RSL-RL 那样把全部梯度手工 flatten 后发起一次阻塞 all-reduce。PyTorch DDP 把梯度组织为 buckets，bucket ready 时做平均 reduction，并可与 backward 重叠；参数/梯度/buffer 的基础语义见 PyTorch DDP 文档。项目没有覆盖 bucket_cap_mb；其声明的最低 Accelerate 版本中 DDP kwargs 默认 bucket_cap_mb=25、broadcast_buffers=True：gear_sonic/pyproject.toml#L93-L103、Accelerate v1.3.0 dataclasses.py#L138-L176。

各 rank 参数能持续一致，是因为它们从相同 DDP 初始参数开始，每个 step 获得相同的平均 gradient，并用同样的 optimizer state 和 LR 做相同更新；不是因为每轮重新同步 parameter tensor。在各 rank local batch size 相等时，DDP 对 rank-local mean gradient 求平均，等价于对这些 rank 的 mini-batch 并集求全局 mean gradient。

默认同步频率

默认配置：

PPO epochs E                  = 5
mini-batches M                = 4
per_device_train_batch_size   = local_mini_batch_size
micro-batches q               = 1
gradient_accumulation_steps A = 1

配置证据：ppo_im_phc.yaml#L9-L18、trl/ppo.yaml#L4-L15、ppo_trainer.py#L540-L545。

所以每个 iteration：

5 PPO epochs * 4 mini-batches = 20 backward/sync phases
                                      20 actual optimizer steps

每个 rollout transition 在每个 epoch 恰好进入一个 mini-batch，因此被重复用于 5 次 PPO update。这里的“20 次 sync phase”通常会展开为多次 bucket collective，而不是恰好 20 个 NCCL all-reduce 调用。

非默认 micro-batch / gradient accumulation

若 P=per_device_train_batch_size 严格整除 N/M，则每个 mini-batch 的 micro-batch 数：

q = (N / M) / P
backward calls / iteration = E * M * q

accelerator.accumulate(model) 在 accumulation boundary 之外使用 DDP no_sync，Accelerate 包装后的 optimizer 也会抑制真实 step() 和 zero_grad()；每 A=gradient_accumulation_steps 个 accumulate call 才同步梯度并实际更新：Accelerate v1.3.0 accelerator.py#L1061-L1103、Accelerate v1.3.0 optimizer.py#L111-L121、Accelerate v1.3.0 optimizer.py#L142-L180。

有三个边界需要特别说明：

1. accumulation counter 按每个 micro-batch 连续计数，并不天然在 PPO mini-batch 或 iteration 边界清零；若总调用数不能被 A 整除，累积窗口可能跨逻辑边界。
2. 项目在每个 micro-batch 都调用 _gradient_clipping()，没有用 accelerator.sync_gradients 做条件保护：ppo_trainer.py#L2044-L2094。因此 A>1 时，中间尚未全局同步的本地累计梯度也可能先被检查/clip；这不是标准“只在 accumulation boundary clip 一次”的语义。默认 A=1 不受影响。
3. num_micro_batches 用整数除法计算，但实际 loop 使用 range(..., P)。若 P 不整除 local mini-batch，统计 buffer 的尺寸与实际循环次数可能不一致；配置应显式保证整除。

因此 production 默认可以精确说是 20 次同步更新；非默认 accumulation 需要先修正/验证上述 boundary，不能只把次数机械写成 E*M*q/A 就认为数值语义等价。

4. 每个 iteration 的通信时序

默认 SONIC 一轮可概括为：

每个 rank 独立 rollout 24 steps
  -> 本地 value / GAE
  -> 1 次 global advantage all-gather + normalization
  -> 5 epochs * 4 mini-batches:
       DDP forward（默认同步 registered buffers）
       1 次 KL scalar all-gather
       DDP backward（若干 gradient-bucket all-reduce）
       local optimizer.step
  -> RunningMeanStd 显式同步
  -> adaptive sampling 本地重算；每 200 轮做一次全局 stats gather
  -> loss/reward 等指标 gather，rank 0 写日志/checkpoint

adaptive KL 的代码先 gather local KL mean，再由每个 rank 用同一个 global KL 执行相同的 LR 分支，因此不需要额外 LR broadcast：ppo_trainer.py#L1373-L1384、ppo_trainer.py#L2154-L2178。

所有 rank 必须按相同顺序执行这些 collective。rollout 快慢不一致不会让算法变成异步训练；快 rank 会在下一次 advantage/KL/gradient collective 等待慢 rank。

5. 可训练参数之外的状态

RunningMeanStd：DDP buffer broadcast + 显式 crude average

RunningMeanStd 把 running_mean、running_var 和 count 注册为 buffers：running_mean_std.py#L13-L38。rollout 阶段 model 为 eval，所以这些统计不会更新；PPO update 阶段 model 为 train，每个本地 micro-batch forward 会更新统计：ppo_trainer.py#L1987-L2003、running_mean_std.py#L93-L149。

普通 DDP 的 broadcast_buffers=True 会在 forward 开头把 registered buffers 从 rank 0 广播到其他 rank。默认 A=1 时，这意味着每个训练 forward 开始前，各 rank 的 normalizer state 先以 rank 0 为 authoritative copy；forward 内又会被本地 batch 更新而产生差异。它并不是“各 rank 从 iteration 开头独立累计完整一轮，末尾才首次通信”。

PPO update 后，trainer 又调用 sync_running_mean_std()：

• 前 200 iteration 无条件每轮同步；
• 之后按 sync_running_mean_std_freq；
• 仓库没有覆盖该值，fallback 是 1，所以默认仍是每轮。

证据：ppo_trainer.py#L1818-L1823、ppo_trainer.py#L1936-L1959。

显式 sync_across_gpus() 只 gather mean 和 variance，然后分别对 rank 维做简单平均：running_mean_std.py#L53-L77。其边界是：

• 没有把 count 纳入显式 merge；count 仍会随默认 DDP buffer broadcast 从 rank 0 传播；
• 没有按各 rank count 加权；
• 没有加入各 rank mean 差异对应的 between-group variance；
• 因而不是严格的全局 moment merge，源码注释也明确称为 “crude sync”。

SONIC universal-token actor 配置 running_mean_std: false，critic 配置为 true，所以默认主要是 critic normalizer 走上述路径：all_mlp_v1.yaml#L22-L30、critics/mlp.yaml#L1-L13。

Adaptive motion sampling

每轮 update 后都会重新计算 sampling probability，但只有在 (global_step + 1) % 200 == 0 时才设置 sync_across_gpus=True：ppo_trainer.py#L1961-L1973。

同步实现把 episode counts 和 failure counts 拼接后 gather，对 rank 维简单平均，再计算 failure rate/sampling probability：motion_lib_base.py#L2503-L2533。所以默认 199 个 iteration 窗口内，各 rank 会根据自己的失败分布使用不同 sampler；第 200 轮才重新对齐统计。

日志口径不是完全统一

PPO loss、KL、clip fraction、value loss、entropy 和 auxiliary loss 使用 gather_for_metrics 做全局平均：ppo_trainer.py#L1569-L1633、ppo_trainer_aux_loss.py#L268-L304。reward/episode length 也 gather rank-local scalar 后取平均：ppo_trainer.py#L1840-L1863。

但部分 episode info / env tensor 指标直接在本地生成，再由 rank 0 写 W&B；W&B callback 本身只允许 world-process-zero 写入：ppo_trainer.py#L1864-L1908、wandb_callback.py#L15-L19。因此日志是“核心 PPO metric 多数全局聚合，部分 env metric 仍是 rank-0 本地视角”，不能一概而论。

6. Checkpoint / resume

只有 world-process-zero 保存 checkpoint。默认每 50 step 覆盖 last.pt，常规 checkpoint 频率由 callback 配置决定：model_save_callback.py#L16-L31、model_save_callback.py#L56-L88。

checkpoint 包含 policy、value model、optimizer、LR scheduler、trainer state/args、rank 0 的 env state（当前 wrapper 只返回 motion-lib adaptive-sampling state），以及可选 discriminator state。保存逻辑见 model_save_callback.py#L117-L160，env state 边界见 manager_env_wrapper.py#L1072-L1082。

load 发生在 accelerator.prepare() 之后；每个 rank 各自 torch.load 同一个路径，unwrap DDP model 后加载 policy/value，并在 resume 时加载 optimizer、scheduler、env 和 trainer state：ppo_trainer.py#L390-L421、ppo_trainer.py#L2180-L2245。load 后没有再从 rank 0 broadcast，一致性依赖所有 rank 解析并读取相同 checkpoint。

checkpoint 不保存 rank-local rollout storage、完整 simulator state 或完整 RNG state，所以 resume 是训练状态恢复，不是逐 transition 的 bitwise continuation。

7. 其他算法：没有 SAC/off-policy production path

当前 trainer 目录只有 TRLPPOTrainer 和继承它的 TRLAuxLossPPOTrainer。Hydra trainer 配置也只指向这两个 class：trainer/trl.yaml、trainer/trl_ppo_aux.yaml。aux trainer 只是把 reconstruction/latent loss 加到同一个 PPO loss 后一起 backward，仍使用同一份 on-policy rollout 和同一套 DDP 同步：ppo_trainer_aux_loss.py#L6-L31、ppo_trainer_aux_loss.py#L158-L196。

全仓搜索没有 SAC、TD3、DQN、replay buffer 或 off-policy runner/storage。出现的 “replay” 是传感器/策略回放，不是 RL replay buffer。

仓库确实保留了一些 distillation/DAgger 脚手架：

• use_dagger 时入口会实例化 teacher_actor 并作为 ref_model 传给 trainer；
• wrapper 支持 policy_distill / policy_distill_ppo forward mode；
• distill_only 会阻止创建 value model。

证据：train_agent_trl.py#L381-L409、train_agent_trl.py#L451-L469、ppo_trainer.py#L134-L180。

但当前主 update 固定请求 modes=["policy", "value"]，并固定计算 PPO/value loss；trainer 内 ref_model 初始化后也没有被 rollout/update 调用。仓库配置中没有 use_dagger、distill_only 或 teacher_actor 的可运行 experiment。尤其 distill_only=True 省略 value model，而 wrapper 的 "value" mode 会直接调用 self.value_model.evaluate()：ppo_trainer.py#L222-L223、ppo_trainer.py#L1288-L1304。

所以当前能下的结论是：production training path 是 PPO / PPO+aux；distillation/DAgger 是尚未接通的扩展痕迹，更不是 SAC 式 off-policy RL。

8. 与前两份调研的对照及对 UniLab 的启示

与 Holosoma FastSAC #975 和 MJLab / RSL-RL #976 对照：

维度	GR00T SONIC PPO	MJLab / RSL-RL PPO	Holosoma FastSAC
num_envs	每 rank 数量	每 rank 数量	全局数量，再拆到 rank
数据容器	rank-local rollout	rank-local rollout	rank-local replay shard
raw transition 通信	无	无	无
advantage normalization	默认全局 gather	rank-local	不适用
主梯度同步	DDP buckets	手写 flat-gradient all-reduce	手写 flat-gradient all-reduce
默认同步节奏	20 个 PPO backward phase/iteration	20 次 flat all-reduce/iteration	每个 actor/critic/alpha optimizer update
参数广播	DDP setup	learn() startup	setup
normalizer	DDP buffer broadcast + 显式 crude average	rank-local	全局 moments all-reduce
off-policy	无	无	SAC

三者共同点仍是 replicated learner + rank-local simulator/data + synchronous gradient averaging，但状态边界差异很大，不能只看“都做 all-reduce”就认为 execution semantics 相同。

对 UniLab 后续设计最直接的启示：

1. API 必须明确 num_envs 是 global budget 还是 per-rank budget；两者扩卡后的样本量完全不同。
2. “不交换 rollout”不等于“除了 gradient 没有通信”；advantage normalization、KL、normalizer、sampler 和 metrics 都可能有不同频率的 collective。
3. 如果使用 DDP，必须明确 broadcast_buffers contract。若 normalizer 需要严格全局统计，应同步 sufficient statistics（count/sum/squared-sum 或正确 merge moments），不要同时依赖默认 rank-0 buffer broadcast 和末尾简单平均。
4. gradient accumulation 的同步边界、clip 时机、PPO mini-batch 边界必须共同设计；只套 no_sync 不足以保证数值语义正确。
5. checkpoint 和日志需要明确代表全局状态还是 rank-0 snapshot，尤其是 sampler/curriculum/simulator 等非参数状态。
6. 若未来支持 off-policy 多 learner，仍需另行定义 replay ownership、跨 rank sampling、actor/critic/temperature/target 的独立同步周期；本仓库的 PPO DDP 不能直接回答这些问题。

仓库没有 distributed training E2E/numerical-equivalence test；gear_sonic/tests 目前只有 input-reader 测试。因此本文结论来自静态源码和依赖语义审计，尚未由多卡 benchmark 验证通信量、scaling efficiency 或跨 rank 数值一致性。

一句话定性：GR00T-WholeBodyControl / SONIC 是每卡独立仿真与 rollout、DDP 每个 PPO update 同步平均梯度的同步数据并行；默认每轮 20 个梯度同步 phase，不做周期参数平均。它额外全局同步 advantage、KL、normalizer 和周期 sampler 状态，且当前没有 SAC/off-policy trainer。