Skip to content

feat(flashsac): expose learner CUDA Graph controls - #2117

Merged
TATP-233 merged 1 commit into
Motphys:mainfrom
EtherealTide:feat/issue-94-flashsac-controls-draft
Oct 10, 2026
Merged

TATP-233 merged 1 commit into
Motphys:mainfrom
EtherealTide:feat/issue-94-flashsac-controls-draft

Conversation

@EtherealTide

@EtherealTide EtherealTide commented Oct 10, 2026 •

Copy link
Copy Markdown
Contributor

Summary

This PR exposes the three FlashSAC learner controls needed to distinguish Inductor compilation, objective compile scope, and learner-owned whole-cycle CUDA Graph capture/replay. The default accelerated path for a single NVIDIA CUDA learner remains unchanged, while users can now opt out explicitly for debugging and compatibility.

What changed

  • Added compile_full_objectives to the UniLab FlashSAC configuration contract.
  • Added use_whole_cycle_cuda_graph as an explicit learner CUDA Graph control.
  • Preserved the existing defaults: use_compile=true, compile_full_objectives=true, and use_whole_cycle_cuda_graph=true.
  • Added structured configuration coverage and updated the English and Chinese FlashSAC documentation.
  • Synced the implementation with the released unilab-rl==1.4.11 baseline and the current UniLab main branch.

Why

These settings control different runtime layers:

  1. use_compile enables or disables Torch Inductor compilation.
  2. compile_full_objectives selects whether Inductor compiles the complete actor/critic objective or a smaller loss region.
  3. use_whole_cycle_cuda_graph controls whether the learner captures and replays the full update cycle as a CUDA Graph.

Previously, UniLab did not expose the learner-level CUDA Graph choice. This made it difficult to disable the graph for debugging or compatibility while retaining Inductor, and made the meaning of compile_full_objectives ambiguous.

Configuration examples

Configuration Result
algo.algo_params.use_compile=false Eager update path; compile-dependent whole-cycle graph is disabled.
algo.algo_params.compile_full_objectives=false Inductor remains enabled, but compilation is limited to the smaller objective/loss region.
algo.algo_params.use_whole_cycle_cuda_graph=false Inductor remains enabled, but learner-owned whole-cycle CUDA Graph capture/replay is disabled.
All three options at their defaults Existing accelerated single-NVIDIA-CUDA path.
# Disable compilation and use the eager compatibility path
uv run train --algo flashsac --task g1_walk_flat --sim mujoco \
  algo.algo_params.use_compile=false
# Keep Inductor, but disable whole-cycle CUDA Graph capture/replay
uv run train --algo flashsac --task g1_walk_flat --sim mujoco \
  algo.algo_params.use_whole_cycle_cuda_graph=false

Linked work

Validation

The branch was rebased onto the latest UniLab main, and the previous uv.lock conflict was resolved.

Passed locally:

  • Focused configuration and documentation tests: 20 passed.
  • Ruff format: passed.
  • Ruff lint: passed.
  • Pyright: 0 errors; one local environment warning for an optional uni_rl.algos.fast_sac import.
  • Sphinx documentation build: succeeded previously with 6 pre-existing toctree warnings.

Environment-limited checks (not code failures):

  • make test-all could not complete because the local uv CUDA dependency download exceeded the temporary disk quota.
  • The local mypy process exited with code 139 before reporting diagnostics.
  • Full pytest collection hit an existing Numba cache-locator error because this worktree is mounted read-only for cache files.
  • An offline Sphinx re-run could not fetch sphinx-copybutton; the earlier successful build remains the available docs evidence.

Remote CI should be reviewed on the current PR head before moving this Draft PR to Ready for review.

Impact

  • Backend: no backend behavior changes.
  • Platform: primarily NVIDIA CUDA/Linux; CPU/eager configurations remain available.
  • Training: default accelerated behavior is preserved; explicit compile and graph opt-outs are now supported.
  • Public configuration: adds use_whole_cycle_cuda_graph and documents the existing compile-scope controls.

Checklist

  • Added or updated focused tests.
  • Updated English and Chinese FlashSAC documentation.
  • Linked the driving issue and released core dependency.
  • Confirmed compatibility with unilab-rl==1.4.11.
  • Rebased onto the latest UniLab main.
  • Recorded exact local validation results and environment limitations.
  • Confirm all current-head remote CI checks are green.
  • Run the full repository gate successfully in an environment with writable caches and sufficient disk space.
  • Move the PR from Draft to Ready for review after the checks above are complete.

@EtherealTide
EtherealTide marked this pull request as ready for review October 10, 2026 05:13
@TATP-233
TATP-233 merged commit 7b2abe0 into Motphys:main Oct 10, 2026
13 of 14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants