[TRTLLM-14541][fix] VisualGen: byte-identical tune-run vs cache-load on multi-GPU#16782
[TRTLLM-14541][fix] VisualGen: byte-identical tune-run vs cache-load on multi-GPU#16782luyiyun1021 wants to merge 3 commits into
Conversation
…cache round-trips An enum tactic such as Fp4QuantTactic serializes via repr() to "<Fp4QuantTactic.TRTLLM: -1>", which ast.literal_eval cannot parse, so loading any autotuner cache that contains it crashes. Serialize the enum's underlying value instead; an IntEnum compares equal to it on reload. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe changes add enum-compatible autotuner cache serialization, cross-rank profiling-cache merging, scoped CUDA graph capture suppression, full-world visual generation mappings, and a two-phase distributed warmup flow. ChangesVisual generation autotuning
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant PipelineLoader
participant BasePipeline
participant CUDAGraphRunner
participant AutoTuner
participant DistributedWorld
PipelineLoader->>BasePipeline: disable CUDA graph capture
PipelineLoader->>AutoTuner: run autotuned warmup
AutoTuner->>CUDAGraphRunner: execute uncaptured keys eagerly
AutoTuner->>DistributedWorld: gather profiling caches
DistributedWorld-->>AutoTuner: return rank cache dictionaries
AutoTuner->>AutoTuner: apply fastest entries and save cache
PipelineLoader->>BasePipeline: run final warmup
BasePipeline->>CUDAGraphRunner: restore capture states
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/visual_gen/pipeline.py (1)
210-225: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the required return annotation.
no_cuda_graph_capture()is a new function but has no return annotation. AddIterator[None](or the repository’s equivalent) for the@contextlib.contextmanagergenerator.As per coding guidelines, every function must be annotated.
Proposed fix
+from collections.abc import Iterator + `@contextlib.contextmanager` -def no_cuda_graph_capture(self): +def no_cuda_graph_capture(self) -> Iterator[None]:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/pipeline.py` around lines 210 - 225, Update no_cuda_graph_capture with the repository’s standard generator return annotation, using Iterator[None] (or the equivalent typing symbol) while preserving its existing context-manager behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/visual_gen/pipeline_loader.py`:
- Around line 40-55: Update _VisualGenAutotuneDist.__init__ to type mapping as
Mapping, and annotate tp_cp_allgather’s obj parameter as object with a
list[object] return type. Add or reuse the appropriate Mapping import, keeping
the override free of Any.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 210-225: Update no_cuda_graph_capture with the repository’s
standard generator return annotation, using Iterator[None] (or the equivalent
typing symbol) while preserving its existing context-manager behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7d7e8fea-e77c-48b7-a459-77c158593170
📒 Files selected for processing (6)
tensorrt_llm/_torch/autotuner.pytensorrt_llm/_torch/visual_gen/cuda_graph_runner.pytensorrt_llm/_torch/visual_gen/mapping.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/_torch/visual_gen/pipeline_loader.pytests/unittest/_torch/misc/test_autotuner.py
1f00262 to
10270a2
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #61235 [ run ] triggered by Bot. Commit: |
10270a2 to
f526572
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tensorrt_llm/_torch/autotuner.py (1)
777-786: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNarrow the cache-serialization exception handler.
This changed serialization path catches every exception, although the expected round-trip failures are parse or assertion failures. Restrict it to the expected exception types so unrelated defects are not downgraded to a warning.
As per coding guidelines, “Avoid broad exception handling; catch specific exceptions instead of using bare
except:.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/autotuner.py` around lines 777 - 786, In the tactic round-trip validation within the autotuner cache-serialization path, narrow the handler around ast.literal_eval and the assertion to catch only the expected parsing and assertion exception types. Preserve the existing warning_once behavior for those failures while allowing unrelated exceptions to propagate.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/autotuner.py`:
- Around line 311-315: Update the autotune flow around setup_distributed_state
and cache I/O so rank is derived only after
autotuner.setup_distributed_state(*unify_tactics) installs the distributed
mapping. Use the resulting rank, or unify_tactics[0].rank, for subsequent cache
load/save operations while preserving cross-rank tactic unification.
In `@tests/unittest/_torch/misc/test_autotuner.py`:
- Around line 1270-1305: Extend the autotuner tests beyond direct cache merging:
add a profiling-cache save/load round-trip using an enum tactic and a focused
autotune context test that passes unify_tactics and verifies cross-rank
unification occurs before persistence. Anchor these additions to the existing
AutoTuner and profiling_cache test setup, and include the required changed-test
coverage and coverage verdict.
---
Nitpick comments:
In `@tensorrt_llm/_torch/autotuner.py`:
- Around line 777-786: In the tactic round-trip validation within the autotuner
cache-serialization path, narrow the handler around ast.literal_eval and the
assertion to catch only the expected parsing and assertion exception types.
Preserve the existing warning_once behavior for those failures while allowing
unrelated exceptions to propagate.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a558b706-3094-49a8-9579-13c86f592b84
📒 Files selected for processing (6)
tensorrt_llm/_torch/autotuner.pytensorrt_llm/_torch/visual_gen/cuda_graph_runner.pytensorrt_llm/_torch/visual_gen/mapping.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/_torch/visual_gen/pipeline_loader.pytests/unittest/_torch/misc/test_autotuner.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py
- tensorrt_llm/_torch/visual_gen/mapping.py
- tensorrt_llm/_torch/visual_gen/pipeline.py
|
/bot run --disable-fail-fast |
|
PR_Github #61241 [ run ] triggered by Bot. Commit: |
|
PR_Github #61235 [ run ] completed with state |
f526572 to
4c9d830
Compare
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/autotuner.py`:
- Around line 282-286: Update the autotuning context around post_tune_merge_dist
and the corresponding cleanup near the later exit block to save the singleton
autotuner’s prior mapping and _dist before setup_distributed_state, then restore
both after merge/cache persistence completes. Use a nested finally so
restoration occurs on success and failure, preventing later sessions from
retaining the temporary full-world distributed state.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 00511959-cbe7-4485-964a-5522bdab91f4
📒 Files selected for processing (6)
tensorrt_llm/_torch/autotuner.pytensorrt_llm/_torch/visual_gen/cuda_graph_runner.pytensorrt_llm/_torch/visual_gen/mapping.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/_torch/visual_gen/pipeline_loader.pytests/unittest/_torch/misc/test_autotuner.py
🚧 Files skipped from review as they are similar to previous changes (4)
- tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py
- tensorrt_llm/_torch/visual_gen/pipeline.py
- tensorrt_llm/_torch/visual_gen/mapping.py
- tensorrt_llm/_torch/visual_gen/pipeline_loader.py
|
PR_Github #61255 [ run ] triggered by Bot. Commit: |
|
PR_Github #61241 [ run ] completed with state |
Add AutoTuner.post_tune_merge_tactics(): one collective that all-gathers every rank's whole profiling cache and keeps the fastest tactic per key. Add an autotune(post_tune_merge_dist=...) option that runs it once at context exit, before the cache is saved: it attaches the dist's mapping up front (for a correct rank) but wires the transport only at exit, so _is_distributed() stays False and per-op sync is skipped while tuning, then restores the prior mapping and dist so the temporary full-world state does not leak. Unlike the per-op distributed_tuning_strategy sync it runs once and does not require the tuned ops to be invoked symmetrically across ranks, so it is safe for non-SPMD warmup. Reuses Distributed.tp_cp_allgather and the min-time winner rule. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
… are tuned and merged VisualGen warmup ran torch.compile, autotune, and CUDA-graph capture in one pass, so each rank baked its own in-flight, pre-merged tactic picks into the graphs. A fresh tuning run's output therefore differed from a cache load, and ranks silently diverged. Split warmup into two phases. Phase 1 tunes eagerly with capture disabled (CUDAGraphRunner.capture_enabled + BasePipeline.no_cuda_graph_capture()) and hands autotune()'s post_tune_merge_dist option a torch.distributed communicator (_VisualGenAutotuneDist) built on a world-sized mapping (VisualGenMapping.to_autotuner_mapping); the autotuner then merges tactics across all ranks in one collective at context exit. Phase 2 re-runs warmup to capture graphs from the finalized, merged tactics. The communicator gathers over the default world group because the autotuner's own tp_cp_allgather runs over the mapping's TP process group, and VisualGen's build-once device mesh has a single-rank TP axis (its parallelism is on cfg/ulysses), so that gather would return only the local cache. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
|
/bot run --disable-fail-fast |
4c9d830 to
835943a
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/misc/test_autotuner.py`:
- Line 1: Add the repository-standard NVIDIA copyright header at the top of
tests/unittest/_torch/misc/test_autotuner.py, using the 2026 year and placing it
before the existing import enum statement.
- Line 1330: Replace the direct profiling_cache.cache.clear() call in the
persistence test with the existing AutoTuner.clear_cache() reset method,
ensuring cache data and independent_op/excluded_op metadata are all cleared
before the test.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d70cf5ab-d911-46b5-a3fc-2afff1f13dfb
📒 Files selected for processing (6)
tensorrt_llm/_torch/autotuner.pytensorrt_llm/_torch/visual_gen/cuda_graph_runner.pytensorrt_llm/_torch/visual_gen/mapping.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/_torch/visual_gen/pipeline_loader.pytests/unittest/_torch/misc/test_autotuner.py
🚧 Files skipped from review as they are similar to previous changes (5)
- tensorrt_llm/_torch/visual_gen/pipeline.py
- tensorrt_llm/_torch/visual_gen/cuda_graph_runner.py
- tensorrt_llm/_torch/visual_gen/mapping.py
- tensorrt_llm/_torch/visual_gen/pipeline_loader.py
- tensorrt_llm/_torch/autotuner.py
| @@ -1,3 +1,4 @@ | |||
| import enum | |||
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add the required NVIDIA copyright header/year.
This modified file begins with an import and has no copyright header. Add the repository-standard NVIDIA header with the 2026 year. As per coding guidelines, “Add the NVIDIA copyright header to all new files and update the copyright year on modified files.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/misc/test_autotuner.py` at line 1, Add the
repository-standard NVIDIA copyright header at the top of
tests/unittest/_torch/misc/test_autotuner.py, using the 2026 year and placing it
before the existing import enum statement.
Source: Coding guidelines
| # autotune(post_tune_merge_dist=...) must merge across ranks and persist the | ||
| # merged winner, then restore the singleton's prior distributed state. | ||
| tuner = AutoTuner.get() | ||
| tuner.profiling_cache.cache.clear() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reset the complete singleton profiling-cache state.
Clearing only .cache leaves independent_op and excluded_op from earlier tests. Those sets influence cache partitioning/exclusion during save_cache, making this persistence test order-dependent.
Proposed fix
- tuner.profiling_cache.cache.clear()
+ tuner.clear_cache()Based on supplied autotuner context, AutoTuner.clear_cache() resets the cache and both associated metadata sets.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| tuner.profiling_cache.cache.clear() | |
| tuner.clear_cache() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/misc/test_autotuner.py` at line 1330, Replace the
direct profiling_cache.cache.clear() call in the persistence test with the
existing AutoTuner.clear_cache() reset method, ensuring cache data and
independent_op/excluded_op metadata are all cleared before the test.
|
PR_Github #61278 [ run ] triggered by Bot. Commit: |
|
PR_Github #61255 [ run ] completed with state |
| ), | ||
| ): | ||
| pipeline.warmup() | ||
| pipeline.warmup() |
There was a problem hiding this comment.
Is this second warmup also needed for model running on single GPU?
| tactic_str = repr(tactic.value) if isinstance( | ||
| tactic, enum.Enum) else repr(tactic) |
There was a problem hiding this comment.
This only works when the Enum is IntEnum. If tactic is any other types of enum, the following assertion will always fail.
|
PR_Github #61278 [ run ] completed with state
|
3 participants
Description
Multi-GPU VisualGen (LTX-2, Wan 2.2) DiT output was nondeterministic: a fresh autotuning run produced different frames than a persistent-cache load of the same run, and different ranks silently diverged. This PR fixes the three root causes so that a tune-run is byte-identical to a cache-load, and every rank runs the identical kernels.
1. CUDA graphs were captured while the autotuner was still tuning. Warmup ran torch.compile + autotune + CUDA-graph capture in one pass, so each rank baked its own in-flight tactic choices into the graphs before tactics were finalized. Rank-level
nsys --cuda-graph-trace=nodeconfirmed this: a rank whose local tactic for a shape was cutlass had no cuBLASLt (nvjet_*) kernel in its graph at that shape, whereas every cache-load rank uniformly baked the cache's tactic.2. NVFP4 GEMM tactic selection flips between numerically-distinct backends on tiny perf gaps. The GEMM autotuner picks among cutlass / cuBLASLt / cuda_core by measured time. On several LTX/Wan shapes cutlass and cuBLASLt are within timing noise of each other, so the winner flips per-launch and per-rank; the two backends are numerically distinct (~1e-4/element), so a flip changes the output. (Intra-backend flips are bit-exact, which is how the cross-backend flip was isolated: forcing a single backend is fully deterministic.)
3. The autotuner's cross-rank tactic sync never engaged for VisualGen (mapping/mesh gap). The autotuner unifies per-rank tactics through its
Distributedbackend keyed off aMapping. VisualGen'sto_llm_mapping()reportstp_size == 1(its parallelism is on cfg/ulysses), so the autotuner saw a world of 1 and never synced. Even forcing a world-sized mapping is not enough: the autotuner'stp_cp_allgathergathers over the mapping's TP process group, and VisualGen'sDeviceMeshTopologyImpl.device_meshis a build-once class singleton whose TP axis spans a single rank — so the gather returned only the local cache and the merge was a silent no-op (verified:gathered=1on 8 ranks).Fix — two-phase warmup + a one-shot world-domain post-tune tactic merge, generic to all VisualGen pipelines (mirrors the LLM
_run_autotuner_warmup, which also disables graph capture while tuning and syncs tactics through the autotuner's distributed backend):CUDAGraphRunner.capture_enabledflag +BasePipeline.no_cuda_graph_capture()context manager: warmup forwards run eagerly (no capture) while tuning.AutoTuner.post_tune_merge_tactics()+autotune(post_tune_merge_dist=dist): a new one-shot collective that all-gathers every rank's whole profiling cache and keeps the fastest tactic per key, run once at theautotune()context exit before the cache is saved. Unlike the per-opdistributed_tuning_strategysync it runs once and does not require the tuned ops to be invoked symmetrically across ranks, so it is safe for VisualGen's non-SPMD warmup (rank-0-only text encoder, parallel-VAE subsets).VisualGenMapping.to_autotuner_mapping()+_VisualGenAutotuneDist: give the autotuner a world-sized mapping so it treats all ranks as one tuning group, and a torch.distributed communicator whose gather spans the default world group (bypassing the VisualGen device mesh's single-rank TP axis).autotune()attaches that mapping up front (for a correct rank) but only wires the communicator at exit, so_is_distributed()stays False and per-op sync is skipped while phase-1 tuning runs; the merge then runs once at exit.PipelineLoader: phase 1 tunes eagerly with capture disabled and passespost_tune_merge_disttoautotune(), which merges tactics across all ranks and saves the cache at context exit → phase 2 re-runs warmup to capture graphs from the finalized, merged tactics.autotuner._serialize_cache_data: serialize enum tactics by.valuesoFp4QuantTacticround-trips (otherwiseast.literal_evalcrashes when loading any cache that contains it).Trade-off: warmup now runs the pipeline forward twice (tune, then capture), which increases load-time warmup. Scope: this makes a tune-run byte-identical to a persistent-cache load and merges tactics across ranks within a run. Full cross-launch determinism (two independent no-cache tuning runs) still requires the persistent autotuner cache to freeze the tactic table, since problem 2's near-tie selection is timing-dependent per launch.
Deterministic cache workflow
Determinism is gated by the
TLLM_AUTOTUNER_CACHE_PATHenv var.autotune()loads the cache when that file already exists (otherwise it tunes), and always re-saves the cross-rank-merged table on context exit — so a build-once-then-load flow reproduces byte-identical output on every run:With this PR the build run itself is already byte-identical to the loads (tune-run == cache-load), so the first, cache-building run is usable output rather than a throwaway. Cross-launch determinism across two independent no-cache runs still requires this persistent cache, since the tactic table depends on per-launch near-tie timing.
Test Coverage
Verified byte-identical tune-run vs cache-load output (mp4 md5) on B200 x8:
5a8bfb2f5a8bfb2f(equal)b06b734fb06b734f(equal)Pre-fix baseline on the same LTX config: tune-run vs cache-load mp4 md5 differed on every launch, and three fresh tuning runs produced three different outputs.
AutoTuner.post_tune_merge_tacticsgets unit coverage intests/unittest/_torch/misc/test_autotuner.py(min-time winner + rank-subset keys kept + single-rank no-op); the end-to-end determinism check is multi-GPU + real-checkpoint, so its scripts are attached to TRTLLM-14541 rather than added to CI.PR Checklist
PR description clearly explains what and why.
PR Follows TRT-LLM CODING GUIDELINES.
Test cases are provided for new code paths.
Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS updated if ownership changes.
Documentation updated as needed.
Update tava architecture diagram if significant design change.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Dev Engineer Review
visual_gen/pipeline_loader.py:pipeline.no_cuda_graph_capture()with tuning graph capture disabled (skip_dynamic_tuning_buckets=True) and an optional distributed post-merge hook.visual_gen/cuda_graph_runner.py(capture_enabled) and a pipeline-scoped context managerBasePipeline.no_cuda_graph_capture()that snapshots/restores runner state viafinallyto avoid leaking capture settings.autotune(..., post_tune_merge_dist=...)andAutoTuner.post_tune_merge_tactics()to all-gather per-rank in-memory profiling caches and merge by cache key while keeping the minimum recorded time.AutoTunerProfilingCache._serialize_cache_datanow serializesrepr(tactic.value)forenum.Enumtactics, matching the existing deserialization path and preventing enum literal round-trip failures.VisualGenMapping.to_autotuner_mapping()(world-sized autotuning group) and_VisualGenAutotuneDist(object all-gather viadist.all_gather_object) to unify autotuner state across ranks.QA Engineer Review
Test changes
tests/unittest/_torch/misc/test_autotuner.pytest_post_tune_merge_tactics_min_time_and_subset_kepttest_post_tune_merge_tactics_single_rank_nooptest_profiling_cache_enum_tactic_roundtriptest_autotune_post_tune_merge_before_saveCoverage in
tests/integration/test_lists/tests/integration/test_lists/(no matches found for the test names).Verdict
insufficient