perf: SAC replay host sample + H2D 单卡 learner 热路径 — Tier1 聚焦 motion_tracking

[TATP-233]

Work type

benchmark / perf

Area

G1 / Motrix / SAC / replay placement / benchmark

Context

Related: #665。#665 覆盖 Env update_state / reset_done 侧 overhead；本 issue 跟进 SAC learner 侧 replay boundary：host replay 上的 random gather、pinned staging、H2D transfer，以及 GPU-resident replay / GPU learner cache 的收益空间。

Benchmark script 已在 #693 提交：benchmark/benchmark_sac_replay_buffer_sampling.py。

推进分两阶段：

Tier	Scope	主任务	目标
Tier1	单卡 SAC	uv run train --algo sac --task g1_motion_tracking --sim motrix	建立 motion_tracking 单卡 replay sample/H2D 基线，完成单卡 replay placement 原型验证
Tier2	2/4/8GPU SAC	motion_tracking first，其他 SAC 任务作为覆盖	评估 per-rank/global replay traffic、rank-local GPU cache、sharded/bulk refresh 方案

Tier1 主案例：g1_motion_tracking + motrix

最终 Hydra compose 关键值：

项	值
command	uv run train --algo sac --task g1_motion_tracking --sim motrix
training.task_name	G1MotionTrackingSAC
training.env_steps_per_sync	1
algo.num_envs	2,048
algo.batch_size	8,192
algo.replay_buffer_n	512
algo.updates_per_step	4
algo.learning_starts	1
algo.use_symmetry	false
obs_dim	160
action_dim	29
critic_dim	289
packed row width	930 float32
row bytes	3,720 B
replay capacity rows	1,048,576
replay _storage size	3,900,702,720 B = 3.63 GiB

Collector 写入量和 learner 采样量：

阶段	计算	transition rows	bytes
单次 env.step / ReplayBuffer.add	num_envs	2,048	7,618,560 B = 7.27 MiB
单个 collector sync chunk	num_envs * env_steps_per_sync	2,048	7.27 MiB
learner 启动阈值	max(batch_size, learning_starts * num_envs)	8,192	30,474,240 B = 29.06 MiB
启动前 sync chunks	ceil(8,192 / 2,048)	4 chunks	29.06 MiB total
稳态每个 learner iter collector 新增	num_envs * env_steps_per_sync	2,048	7.27 MiB
稳态每个 learner iter 每 rank 采样/H2D	batch_size * updates_per_step	32,768	121,896,960 B = 116.25 MiB
单卡 learner sample:new ratio	32,768 / 2,048	16x	16x

Tier1 单卡下，collector 每个 sync chunk 写入 7.27 MiB，learner 每 iter 从 replay 读取并传输 116.25 MiB。这个 16x replay read/write 比例来自 SAC replay reuse，是 replay placement 优化的核心 workload。

已有 benchmark 参考

当前 #693 的已完成 microbenchmark 可作为 replay placement harness 的趋势参考。该已有 workload 单卡每 iter H2D 为 107.5 MiB，和 motion_tracking 的 116.25 MiB 处于同一量级；Tier1 将以 motion_tracking 重跑并作为主结果。

单卡、capacity range 结果，5 次 warmup 后 20 次 measured repeat 取均值：

Capacity rows	CPU sample	H2D	CPU sample+H2D end-to-end	GPU buffer sample	Speedup
1,048,576	6.098 ms	2.070 ms	5.999 ms	0.119 ms	50.5x
4,194,304	6.298 ms	2.071 ms	6.233 ms	0.119 ms	52.6x
16,777,216	5.998 ms	2.071 ms	6.313 ms	0.120 ms	52.8x
67,108,864	5.997 ms	2.071 ms	5.990 ms	0.116 ms	51.6x

同一 workload 的多卡参考数据放入 Tier2：

GPUs	per-rank H2D	global H2D	CPU sample+H2D end-to-end	GPU buffer sample	Speedup
1	107.5 MiB	107.5 MiB	5.999 ms	0.119 ms	50.5x
2	107.5 MiB	215.0 MiB	6.145 ms	0.256 ms	24.0x
4	107.5 MiB	430.0 MiB	10.203 ms	0.432 ms	23.6x
8	107.5 MiB	860.0 MiB	20.034 ms	0.837 ms	23.9x

Tier2 解释模型：每个 rank 保持相同 per-rank batch，world size 增大时 global H2D 和 host random gather 并发压力按 rank 数增长。8GPU 的 860 MiB global H2D 暴露 shared host memory bandwidth、PCIe/xGMI path、driver submit、pinned staging、rank synchronization 的组合成本。

