Issue662

[DH13768095744]

Reduce FastSAC learner launch fragmentation

背景

这个 PR 处理的是 FastSAC 在 sac/g1_walk_flat/mujoco 上暴露出的 learner-side GPU duty cycle 低、tiny kernel 过多、CUDA launch API bound 的问题，并用 sac/g1_walk_rough/mujoco 做相邻任务复核。

核心判断是：当前瓶颈主要不是稳定态 replay/H2D starvation，而是 SAC learner update 被切成大量微小 CUDA kernels；这些 kernels 在 RTX 5090 D 这类大 GPU 上无法填满 SM，并且 CUDA runtime launch API 的累计成本高于 GPU kernel 设备端执行时间。

本 PR 的优化范围保持在 FastSAC learner、off-policy runtime、replay staging owner layer 内，不把长期业务规则放到 scripts/，也不改变 env/backend contract、runner lifecycle 或 SAC 的数学语义。

Summary

• Reduce FastSAC target update、alpha update、symmetry/replay staging、critic/actor graph boundary 等路径的 launch fragmentation。
• Add opt-in CUDA graph replay for critic / actor update，并配套 capture-safe packed staging。
• Use a single H2D staging path for SAC graph inputs, then perform graph-friendly D2D staging on GPU side.
• Preserve replay sampling semantics, SAC update order, optimizer semantics, target network semantics, and reward/task configuration.
• Provide profile evidence across E2E time, reward smoke, Nsight launch API time, tiny kernel count, GPU metrics, copy ranges, and issue-node ablation.

实验环境

Figure 1. 实验环境。 Profile 在 UniLab 本地 unilab 环境中执行，硬件为 NVIDIA GeForce RTX 5090 D。后续 GPU utilization、SM active、waves/SM 与 tiny-kernel 结论都应按该硬件规模解释；这些数字用于说明 launch fragmentation 的实现形态，不应被读成跨硬件的绝对性能承诺。

Root Cause Evidence

Timeline trace

Figure 2. Timeline trace 约束根因。 Trace 证据显示 learner 稳态通常没有长时间等待 replay batch 或 H2D completion；主要耗时集中在周期性 learner/update_critic、learner/update_actor 和相关 update window。因此后续优化优先聚焦 learner launch fragmentation，而不是先重写 replay/H2D path。

Nsight Systems

Figure 3. Nsight Systems 系统级证据。 Profile 中 CUDA kernels 数量很大且平均时长很短；CUDA launch API 的累计时间高于 GPU kernel 设备端累计时间。这说明 bottleneck 是频繁 launch 小 work，而不是少数大 kernel 计算慢。

Figure 4. Launch API vs kernel vs H2D。 Launch API time 明显大于 CUDA kernel total time 和 H2D memcpy total time，支持优先减少 kernel launch 数量和 launch API 时间。这个图不表示 H2D 永远不是瓶颈，只说明本 issue 的已测 case 不支持 H2D 为主因。

Figure 5. Launch API time 消融柱状图。 该图从 CUDA runtime API 侧展示不同优化节点对 launch API time 的影响。局部 tensor/update 清理收益有限，CUDA graph 和 graph-boundary reduction 才带来数量级下降。

Nsight Compute

Figure 6. NCU kernel-level 摘要。 Sampled kernels 的 duration、SM throughput、achieved occupancy 和 waves/SM 都显示大量 kernel 本身 underfilled。也就是说，低 GPU duty cycle 不只是外部等待或采样窗口造成，kernel launch 规模本身就不足以喂满大 GPU。

Figure 7. NCU sampled kernel efficiency。 图中同时展示 duration、SM throughput、occupancy 和 waves/SM。大量 kernel 位于 <10us 和低 waves/SM 区域，说明优化方向应是减少碎片化 launch，而不是只调 profiler 采样窗口或依赖 nvidia-smi 解释。

torch.compile coverage

Figure 8. torch.compile on/off matrix。 torch.compile 可以显著减少 eager pointwise/reduction fragmentation，但无法覆盖 optimizer、target update、symmetry、graph boundary、collector/replay/sync 等路径。因此 compile 是有效缓解项，但不是完整解决方案。

