调研报告：Holosoma FastSAC 多卡训练的数据并行与同步语义

[TATP-233]

摘要

直接回答标题中的几个问题：

• Holosoma FastSAC 采用 一张 GPU 一个进程、模型按 rank 复制、环境与 replay 按 rank 分片 的同步数据并行。
• 每个 rank 只采集、存储和采样自己的 transition；不 all-gather 数据，也不共享 replay buffer。
• 稳态训练同步的是 梯度：critic、actor 和自动温度系数 log_alpha 都在各自 optimizer step 前做阻塞式 all_reduce(SUM) / world_size。
• 参数同步只发生在初始化：rank 0 把 actor、critic 参数和 log_alpha 广播给其他 rank。之后没有“每 N 步广播一次参数”；各 rank 依靠相同初值、相同平均梯度和相同 optimizer 状态保持参数一致。
• 同步周期不是“一轮 rollout 后统一同步一次”，而是 每一个实际 optimizer update 都同步一次梯度。默认每个环境 step 做 8 次 critic/alpha 更新、2 次 actor 更新，因此默认对应 8 次 critic 梯度 all-reduce、8 次 alpha 标量 all-reduce、2 次 actor 梯度 all-reduce。
• 它没有用 DistributedDataParallel 包模型，而是手写 collective：每个网络把所有已有梯度拼成一个 flat tensor，一次 NCCL all-reduce 后再拷回各参数梯度。

这是一份源码审计，不是多卡性能 benchmark。相关结论基于本地 6e146b0；同时核对了 2026-08-14 的远端 main@6761e2b，两者之间的提交没有改动本文涉及的 FastSAC、训练入口或分布式代码。

审计基线：amazon-far/holosoma@6761e2b。

1. 启动与进程拓扑

官方示例用 torchrun --nproc_per_node=4 启动 4 个进程，并把 training.num_envs 定义为全局环境数：README.md#L211-L219。

每个进程从 WORLD_SIZE / RANK / LOCAL_RANK 读取拓扑，默认使用 NCCL 初始化 process group，然后绑定 cuda:LOCAL_RANK：train_agent.py#L59-L97。

这里有一个容易混淆的入口语义：普通 train_agent.py 是否进入多卡路径由 WORLD_SIZE > 1 决定，也就是要由 torchrun 建立进程；training.multigpu 本身不会在这个入口里自动 spawn。nightly wrapper 才会根据该配置重新用固定 4 卡 torchrun 拉起：nightly.py#L68-L88。

配置的全局环境数 E 会在每个 rank 上做整数除法：

local_envs = floor(E / W)
actual_global_envs = W * local_envs

源码直接使用 num_envs // world_size，没有分配余数，也没有整除校验：train_agent.py#L245-L260。每个 rank 的 seed 还会加上 global rank，因此在相同模型下产生不同的环境随机性和动作采样：train_agent.py#L195-L199。

2. 数据如何并行：本地 replay shard，不传 transition

每个 rank 都创建自己的 SimpleReplayBuffer，形状以本地 env.num_envs 为第一维，容量是每个环境 buffer_size 个时间槽：fast_sac_agent.py#L339-L348、fast_sac_utils.py#L15-L53。

每个外层迭代中，各 rank 独立执行一个 vectorized env step，把本地 transition 写进本地 buffer；这里没有任何 distributed collective：fast_sac_agent.py#L681-L731。采样也只访问本地 tensor，而且不是在扁平 replay 上一次抽样：它为每个本地环境分别抽取同样数量的随机时间索引，再 reshape 成本地 batch：fast_sac_utils.py#L80-L110。

因此数据流是：

rank r: local envs -> local replay shard -> local minibatch -> local backward
                                                        |
all ranks:                                  gradient all-reduce
                                                        |
rank r:                                      local optimizer.step

全局 batch 的实际计算

源码把 args.batch_size 注释为 global batch，但先算“每个环境抽几条”：