Host Replay Boundary 开销来源

• CPU random gather 对 packed replay storage 做大 batch 非连续访问，cache pressure 明显。
• Pinned host staging 后仍有 116.25 MiB/rank/iter 的 H2D boundary。
• SAC learner 每 iter 消费多批 replay rows；motion_tracking Tier1 单卡读写比为 16x。
• GPU-resident replay / GPU learner cache 将 learner sample 保持在 GPU memory domain，microbenchmark 中同量级 workload 已看到 50x 级单卡 sample boundary speedup。

Deliverable

• SAC replay placement Tier1 单卡方案：host authoritative replay + GPU learner cache，或 GPU-resident learner replay。
• Benchmark 支持 --task g1_motion_tracking，并输出 motion_tracking 单卡 CPU sample、H2D、CPU sample+H2D、GPU sample 柱状图。
• 单卡端到端训练报告：throughput、learner update time、replay sample/H2D exposed time、collector wait time。
• Tier2 计划文档：2/4/8GPU 的 per-rank/global traffic 指标、rank-local cache/sharding 方向、显存边界。

推荐路径

Tier1 / P0：motion_tracking 单卡 replay placement

• 将 benchmark 的 config compose 从固定 owner config 扩展为 task 参数，默认主案例设为 g1_motion_tracking。
• 基于 motion_tracking shape 复现 capacity range：2 ** i for i in range(20, 31, 2)。
• 实现单卡 GPU learner cache / GPU-resident replay prototype，并记录 exposed replay boundary time。
• 保持 collector 写入、runner lifecycle、ReplayBuffer contract 清晰：collector 继续写 authoritative transition rows，learner 热路径使用 GPU-side batch source。

Tier2 / P1：多卡 replay traffic scaling

• 在 Tier1 指标稳定后扩展到 2/4/8GPU。
• 记录 per-rank H2D、global H2D、rank barrier、collector wait、learner sample exposed time。
• 评估 rank-local GPU cache、sharded replay view、periodic bulk refresh 三类方案。
• 报告 global sample:new ratio：2GPU 32x、4GPU 64x、8GPU 128x（motion_tracking 默认配置）。

P2：batch bytes 和 field layout 优化

• 梳理 SAC learner 实际使用字段，压缩热路径 batch payload。
• 评估 packed layout、field slicing、dtype policy、prefetch overlap 的组合收益。
• batch reuse / shared indices 以 SAC 更新统计一致性验证为准。

Definition of done

• Tier1 单卡：uv run train --algo sac --task g1_motion_tracking --sim motrix 对应的 replay workload 有完整 microbenchmark 和端到端训练指标。
• 单卡 benchmark 图表包含 CPU sample、H2D、CPU sample+H2D、GPU sample，标题、标记、数值标签布局清晰。
• 单卡 replay placement prototype 给出 host path 与 GPU-side path 的 exposed replay boundary 对比。
• Tier2 进入条件：Tier1 指标、实现边界、显存预算文档齐备。
• make test-all 通过后再提交实现 PR。

Validation plan

Tier1 benchmark 目标命令（benchmark 增加 --task 参数后）：

uv run benchmark/benchmark_sac_replay_buffer_sampling.py \
  --task g1_motion_tracking \
  --sim motrix \
  --gpu-counts 1 \
  --capacity-exponents 20,22,24,26,28,30 \
  --warmup 5 \
  --repeat 20 \
  --out-json benchmark/outputs/sac_replay_buffer_sampling/results_motion_tracking_motrix_1gpu.json

Tier1 plot-only：

uv run benchmark/benchmark_sac_replay_buffer_sampling.py \
  --plot-only benchmark/outputs/sac_replay_buffer_sampling/results_motion_tracking_motrix_1gpu.json

Tier1 end-to-end：

uv run train --algo sac --task g1_motion_tracking --sim motrix training.num_gpus=1

Tier2 benchmark 目标命令：

uv run benchmark/benchmark_sac_replay_buffer_sampling.py \
  --task g1_motion_tracking \
  --sim motrix \
  --gpu-counts 1,2,4,8 \
  --capacity-exponents 20,22,24,26 \
  --warmup 5 \
  --repeat 20 \
  --out-json benchmark/outputs/sac_replay_buffer_sampling/results_motion_tracking_motrix_multigpu.json

Dependencies and blockers

• Related: perf: 定位 Env overhead 并行化实测天花板 — Numba prange 融合 kernel 在 Xeon 8568Y+ 到 31x #665
• Benchmark PR: benchmark: add SAC replay sampling placement test #693
• Tier1 需要 benchmark task selector 和 motion_tracking replay shape validation。