perf(training): combine cuda_graphs with compile - #1479
Merged
Merged
Conversation
Borda
commented
Sep 17, 2026
Member
- Drop the ModelConfig validator that rejected cuda_graphs=True together with compile=True; the cuda_graphs docstring now describes the combined semantics
- RFDETRModelModule passes the triton.cudagraphs compile option when both flags are set, so Inductor cudagraph trees record and replay the compiled forward and backward; the eager CudaGraphTrainingRunner is never wrapped around the OptimizedModule
- New _inductor_cudagraphs_unsupported_reason gate falls back to compile-only with a warning for segmentation, keypoints, gradient checkpointing, grad_accum_steps > 1, and multi-device or multi-node configs; cudagraph trees allocate gradient outputs inside the graph pool, so accumulation across replays is unsupported
- _configure_cuda_graph_runner raises RuntimeError at train start when the trainer resolved accumulate_grad_batches > 1 or world_size != 1 on the Inductor path, since trainer_kwargs and devices="auto" can override the construction-time gate
- training_step and validation_step call torch.compiler.cudagraph_mark_step_begin() on the Inductor path so Lightning-held loss tensors and COCO-eval results do not trip the overwritten-output check
- Tests: config accepts both flags; compile options include triton.cudagraphs; fallback cases warn; runner stays off and train start stops outside scope; step marks fire only on the Inductor path
- Docs: advanced guide gains a "Combining CUDA graphs with compilation" section with the RTX PRO 6000 matrix (batch 4: 1.20x over compile; batch 64: parity) and scope limits; training-parameters, AGENTS.md, and the Unreleased CHANGELOG entry drop the mutual-exclusion wording; COCO2017 cookbook adds a CUDA_GRAPHS knob defaulting on for batch <= 16, the measured table, and a "What is left" note on the off-GPU share
- Drop the ModelConfig validator that rejected cuda_graphs=True together with compile=True; the cuda_graphs docstring now describes the combined semantics - RFDETRModelModule passes the triton.cudagraphs compile option when both flags are set, so Inductor cudagraph trees record and replay the compiled forward and backward; the eager CudaGraphTrainingRunner is never wrapped around the OptimizedModule - New _inductor_cudagraphs_unsupported_reason gate falls back to compile-only with a warning for segmentation, keypoints, gradient checkpointing, grad_accum_steps > 1, and multi-device or multi-node configs; cudagraph trees allocate gradient outputs inside the graph pool, so accumulation across replays is unsupported - _configure_cuda_graph_runner raises RuntimeError at train start when the trainer resolved accumulate_grad_batches > 1 or world_size != 1 on the Inductor path, since trainer_kwargs and devices="auto" can override the construction-time gate - training_step and validation_step call torch.compiler.cudagraph_mark_step_begin() on the Inductor path so Lightning-held loss tensors and COCO-eval results do not trip the overwritten-output check - Tests: config accepts both flags; compile options include triton.cudagraphs; fallback cases warn; runner stays off and train start stops outside scope; step marks fire only on the Inductor path - Docs: advanced guide gains a "Combining CUDA graphs with compilation" section with the RTX PRO 6000 matrix (batch 4: 1.20x over compile; batch 64: parity) and scope limits; training-parameters, AGENTS.md, and the Unreleased CHANGELOG entry drop the mutual-exclusion wording; COCO2017 cookbook adds a CUDA_GRAPHS knob defaulting on for batch <= 16, the measured table, and a "What is left" note on the off-GPU share --- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Borda
requested review from
SkalskiP,
isaacrob and
probicheaux
as code owners
September 17, 2026 14:40
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## develop #1479 +/- ##
=======================================
Coverage 88% 88%
=======================================
Files 128 128
Lines 16467 16509 +42
=======================================
+ Hits 14477 14517 +40
- Misses 1990 1992 +2 🚀 New features to boost your workflow:
|
- Column headers of the CUDA graphs × compile matrix in the advanced guide and the COCO2017 cookbook now carry units in brackets ([img], [img/s], [×], [-]); the "(images per second)" prose that used to carry the unit is dropped --- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Compile-only fallbacks can still wrap the compiled model with the eager runner, and resolved CPU placement is not guarded.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Enables combining torch.compile with Inductor CUDA graph trees for supported training configurations.
Changes:
- Adds Inductor cudagraph configuration, runtime guards, and step markers.
- Adds configuration and lifecycle tests.
- Updates training documentation, cookbook, changelog, and agent guidance.
File summaries
| File | Description |
|---|---|
src/rfdetr/training/module_model.py |
Implements combined compilation and graph replay. |
src/rfdetr/config.py |
Allows both acceleration flags. |
tests/training/test_module_model.py |
Tests configuration, fallbacks, and lifecycle behavior. |
tests/training/test_cuda_graph_step.py |
Tests combined configuration acceptance. |
docs/learn/train/training-parameters.md |
Documents combined semantics. |
docs/learn/train/advanced.md |
Adds benchmarks and usage guidance. |
docs/cookbooks/train-coco2017.ipynb |
Adds a batch-dependent CUDA graphs option. |
CHANGELOG.md |
Records the new capability. |
AGENTS.md |
Updates architecture guidance. |
Review details
- Files reviewed: 9/9 changed files
- Comments generated: 3
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Add check to disable Inductor CUDA graph replay for non-CUDA devices. Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Changes: - Track successful compilation independently from Inductor CUDA graph-tree activation so unsupported combined configurations cannot wrap an OptimizedModule in the eager CUDA graph runner. - Add lifecycle regressions for gradient accumulation and resolved CPU placement, plus a GPU-marked two-step production training-step replay test. Impact: - Unsupported compile+cuda_graphs requests now remain compile-only instead of risking double CUDA-graph capture. - Future changes are checked for CPU device resolution and real combined replay/gradient/optimizer behavior. Verification: - Focused CPU training suite: 238 passed, 2 skipped. - Pre-commit lint, formatting, docs, and whitespace checks passed. Residual limits: - The new GPU regression requires GPU CI; local CUDA execution was unavailable. - Mypy still reports the pre-existing unrelated src/rfdetr/models/weights.py:691 arg-type error. --- Co-authored-by: Codex <codex@openai.com>
Changes: - Move the `self.optimizers()` call and its list unwrap in `training_step` inside the `_use_manual_optimization` branch, the only consumer of the optimizer. - Add `TestTrainingStep.test_automatic_optimization_never_touches_trainer_strategy`, a CPU regression that runs a detection `training_step` with a `SimpleNamespace(accumulate_grad_batches=1)` trainer and no `optimizers` stub. Impact: - The automatic-optimization path no longer dereferences `trainer.strategy._lightning_optimizers`; this unblocks the GPU test `test_compiled_cudagraph_trees_replay_two_optimizer_steps`, which failed with `AttributeError: 'types.SimpleNamespace' object has no attribute 'strategy'`. - Keypoint manual optimization keeps the same optimizer handling. Verification: - New regression test fails before the change with the CI error and passes after. - `tests/training` CPU suite: 1375 passed, 23 skipped. --- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
7 tasks
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.