k = max(args.batch_size // local_envs // world_size, 1)

随后每个 rank 的单次 update batch 是 local_envs * k，所有 rank 合起来是：

B_effective = actual_global_envs * max(floor(B_config / actual_global_envs), 1)

证据：fast_sac_agent.py#L730-L752、fast_sac_agent.py#L537-L624。

所以只有在环境数和 batch 配置满足整除关系时，实际全局 batch 才严格等于配置值。例如框架默认 E=4096、FastSAC 默认 B=8192，在 W=4 时每卡 1024 个环境、每环境抽 2 条，每卡 batch 2048、全局正好 8192。如果 B < actual_global_envs，由于下限为 1，实际 batch 会至少等于实际全局环境数。配置证据：experiment.py#L55-L73、config_values/algo.py#L64-L109。注册的 FastSAC 默认还开启 symmetry，准备 batch 时会把每次 update 的样本数乘 2。

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

答案是“两者都有，但阶段不同”：

对象	机制	周期
actor 参数	rank 0 broadcast	setup 时一次
critic 参数	rank 0 broadcast	setup 时一次
log_alpha 值	rank 0 broadcast	setup 时一次
target critic	从已同步 critic 本地 load_state_dict	setup 时一次；之后不广播
critic 梯度	flat tensor all_reduce(SUM) / W	每次 critic update
actor 梯度	flat tensor all_reduce(SUM) / W	每次 actor update
log_alpha 梯度	标量 all_reduce(SUM) / W	每次 critic update，且 use_autotune=true 时
optimizer moments	无 collective	各 rank 对相同平均梯度执行相同 step，隐式保持一致
target critic 参数	本地 Polyak update	每次 critic update，无 collective

初始化广播见 fast_sac_agent.py#L354-L380。手写梯度同步会把一个网络的全部现有梯度 flatten 成单个 tensor，只发起一次阻塞式 all-reduce，然后除以 world size 并 scatter 回各梯度 tensor：fast_sac_agent.py#L382-L402。

critic 与 alpha 的调用位置在 fast_sac_agent.py#L446-L478，actor 的调用位置在 fast_sac_agent.py#L513-L529。all-reduce 发生在 AMP backward() 之后、unscale_ 和 gradient clipping 之前；也就是说通信的是 scaled gradient，随后每个 rank 本地 unscale、clip、step。

这不是定期参数平均，也不是 parameter-server 模式。同步 all-reduce 默认 async_op=False，因此每次 optimizer update 都是 rank 间 rendezvous；慢 rank 会在该 collective 处拖住其他 rank。参数之所以能持续一致，是因为所有 rank：

1. 从相同参数和 optimizer 状态开始；
2. 对各自 minibatch 求本地梯度；
3. 在 step 前拿到完全相同的平均梯度；
4. 各自在相同状态上执行同一个 AdamW step。

在各 rank 本地 batch 大小相同的前提下，这个“平均 rank-local mean gradient”等价于对各 rank minibatch 并集求全局 mean gradient。

4. 默认配置下的精确同步频率

学习开始条件是 global_step > learning_starts。进入稳态后，每个外层环境 step 会先一次性采一个大 batch，再拆成 num_updates 个 update batch：fast_sac_agent.py#L730-L752。

注册配置默认：

num_updates = 8
policy_frequency = 4
use_autotune = true

因此每个外层环境 step：

• critic 更新 8 次，每次 1 个 critic flat-gradient all-reduce；
• alpha 更新 8 次，每次 1 个标量-gradient all-reduce；
• actor 判断条件实际写的是 i % policy_frequency == 1，所以 i=1,5 更新，共 2 个 actor flat-gradient all-reduce；
• target critic 在每个 inner update 后本地 Polyak 更新，共 8 次，无通信。

若 num_updates == 1，actor 才改用 global_step % policy_frequency == 0。这些条件和 target 更新位置见 fast_sac_agent.py#L738-L775。因此“actor 每 4 次更新同步”只是默认配置下的直观描述，准确语义应以这两个分支为准。

5. 模型参数之外还有哪些同步

Observation normalization

默认启用的 EmpiricalNormalization 会对本地 sum 和 squared-sum 做 all-reduce，以全局 batch 更新 mean/variance/count：fast_sac_utils.py#L274-L329。

训练循环在 rollout inference 时传 update=False，所以采集阶段不更新统计；在大 batch 拆分前依次 normalize obs / next_obs / critic_obs / next_critic_obs，默认每个外层学习迭代共触发 4 次统计 all-reduce，而不是每个 inner update 都触发：fast_sac_agent.py#L537-L595、fast_sac_agent.py#L686-L689。

Curriculum

如果启用 curriculum，每个外层 rollout 前都会调用同步；具体实现是把 rank 0 的 average-episode tracker 和 penalty scale 广播给其他 rank，并不是对各 rank curriculum 指标求平均：fast_sac_agent.py#L681-L684、locomotion_manager.py#L223-L236。

Checkpoint 与日志

只有 rank 0 定期保存 checkpoint/ONNX：fast_sac_agent.py#L794-L811。恢复时 setup 的初始化广播先发生，随后所有 rank 各自加载同一 checkpoint；load() 会恢复 actor、critic、target、normalizer、三个 optimizer、GradScaler 和 global step，但没有恢复/交换 replay buffer，也没有 load 后的再次广播：train_agent.py#L281-L299、fast_sac_agent.py#L626-L648。一致性依赖所有 rank 读取相同 checkpoint。

另一个可观测性边界是：episode 统计、TensorBoard/W&B 写入只取 rank 0，本地 loss/reward 没有额外做全局 reduce；吞吐量才用 num_gpus 乘回全局口径：logging_utils.py#L120-L149、logging_utils.py#L181-L203、logging_utils.py#L280-L342。所以 rank 0 日志不是全局数据分布的严格聚合指标。

6. 一句话定性

Holosoma FastSAC 是 replicated model + rank-local environment/replay shard + synchronous gradient averaging on every optimizer step。它不是共享 replay、不是周期参数平均、不是异步 parameter server，也不是 DDP wrapper；只是训练语义与同步 data parallel 相同，collective 由算法代码手工完成。

仓库的多卡 E2E 测试确实会用 torchrun 覆盖 FastSAC workflow，但主要验证训练/评估流程成功，不验证上述 collective 的数值等价性或通信性能：test_training.py#L77-L103、test_training.py#L126-L155。

7. 对 UniLab 的边界性启示

这套方案成立的前提是每个 rank 同时拥有一套 simulator、replay 和 learner，并在每个 optimizer step 等待所有 rank。UniLab 当前 off-policy production 拓扑是单 collector + learner-owned inference，且 #956 已明确把多 learner/NCCL、跨 learner replay 和多 GPU support claim 列为 non-goal，因此本文不建议把 Holosoma 路径直接移植进当前 execution path。

如果未来另立多 GPU roadmap，这次审计至少说明需要同时做出以下 contract 决策，不能只回答“用 DDP 还是同步参数”：

1. replay 是 rank-local shard、中心化，还是需要跨 rank 采样；
2. num_envs 和 global batch 的余数/最小 batch 如何定义；
3. critic、actor、temperature 是否每次 optimizer step 都同步；
4. normalizer、target network、optimizer/GradScaler、curriculum、checkpoint 的一致性由谁维护；
5. 日志是 rank 0 本地视角还是全局聚合视角。

相关讨论：

• 调研报告：Off-policy rollout 应该使用哪个 Actor？——Holosoma FastSAC 与 UniLab 的实践 #958 已解释 Holosoma 中 rollout actor 与 learner actor 的所有权关系；本文补充多 rank 之间到底同步什么、何时同步。
• roadmap(offpolicy): learner-owned GPU inference（单 collector，不做多 shard） #956 记录了 UniLab 当前 off-policy production 拓扑和多 GPU non-goal，本文不改变该决策。