perf(closed_loop): torch.compile the evaluation forward pass - #342
Closed
HansRobo wants to merge 1 commit into
Closed
perf(closed_loop): torch.compile the evaluation forward pass#342HansRobo wants to merge 1 commit into
HansRobo wants to merge 1 commit into
Conversation
HansRobo
changed the title
HansRobo
marked this pull request as draft
July 31, 2026 00:08
HansRobo
force-pushed
the
feat/closed-loop-compile-bitexact
branch
from
July 31, 2026 00:28
1067918 to
413c679
Compare
HansRobo
changed the title
HansRobo
force-pushed
the
feat/closed-loop-compile-bitexact
branch
from
July 31, 2026 05:35
413c679 to
c54cd24
Compare
HansRobo
changed the title
HansRobo
force-pushed
the
feat/closed-loop-compile-bitexact
branch
from
July 31, 2026 05:38
c54cd24 to
7825f62
Compare
|
Is this PR compatible with the following PR? |
Member
Author
|
@SakodaShintaro |
Compile the encoder and the DiT with the cudagraphs backend for the duration of a closed-loop evaluation. Full-scale evaluation runs about 2x faster. A compiled model cannot take nn.MultiheadAttention's fused fastpath, so it rounds differently in the last float32 bit. A closed loop feeds that back into its own next input, so closed-loop metrics shift once.
HansRobo
force-pushed
the
feat/closed-loop-compile-bitexact
branch
from
July 31, 2026 07:17
7825f62 to
510317b
Compare
HansRobo
marked this pull request as ready for review
July 31, 2026 07:21
yhisaki
approved these changes
Aug 7, 2026
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
4 participants
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.
Compiles the encoder and the DiT for the duration of a closed-loop evaluation. Full-scale evaluation goes from 2,869 s to 1,443 s — 1.99x.
Measured over four sites and both object modes, 64,227 steps, identical in both arms, with eight evaluations running concurrently on one node with 28 cores and one GPU each — the configuration closed-loop validation runs in. Spread over four repeats was 0.6% before and 0.1% after. The baseline ran with the MHA fastpath already disabled; that switch on its own moved wall time within ±0.3%.
Closed-loop metrics shift by the last float32 bit — a compiled model cannot take
nn.MultiheadAttention's fused fastpath, and the loop amplifies the difference into its own next input.What changed
New
scenario_generation/inference_compile.py.compiled_for_inference()compiles and restores the model on exit — the training loop hands over its live model, and a leftover wrapper would put_orig_mod.prefixes in the next checkpoint.mark_inference_step()opens one cudagraph step per inference.For reviewers
The fastpath switch is entered inside
compiled_for_inference()rather than left to the caller: it is not in dynamo's guard set, so setting it after compilation leaves the traced graph in place while the flag reads back as changed.inductorwas tried and rejected — a second source of numeric difference on top, and no faster.