perf(flashsac): add opt-in DP CUDA graphs and BF16 gradient wire

Author: Cudrix-dev

AGENTS.md

@@ -29,6 +29,22 @@ architecture changes.
 
 ## Architecture contracts
 
+- FlashSAC DP whole-cycle CUDA graphs require explicit
+  `algo.algo_params.dp_whole_cycle_cuda_graph=true` and synchronizer lifecycle
+  hooks. Capture failure must raise on DP; rank-local eager fallback can hang
+  collectives. Release graphs before closing the process group. On the compiled
+  CUDA path without GradScaler, loss validity
+  travels with gradients so nonfinite losses with finite derivatives skip the
+  same optimizer and target update on every rank.
+- `training.dp_gradient_wire_dtype` defaults to `float32`. Optional `bfloat16`
+  compresses FP32 groups with at least 1024 elements, restores FP32 before
+  averaging, and changes numerical precision independently of AMP. Do not claim
+  convergence equivalence from short throughput runs. Distributed communication
+  remains in `rl/runtime/distributed`, never downstream task packages.
+- Multi-GPU MPS uses one named daemon per selected physical GPU. Select each
+  rank's pipe and log before CUDA initialization; inherited rank-zero paths must
+  be replaced in spawned ranks. Keep explicit GPU selection in launch docs.
+
 - UniLab provides environments and adapters, and now also owns the RL
   algorithm and training-execution stack in `unilab.rl` (merged from the
   unilab_rl repo @ 3654e4b). Its layers: `contracts` (env injection and

README.md

@@ -111,6 +111,15 @@ and [support matrix](https://unilabsim.github.io/UniLab-doc/en/5-reference/5-sup
 
 ## Ecosystem
 
+For multi-GPU FlashSAC, select the GPU UUID list with `CUDA_VISIBLE_DEVICES`;
+the native launcher starts one rank per visible GPU. Keep `algo.num_envs` and
+`algo.batch_size` fixed per rank. Opt into whole-cycle DP graphs with
+`+algo.algo_params.dp_whole_cycle_cuda_graph=true`; communication remains FP32
+unless `training.dp_gradient_wire_dtype=bfloat16` is explicitly selected.
+See [execution and precision contracts](docs/flashsac_dp.md) and the
+[FlashMimic launch instructions](https://github.com/Motphys/FlashMimic#train)
+for per-GPU MPS services and the live training panel.
+
 UniLab is designed to be a shared task and training surface for robot-specific
 repositories. They can ship robot recipes independently while consuming the
 same task, backend, and RL contracts. Current downstream examples:

README_zh.md

@@ -104,6 +104,14 @@ uv run --extra genesis eval --algo flashsac --task g1_motion_tracking --sim gene
 
 ## 生态
 
+多卡 FlashSAC 通过 `CUDA_VISIBLE_DEVICES` 指定 GPU UUID 列表，原生 launcher
+为每张可见卡启动一个 rank。`algo.num_envs` 和 `algo.batch_size` 均保持每 rank
+固定。使用 `+algo.algo_params.dp_whole_cycle_cuda_graph=true` 显式启用 DP 整周期
+CUDA graph；通信默认 FP32，可通过 `training.dp_gradient_wire_dtype=bfloat16`
+显式选择有损 BF16。详见[执行与精度约定](docs/flashsac_dp.md)和
+[FlashMimic 启动说明](https://github.com/Motphys/FlashMimic#train)，包含逐卡 MPS
+服务管理和训练面板。
+
 UniLab 被设计为机器人专属仓库共享的任务与训练界面。下游仓库可以独立发布机器人
 recipe，同时消费同一套 task、backend 和 RL contract。目前的下游示例：
 

docs/flashsac_dp.md

@@ -0,0 +1,53 @@
+# FlashSAC data-parallel CUDA graphs and wire precision
+
+## Scope and decision
+
+This change ports the measured FlashMimic multi-GPU optimization to UniLab
+1.5.0 (`03b75353`) with UniSim 2.1.0 (`6994cd1`). The intended PR base is `main`.
+The requested scope is opt-in DP graph execution and optional BF16 gradient
+communication; the existing runner owns process lifecycle and synchronization.
+No new runner or backend protocol is introduced. Defaults remain FP32 and DP
+eager execution. The actor/temperature fusion experiment is excluded because
+it did not provide a stable throughput improvement.
+
+The proposed contract is recorded in
+[ADR-0020](sphinx/source/adr/ADR-0020-flashsac-dp-graph-and-wire-precision.md).
+
+## Execution contract
+
+Set `+algo.algo_params.dp_whole_cycle_cuda_graph=true` to let FlashSAC capture
+the update cycle with NCCL reductions. The runner attaches synchronizer hooks
+around capture and before replay. Replacing callbacks, hooks, or checkpoint
+state invalidates the graph. A DP capture failure raises rather than allowing
+one rank to fall back to eager collectives. A loss-validity sentinel is packed
+with each optimizer's gradients, including losses whose derivatives are finite.
+On the compiled CUDA path without GradScaler, every rank therefore gates that
+optimizer and the critic target consistently. CPU, eager CUDA and FP16 scaler
+paths retain their existing finite-check behavior.
+
+This is raw CUDA graph capture around compiled objective/backward segments and
+optimizer updates. It is not a claim that the entire learner was compiled into
+one `torch.compile(fullgraph=True)` graph.
+
+`training.dp_gradient_wire_dtype=float32` preserves the default reduction.
+Explicit `bfloat16` packs FP32 groups of at least 1024 elements in BF16, reduces
+their SUM, restores gradients to FP32, then divides in FP32. Smaller groups and
+other source dtypes retain their original dtype. Parameters and optimizer states
+retain their existing dtype. Layout and precision must match across all ranks.
+This is a lossy numerical option, independent of AMP; long-run convergence
+equivalence has not been demonstrated.
+
+## Evidence and acceptance
+
+Historical experiments used UniLab `c71a7a20`, UniSim `26d17b4e`, FlashMimic
+`d3990125`, with env=2048, batch=2048 per rank and six updates per cycle. On the
+tested eight-GPU host, BF16 wire increased median throughput from 346875 to
+388227 steps/s (+11.9%). Relative gradient L2 error was approximately 0.3%.
+These are historical measurements, not benchmarks of the new upstream defaults.
+Learner time still grew with rank count; constant learner time was not achieved.
+
+Acceptance covers default precision, BF16 threshold and mean semantics, stable
+buffer storage, graph ownership and failure behavior, rank-wide nonfinite-loss
+gates, and Linux multi-GPU capture/replay. README commands show explicit GPU
+selection and user-owned MPS. CPU mock tests do not establish NCCL graph support.
+Full repository and current-head remote CI gates remain required before merge.

docs/sphinx/source/adr/ADR-0000-index.md

@@ -31,6 +31,7 @@ orphan: true
 | [ADR-0017 FlashSAC Learner-Side Custom Actor Contract](ADR-0017-flashsac-learner-side-custom-actor-contract.md) | RL algorithm extension / checkpoint semantics | Proposed |
 | [ADR-0018 Collector Tensor Device Sole Placement Knob](ADR-0018-collector-tensor-device-sole-placement-knob.md) | Training config / off-policy placement | Proposed |
 | [ADR-0019 CUDA Learner Forbids CPU Carrier](ADR-0019-cuda-learner-forbids-cpu-carrier.md) | Training config / carrier pairing contract | Proposed |
+| [ADR-0020 FlashSAC DP Graph Ownership And Wire Precision](ADR-0020-flashsac-dp-graph-and-wire-precision.md) | RL algorithm / distributed runtime | Proposed |
 | [ADR-0231 Integration-only K=1 MJWarp Collector CUDA Graph](ADR-0231-k1-mjwarp-collector-cuda-graph-integration-only.md) | Collector CUDA graph / historical pre-registration | Rejected / no-go at M2 |
 | [ADR-0238 State-Update-Only CUDA Graph Boundary](ADR-0238-state-update-only-cuda-graph-boundary.md) | Manager CUDA graph / historical pre-registration | Rejected / no-go |
 

docs/sphinx/source/adr/ADR-0020-flashsac-dp-graph-and-wire-precision.md

@@ -0,0 +1,73 @@
+---
+orphan: true
+---
+
+# ADR-0020 FlashSAC DP Graph Ownership And Wire Precision
+
+- Status: Proposed
+- Date: 2026-10-22
+- Owners: UniLab RL runtime maintainers
+- Supersedes: None
+- Superseded by: None
+
+## Context
+
+FlashSAC disables whole-cycle capture when gradient synchronization is attached,
+although the distributed runner already provides graph lifecycle hooks. Multi-GPU
+profiling also shows that large gradient reductions dominate learner time on
+hosts without GPU peer access. Downstream FlashMimic should consume these runtime
+capabilities rather than patching algorithm or communication implementations.
+
+## Decision
+
+Keep the existing runner and process-group owner. Add an explicit FlashSAC DP
+whole-cycle capture opt-in and a separate runtime wire-precision option. Preserve
+FP32 communication and existing DP eager execution as defaults. This decision
+does not change physics backend contracts or establish BF16 convergence parity.
+
+## Stable Contracts
+
+- `algo.algo_params.dp_whole_cycle_cuda_graph=true` requires synchronizer hooks;
+  changing hooks, gradient callbacks, or checkpoint state invalidates capture.
+- DP capture errors propagate. A rank must not independently fall back to eager
+  collectives. Graphs are released before process-group destruction.
+- Compiled CUDA execution without GradScaler includes loss validity in gradient
+  reduction, so all ranks skip the same invalid optimizer and critic target step.
+- `training.dp_gradient_wire_dtype` accepts `float32` (default) or `bfloat16`.
+  BF16 affects only FP32 groups with at least 1024 elements. Restore FP32 gradients
+  before division; preserve parameter/optimizer dtype and gradient storage.
+- Every rank uses the same layout and precision. The option is independent of AMP.
+- MPS pipes/logs are selected per rank before CUDA initialization. Explicit host
+  daemon directories do not require access to another user's default runtime root.
+
+## Alternatives Considered
+
+- Keep graph execution disabled: retains avoidable learner launch overhead.
+- Enable BF16 by default: rejected because gradient reduction is lossy and
+  long-run convergence equivalence has not been established.
+- Fuse actor and temperature reductions: excluded after no stable measured
+  throughput improvement; it would also change the small-group precision boundary.
+- Put optimized runtime code in FlashMimic: rejected because UniLab owns DP.
+
+## Consequences
+
+Optional CUDA/NCCL execution requires real multi-rank validation in addition to
+CPU regression tests. Historical throughput evidence must identify its exact
+source pins and configuration; it is not a promise for newer backend defaults.
+Capturing a raw CUDA cycle does not imply one fully compiled learner FX graph.
+
+## Evidence In Repo
+
+- `src/unilab/rl/algos/flash_sac/learner.py`
+- `src/unilab/rl/runtime/distributed/sync.py`
+- `src/unilab/rl/runtime/offpolicy/double_buffer_runner.py`
+- `src/unilab/scripts/train_offpolicy.py`
+- `tests/rl/algos/test_flash_sac_dp_graph.py`
+- `tests/rl/runtime/test_dp_wire_precision.py`
+- `tests/algos/test_offpolicy_dp_sync.py`
+- `tests/training/test_cuda_mps_cli.py`
+
+## Related Documents
+
+- {doc}`ADR-0016-rl-reintegration-into-unilab-rl`
+- {doc}`ADR-0014-cuda-mps-cli-lifecycle-owner`

src/unilab/conf/flashsac/config.yaml

@@ -92,6 +92,8 @@ training:
   # process sharing and fails closed if the host daemon is not already valid.
   # Manage an explicit user-owned daemon with `uni-cumps`.
   cuda_process_sharing: null
+  # Optional lossy gradient communication, independent of learner AMP.
+  dp_gradient_wire_dtype: float32
   nan_guard:
     enabled: true
     buffer_size: 100

src/unilab/conf/sac/config.yaml

@@ -70,6 +70,7 @@ training:
   # process sharing and fails closed if the host daemon is not already valid.
   # Manage an explicit user-owned daemon with `uni-cumps`.
   cuda_process_sharing: null
+  dp_gradient_wire_dtype: float32
   nan_guard:
     enabled: true
     buffer_size: 100

src/unilab/rl/algos/flash_sac/double_buffer.py

@@ -129,6 +129,7 @@ def build_double_buffer_runner(
         "use_compile": cfg.algo.algo_params.use_compile,
         "compile_full_objectives": _algo_param(cfg.algo, "compile_full_objectives", None),
         "use_whole_cycle_cuda_graph": _algo_param(cfg.algo, "use_whole_cycle_cuda_graph", None),
+        "dp_whole_cycle_cuda_graph": _algo_param(cfg.algo, "dp_whole_cycle_cuda_graph", False),
         "actor_normalize_parameters": bool(
             _algo_param(cfg.algo, "actor_normalize_parameters", True)
         ),