Figure 9. Compile on/off comparison。 Kernel launches、kernel total time 和 launch API total time 随 compile-on 下降，而 H2D copy total time 基本不变。这进一步说明 compile 的收益来自减少 eager fragmentation，而不是改变数据搬运。

Replay / H2D costs

Figure 10. Replay/H2D benchmark 表。 Replay benchmark 用 SAC G1 Walk MuJoCo 的 replay shape 单独测量采样和 H2D placement 成本。当前 incremental H2D path 的成本远小于 learner update window 级别开销，CPU pre-sample + sampled batch H2D 的替代路径反而更慢。

Figure 11. Replay costs 柱状图。 该图将 CPU random sample、device random sample、incremental H2D 和 sampled batch H2D 放在同一纵轴比较。它支持保留当前 replay incremental H2D path，并把优化重点放在 learner launch structure。

NVTX attribution

Figure 12. NVTX range tiny-kernel 分桶表。 Tiny kernels 集中在 learner owner ranges，例如 target soft update、symmetry augment、alpha update、actor/critic backward 和 graph boundary，而不是均匀散落在 env/backend 或 replay wait path。

Figure 13. Tiny kernels by NVTX range。 全局 <10us kernels 占比较高，并且集中在少数 learner 子阶段。该图把“GPU 利用率低”具体化为可处理的 owner-layer targets：减少逐参数/逐 tensor update、小 op、copy/cat/index fragmentation 和未被 compile 覆盖的小 launch。

Figure 14. Root-cause ranking。 多层证据把 learner tiny-kernel / launch-bound fragmentation、work-size underfeed 和 compile coverage 不完整排在主因前列；replay/H2D starvation 被降级为低置信主因。

Main vs Issue662 AB Profile

AB test 比较 origin/main 与 issue662，覆盖两个任务：

• sac/g1_walk_flat/mujoco
• sac/g1_walk_rough/mujoco

每个 task/variant 运行 3 次，E2E 时间取均值，error bar 表示 min/max。关键 profile 指标包括 tiny kernel 数量、launch API time/count、MemcpyAsync、stream sync、GPU metrics、graph copy range 和 time-reward smoke。

E2E wall time

Figure 15. Main vs issue662 E2E training wall time。 两个 MuJoCo 任务中，issue662 的 200-iteration wall time 相比 main 均显著下降。该图是端到端收益入口图，但收益来源需要和 tiny kernel、launch API 与 issue-node ablation 一起解释。

Figure 16. Issue-node E2E 消融。 早期局部优化节点对 E2E 改善有限，真正的结构性拐点出现在 CUDA graph 化及后续 graph-boundary launch reduction。

Reward smoke

Figure 17. g1_walk_flat_mujoco time-reward smoke。 横轴是训练 wall time，纵轴是 reward。issue662 更快完成同样 iteration 数；短 run 中没有观察到明显 reward 崩坏。该图只作为 smoke guardrail，不证明长期训练质量完全等价。

**Figure 17b. g1_walk_flat_mujoco 默认 iteration 的 main vs issue662 CUDA Graph time-reward 复测。3-run 均值显示 wall time 从 main 的约 468s 降至 issue662 CUDA Graph 的约 270s，说明端到端训练时间收益来自实际启用 graph path，而不是默认配置噪声。

Figure 18. g1_walk_rough_mujoco time-reward smoke。 Rough task 重复相同检查。曲线用于确认性能优化没有在短 run 中引入明显训练异常，但不能替代更长训练、更多 seed 或最终策略质量评估。

Infra metrics

Figure 19. g1_walk_flat_mujoco infra metrics。 Issue662 显著降低 <10us kernels、Launch API time 和 MemcpyAsync count；同时 stream sync time 上升，反映 CUDA graph replay / graph boundary 之后的同步形态变化。总体上主要 launch fragmentation 被压缩，但 residual sync/copy 成本仍是后续观察点。

Figure 20. g1_walk_rough_mujoco infra metrics。 Rough task 复现 flat task 中的 tiny-kernel 和 launch API 下降，说明优化触及共享 learner/runtime launch structure，而不是单个 terrain/task 的偶然现象。

Figure 21. Tiny kernel count 消融。 <10us kernels 在早期局部节点中基本不变，在 CUDA graph 和 graph-boundary reduction 后出现数量级下降。这是本 PR 最直接对应 issue 根因的指标之一。

GPU metrics

Figure 22. g1_walk_flat_mujoco GPU metrics。 Nsight GPU metrics 显示 issue662 中 SMs Active、GR Active、SM Issue 和 Async Copy Engine 的 activity 均提高。该指标用于辅助说明 GPU work 更集中，但不能直接等同于 E2E speedup。

Figure 23. g1_walk_rough_mujoco GPU metrics。 Rough task 中 GPU activity 变化方向与 flat task 一致。优化改善了 launch-bound workload 的提交结构，但并不声称 GPU 已达到大型 supervised DL workload 那样的持续饱和状态。

Graph boundary copy ranges

Figure 24. g1_walk_flat_mujoco graph boundary copy ranges。 Main 没有 graph copy range；issue662 引入 critic/actor graph copy ranges。这是 CUDA graph / packed staging 的边界成本：它留下更可控的 copy kernels，同时整体 tiny kernel 和 launch API count 大幅下降。

Figure 25. g1_walk_rough_mujoco graph boundary copy ranges。 Rough task 的 graph boundary copy 形态与 flat task 接近，说明 residual copy kernels 与共享 SAC graph layout 相关，而不是某个任务配置单独引起。

Cycle-level learner update

Figure 26. Per-cycle GPU learner update latency。 Cycle 定义为相邻 CPU/env simulation 边界之间的训练窗口；图中统计每个 cycle 内 learner GPU update 的 mean、p90 和 p99。CUDA graph 化后，周期性 learner update 时间显著下降，和 tiny kernel / launch API 下降方向一致。

FlashSAC Extension

同一类 CUDA graph 优化已经迁移到 FlashSAC 路径，并保持 opt-in 语义。默认 conf/offpolicy/algo/flashsac.yaml 仍关闭 use_cuda_graph_critic、use_cuda_graph_actor、use_cuda_graph_critic_packed_staging 和 use_cuda_graph_actor_packed_staging；只有显式 override 时，FlashSAC learner 才进入 graph replay 路径。实现上，FlashSAC learner 增加 critic / actor graph capture 与 replay body，runner 将 Hydra 参数传递到 learner，double-buffer runtime 只在 learner 显式声明支持 packed staging 且对应 flag 打开时生成 sac_graph_packed_source。因此，优化边界仍位于 learner / replay staging owner layer，没有把 FlashSAC 的私有能力泄漏到 env 或 backend contract。

该迁移复用了 FastSAC/SAC 中已经验证过的设计原则：graph replay 使用固定地址 staging buffer；shape 变化会重建 graph；packed staging 只改变 graph 输入组织方式，不改变 replay sampling、critic update、actor update、temperature update 或 target soft update 的数学含义。由此，FlashSAC 的修改目标被限定为降低 learner update 的 launch fragmentation，而不是改变策略优化问题本身。

FlashSAC timeline evidence

FlashSAC 的原始 clean timeline 显示，cycle 中的长条主要来自 collector pacing、actor inference、env step 以及 learner update 的交叠压力。Figure 27 给出未启用 FlashSAC graph path 时的稳态 cycle 分解；这说明 FlashSAC 的瓶颈并不只是单个 GPU kernel 慢，而是 collector/learner pipeline 中多个短阶段的重复调度成本。

> **Figure 27. FlashSAC baseline pipeline timeline。** `flashsac/g1_walk_flat/mujoco` clean timeline，`iter=100`，图中 block 宽度按稳态 cycle mean duration 绘制。该图用于定位 baseline 中 actor inference、env step、replay transfer 与 learner update 在一个 cycle 内的相对位置。

在 Motrix 后端上启用 both_graph_packed 后，Figure 28 的 smoke timeline 显示 learner update 被压缩到更窄的 cycle 区间。该图的意义不在于单独证明最终吞吐，而在于显示 CUDA graph + packed staging 对 cycle 内 learner 窗口的结构性影响：critic/actor update 不再以大量碎片化 launch 的形式铺满一个宽窗口，而是更集中地完成。

Figure 28. FlashSAC optimized Motrix smoke timeline。 flashsac/g1_walk_flat/motrix，both_graph_packed，iter=2 smoke run。该图用于快速确认 optimized graph path 在 Motrix 后端可运行，并观察 learner update 在 cycle 内的压缩方向。

Figure 29 将同一 optimized path 延长到 iter=50 后重新统计，避免只依赖极短 smoke run 的瞬时形态。iter50 timeline 中，cycle 仍保留 actor inference 与 env step 的可见占比，说明后续瓶颈会更多转向 collector/env cadence；但 learner update 的 cycle 占位已经显著收敛，这与 CUDA graph 主要压缩 launch-bound learner work 的预期一致。

Figure 29. FlashSAC optimized Motrix iter50 timeline。 flashsac/g1_walk_flat/motrix，both_graph_packed，iter=50，由 analyze_offpolicy_trace.py 基于 clean timeline 统计绘制。该图说明 graph path 的 cycle 级效果能够在更长的 Motrix run 中保持，而不是只出现在 2-iteration smoke trace。

FlashSAC time-reward evidence

从宏观行为看，FlashSAC CUDA graph path 的核心收益是 wall-clock 压缩，而不是 reward 函数或样本效率的改变。MuJoCo 与 Motrix 的 iter=500 time-reward 图均以 3 个 seed 串行运行，TensorBoard reward/mean 按训练 wall time 对齐，均值曲线配合 min/max band 展示。两组后端上，baseline 与 optimized 曲线处于同一行为量级；局部早期区间中 baseline/native 曲线可能先达到较高 reward，因此这些图不支持把 CUDA graph 解读为提升 reward 收敛速度的算法改动。它们支持的结论更窄：在未观察到明显行为漂移的前提下，optimized path 更快完成相同 iteration 预算。

Figure 30. FlashSAC MuJoCo time-reward。 flashsac/g1_walk_flat/mujoco，iter=500，baseline 与 CUDA Graph optimized 各 3 runs；baseline 关闭 critic/actor CUDA graph 与 packed staging，optimized 显式打开四个 graph/staging flag。3-run wall time 均值从 baseline 的约 127.1s 降至 optimized 的约 56.1s；reward 曲线用于行为 smoke，不作为样本效率提升声明。

Figure 31. FlashSAC Motrix time-reward。 flashsac/g1_walk_flat/motrix，iter=500，baseline 与 CUDA Graph optimized 各 3 runs；optimized 使用 critic/actor CUDA graph 与 packed staging。3-run wall time 均值从 baseline 的约 118.1s 降至 optimized 的约 66.2s。曲线显示优化后的训练行为仍与 baseline 处于同一 reward 区间，同时主要收益体现为完成同等训练预算所需 wall time 减少。

Semantic Equivalence / Algorithm Safety

这个 PR 的设计目标是 infra-level launch reduction，而不是改变 SAC 算法。

保持不变的语义包括：

• CPU env interaction、replay buffer 写入、随机采样的时机与样本含义不变。
• H2D 仍然发生在 replay sample 之后；issue662 只改变 graph-friendly staging 组织方式。
• Critic update、actor update、alpha update、target soft update 的数学含义不变。
• CUDA graph replay 使用固定地址 staging buffer，但不改变 tensor value 的来源或 batch 语义。
• training.sim_backend、task/reward/backend 选择仍走 Hydra + owner config，不在脚本层绕开 contract。

Reward smoke 图只支持短 run 中未观察到明显异常；长期策略质量仍应通过更长训练和更多 seed 验证。

[TATP-233]

再适配一下 flashsac 吧，任务指令：

uv run train --algo flashsac --task g1_walk_flat --sim mujoco

[TATP-233]

补充报告：
uv run train --algo sac --task g1_walk_flat --sim mujoco
的端到端训练时间对比

[DH13768095744]

补充报告： UV run train --algo sac --task g1_walk_flat --sim mujoco 的端到端训练时间对比

再适配一下 flashsac 吧，任务指令：

uv run train --algo flashsac --task g1_walk_flat --sim mujoco

都弄了