Repository navigation
feat(flashsac): expose learner CUDA Graph controls - #2117
Merged
TATP-233 merged 1 commit intoOct 10, 2026
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
compile_full_objectivesto the UniLab FlashSAC configuration contract.use_whole_cycle_cuda_graphas an explicit learner CUDA Graph control.use_compile=true,compile_full_objectives=true, anduse_whole_cycle_cuda_graph=true.unilab-rl==1.4.11baseline and the current UniLab main branch.Why
These settings control different runtime layers:
use_compileenables or disables Torch Inductor compilation.compile_full_objectivesselects whether Inductor compiles the complete actor/critic objective or a smaller loss region.use_whole_cycle_cuda_graphcontrols 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_objectivesambiguous.Configuration examples
algo.algo_params.use_compile=falsealgo.algo_params.compile_full_objectives=falsealgo.algo_params.use_whole_cycle_cuda_graph=false# 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=falseLinked work
mainValidation
The branch was rebased onto the latest UniLab
main, and the previousuv.lockconflict was resolved.Passed locally:
20 passed.0 errors; one local environment warning for an optionaluni_rl.algos.fast_sacimport.Environment-limited checks (not code failures):
make test-allcould not complete because the local uv CUDA dependency download exceeded the temporary disk quota.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
use_whole_cycle_cuda_graphand documents the existing compile-scope controls.Checklist
unilab-rl==1.4.11.main.