[TRTLLM-11412][feat] Add offloading support for visual_gen - #14095
Open
rahul-steiger-nv wants to merge 22 commits into
Open
[TRTLLM-11412][feat] Add offloading support for visual_gen#14095rahul-steiger-nv wants to merge 22 commits into
rahul-steiger-nv wants to merge 22 commits into
Conversation
Contributor
📝 WalkthroughWalkthroughThis PR implements CPU offloading for visual-generation pipelines, enabling selective movement of text encoders, VAEs, and transformer blocks between GPU and pinned CPU storage to reduce peak GPU memory consumption. The change spans infrastructure (packing/staging manager), pipeline integration, model-specific wiring for WAN T2V, metadata materialization, memory logging, CLI/config updates, and comprehensive tests. ChangesCPU Offloading for Visual Generation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/visual_gen/offloading.py (1)
377-377: ⚡ Quick winAdd
strict=Truetozip()for safer iteration.The
zip()call pairsself.stageswithself._stage_names, which are constructed together and should always have the same length. Addingstrict=Truewill catch logic errors if they ever fall out of sync.🔒 Proposed fix
- for stage, group_name in zip(self.stages, self._stage_names): + for stage, group_name in zip(self.stages, self._stage_names, strict=True):Based on learnings: In TensorRT-LLM, Python 3.10+ features are available throughout the codebase.
🤖 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/offloading.py` at line 377, The loop uses zip(self.stages, self._stage_names) which silently truncates if the two iterables ever differ in length; update the iteration to zip(..., strict=True) to enforce they stay matched and raise an error if they ever fall out of sync—modify the loop where zip(self.stages, self._stage_names) is used (e.g., in the offloading iteration block) to pass strict=True.
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/offloading.py`:
- Line 377: The loop uses zip(self.stages, self._stage_names) which silently
truncates if the two iterables ever differ in length; update the iteration to
zip(..., strict=True) to enforce they stay matched and raise an error if they
ever fall out of sync—modify the loop where zip(self.stages, self._stage_names)
is used (e.g., in the offloading iteration block) to pass strict=True.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ce462aca-6de3-431b-810d-57c128a49b77
📒 Files selected for processing (8)
examples/visual_gen/README.mdexamples/visual_gen/visual_gen_wan_t2v.pytensorrt_llm/_torch/visual_gen/executor.pytensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.pytensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.pytensorrt_llm/_torch/visual_gen/offloading.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/_torch/visual_gen/pipeline_loader.py
Collaborator
|
Few initial thoughts -
|
NVShreyas
reviewed
May 14, 2026
NVShreyas
reviewed
May 14, 2026
NVShreyas
reviewed
May 14, 2026
Collaborator
Author
I agree that model.to(device) before/after use is simpler and is a good baseline. The reason I didn’t start there is that it has a few downsides for this use case:
So the harness is intended to trade complexity for fewer transfers, pinned CPU backing storage, and larger contiguous copies. That said, I agree the design only makes sense if it avoids stale CUDA views and remains clearly scoped. |
Collaborator
Author
This makes sense, I agree. The current invariant should be tightened so inactive/offloaded weights are either CPU-backed or fail loudly if accessed as CUDA weights. Silent stale CUDA views into a reused arena are not acceptable. I’ll work on this next week. My current direction is to ensure only the active group is bound to the GPU arena, while inactive groups are rebound to CPU storage or otherwise guarded so accidental use errors out. Happy to take pointers on whether you’d prefer CPU-backed inactive params or an explicit sentinel/error tensors. |
rahul-steiger-nv
force-pushed
the
visual_gen_offloading
branch
from
May 18, 2026 13:54
fbc06fd to
92e09f8
Compare
NVShreyas
reviewed
May 18, 2026
NVShreyas
left a comment
NVShreyas
left a comment
Collaborator
There was a problem hiding this comment.
Overall the approach is good and I'm okay with not using module.to(device) based on the explanation you provided. Lets make sure the code is architected in a maintainable and readable manner (including docstrings). It also should be tested with all single GPU features like quantization, caching, torch compile, etc (or raise NotImplementedError explicitly). Let's leave multi-gpu offloading as a future enhancement for now.
rahul-steiger-nv
force-pushed
the
visual_gen_offloading
branch
2 times, most recently
from
May 19, 2026 11:24
05f439e to
9c77ea6
Compare
Collaborator
Author
|
@coderabbitai help |
Contributor
ChatThere are 3 ways to chat with CodeRabbit:
CodeRabbit commands
Other keywords and placeholders
Status, support, documentation and community
|
Collaborator
Author
|
@coderabbitai review |
Contributor
✅ Actions performedReview triggered.
|
Collaborator
Author
|
@coderabbitai full review |
Contributor
✅ Actions performedFull review triggered. |
Collaborator
Author
|
@coderabbitai run pre-merge checks |
Contributor
✅ Actions performedRunning Pre-merge checks. Results will be posted in the summary comment once complete. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (1)
tensorrt_llm/_torch/visual_gen/offloading.py (1)
410-415:⚠️ Potential issue | 🔴 Critical | 🏗️ Heavy liftCached GPU views still leave stale arena-backed tensors reachable.
layout.gpu_viewsis created once and then rebound in and out of modules, but anyParameter/buffer reference captured while a group is active still points atself.gpu_arenaafter_rebind_to_cpu(). Once the next group stages, that old object silently reads unrelated weights instead of failing, so the stale-CUDA-view hazard called out earlier is still present.Please either recreate/poison GPU-backed view objects on each activation or add a hard guard that prevents arena-backed tensors from escaping the active call scope.
Also applies to: 425-466
🤖 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/offloading.py` around lines 410 - 415, Cached GPU views in layout.gpu_views let Parameter/buffer objects keep pointers into self.gpu_arena after _rebind_to_cpu(), causing stale-arena reads; fix by ensuring GPU-backed view objects are not reusable across activations: either (A) recreate GPU views on every activation and/or clear/poison old layout.gpu_views when deactivating, or (B) add a hard guard in _rebind_to_cpu() (and activation entry) that walks module parameters and buffers and replaces any tensor whose storage points at self.gpu_arena with a safe CPU-backed tensor or raises an error if it would escape scope; implement the change around calls to _make_views, layout.gpu_views, and _rebind_to_cpu so no arena-backed tensors can remain reachable after deactivation.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/visual_gen/config.py (1)
351-351: ⚡ Quick winAdd Pydantic field description for the new user-facing flag.
PipelineConfigis user-facing, so this new field should usePydanticFieldwith a description for schema/help consistency.♻️ Proposed fix
- enable_cuda_memory_logging: bool = False + enable_cuda_memory_logging: bool = PydanticField( + False, + description="Enable per-request CUDA peak memory logging in diffusion executor.", + )As per coding guidelines: "Add descriptions to all user-facing Pydantic fields via
Field(description="...")rather than using comments".🤖 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/config.py` at line 351, The new user-facing flag enable_cuda_memory_logging is declared as a bare bool; update PipelineConfig to use PydanticField with a description (e.g. enable_cuda_memory_logging: bool = PydanticField(False, description="Enable detailed CUDA memory logging for debugging and profiling") ) so it appears in the generated schema/help; also add or ensure the PydanticField import is present and follow the same description style used by other PipelineConfig fields.tests/unittest/_torch/visual_gen/test_offloading.py (1)
1-268: No QA integration test-list update is needed for this change set.This PR scope adds unit tests under
tests/unittest/..., so updatingtests/integration/test_lists/qa/*is unnecessary.As per coding guidelines, “If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional.”
🤖 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/visual_gen/test_offloading.py` around lines 1 - 268, Add an explicit note to the PR description or commit message stating that this change only adds unit tests (test_offloading.py) and therefore no QA integration test-list updates are required; include the exact sentence "No QA integration test-list update is needed for this change set." so reviewers and release tooling can detect it, and update any PR checklist/description fields that record whether integration QA lists were modified.
🤖 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/executor.py`:
- Around line 349-355: Change the broad except Exception in the CUDA-memory
helper try/except blocks to catch RuntimeError only: for the block calling
torch.cuda.reset_peak_memory_stats(self.device_id) and the similar block around
lines handling CUDA memory (the second try/except referenced in the comment),
replace "except Exception as e" with "except RuntimeError as e" so only PyTorch
CUDA runtime errors are caught and other exceptions are not masked; keep the
same logger.warning message and error variable name.
In `@tensorrt_llm/_torch/visual_gen/pipeline_loader.py`:
- Around line 280-286: The code is using pipeline.default_offload_stages() which
hard-codes the full default set; instead fetch and pass the configured offload
stages via the same resolution path used by runtime init (the helper used by
initialize_offload_pipeline()). Replace the call to
pipeline.default_offload_stages() inside the
pipeline._filter_available_offload_stages(...) call with the configured-stage
resolver (e.g., the pipeline method/attribute used by
initialize_offload_pipeline(), such as pipeline.offload_stages() or
pipeline._resolve_offload_stages(...)) so the collected parts mirror the runtime
configuration and avoid materializing modules that should not be staged.
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 366-370: The method offload_context currently returns
nullcontext() when self._offload_pipeline is None, which silently masks
misordered initialization; instead detect when offloading is intended (enable
True) but _offload_pipeline is uninitialized and raise a clear exception
indicating offload pipeline not initialized (refer to offload_context and
self._offload_pipeline), instruct callers to run initialize_offload_pipeline()
first; apply the same explicit guard/exception pattern to the other places that
currently check self._offload_pipeline (the subsequent blocks surrounding
offload setup/usage) so they fail fast rather than silently falling back to CPU
behavior.
In `@tests/unittest/_torch/visual_gen/test_offloading.py`:
- Around line 153-156: The regex in the pytest.raises assertion uses
"group.weight" where the dot is a regex metacharacter; update the match argument
in the pytest.raises call in tests/unittest/_torch/visual_gen/test_offloading.py
(the block using pytest.raises(RuntimeError, match="group.weight")) to escape
the dot (e.g., "group\.weight" or use re.escape on the literal) so the assertion
matches the literal field name rather than any character.
---
Duplicate comments:
In `@tensorrt_llm/_torch/visual_gen/offloading.py`:
- Around line 410-415: Cached GPU views in layout.gpu_views let Parameter/buffer
objects keep pointers into self.gpu_arena after _rebind_to_cpu(), causing
stale-arena reads; fix by ensuring GPU-backed view objects are not reusable
across activations: either (A) recreate GPU views on every activation and/or
clear/poison old layout.gpu_views when deactivating, or (B) add a hard guard in
_rebind_to_cpu() (and activation entry) that walks module parameters and buffers
and replaces any tensor whose storage points at self.gpu_arena with a safe
CPU-backed tensor or raises an error if it would escape scope; implement the
change around calls to _make_views, layout.gpu_views, and _rebind_to_cpu so no
arena-backed tensors can remain reachable after deactivation.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/config.py`:
- Line 351: The new user-facing flag enable_cuda_memory_logging is declared as a
bare bool; update PipelineConfig to use PydanticField with a description (e.g.
enable_cuda_memory_logging: bool = PydanticField(False, description="Enable
detailed CUDA memory logging for debugging and profiling") ) so it appears in
the generated schema/help; also add or ensure the PydanticField import is
present and follow the same description style used by other PipelineConfig
fields.
In `@tests/unittest/_torch/visual_gen/test_offloading.py`:
- Around line 1-268: Add an explicit note to the PR description or commit
message stating that this change only adds unit tests (test_offloading.py) and
therefore no QA integration test-list updates are required; include the exact
sentence "No QA integration test-list update is needed for this change set." so
reviewers and release tooling can detect it, and update any PR
checklist/description fields that record whether integration QA lists were
modified.
🪄 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: 00865218-f5b0-46ae-9ce1-912e0f6d4817
⛔ Files ignored due to path filters (2)
examples/visual_gen/output_wan_fp8_offload.aviis excluded by!**/*.aviexamples/visual_gen/output_wan_fp8_offload_quality.aviis excluded by!**/*.avi
📒 Files selected for processing (9)
examples/visual_gen/README.mdexamples/visual_gen/visual_gen_wan_t2v.pytensorrt_llm/_torch/visual_gen/config.pytensorrt_llm/_torch/visual_gen/executor.pytensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.pytensorrt_llm/_torch/visual_gen/offloading.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/_torch/visual_gen/pipeline_loader.pytests/unittest/_torch/visual_gen/test_offloading.py
✅ Files skipped from review due to trivial changes (1)
- examples/visual_gen/README.md
🚧 Files skipped from review as they are similar to previous changes (1)
- examples/visual_gen/visual_gen_wan_t2v.py
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/visual_gen/executor.py (1)
309-323:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winPrevent stale CUDA peak-memory logs on early failures.
If request handling fails before the reset call, the exception-path log can report peak memory from a previous request. Reset at request start (or gate logging on a per-request reset flag).
Proposed fix
def process_request(self, req: DiffusionRequest): """Process a single request.""" log_cuda_memory = self._cuda_memory_logging_enabled() + did_reset_peak_stats = False try: + if log_cuda_memory: + self._reset_cuda_peak_memory_stats() + did_reset_peak_stats = True self._merge_defaults(req) cache_key = self.pipeline.warmup_cache_key( req.params.height, req.params.width, num_frames=req.params.num_frames ) @@ - if log_cuda_memory: - self._reset_cuda_peak_memory_stats() output = self.pipeline.infer(req) - if log_cuda_memory: + if did_reset_peak_stats: self._log_cuda_peak_memory(req.request_id) @@ - if log_cuda_memory: + if did_reset_peak_stats: self._log_cuda_peak_memory(req.request_id)Also applies to: 325-326, 330-331
🤖 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/executor.py` around lines 309 - 323, Move the per-request CUDA peak-memory reset to the start of request handling so stale peak-memory values cannot be logged on early failures: call self._reset_cuda_peak_memory_stats() immediately after determining log_cuda_memory with self._cuda_memory_logging_enabled() (before calling self._merge_defaults or any code that can raise), and ensure the same pattern is applied for the other CUDA-logging sites in this file (the blocks that reference self._reset_cuda_peak_memory_stats() later) so logging is gated by a per-request reset rather than relying on post-success cleanup.
♻️ Duplicate comments (1)
tests/unittest/_torch/visual_gen/test_offloading.py (1)
153-156:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winEscape the dot in the
pytest.raises(match=...)pattern.
group.weightis interpreted as regex (.= any char), so this assertion can pass on unintended messages.Proposed fix
with pytest.raises( RuntimeError, - match="Failed to copy offload tensor 'group.weight'", + match=r"Failed to copy offload tensor 'group\.weight'", ) as exc_info:🤖 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/visual_gen/test_offloading.py` around lines 153 - 156, The regex in the pytest.raises call (pytest.raises(..., match="Failed to copy offload tensor 'group.weight'")) treats the dot in "group.weight" as a wildcard; update the match to escape the dot so it matches a literal period (e.g., use a raw string with the dot escaped or use re.escape on "group.weight") in the pytest.raises call in test_offloading.py to ensure the assertion only matches the exact message.
🧹 Nitpick comments (1)
tests/unittest/_torch/visual_gen/test_offloading.py (1)
1-268: QA list update scope is unnecessary for this PR surface.These changes are confined to
tests/unittest/...; notests/integration/test_lists/qa/*update is needed in this PR.As per coding guidelines, "If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional."
🤖 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/visual_gen/test_offloading.py` around lines 1 - 268, The PR touches only unit test files under tests/unittest and therefore does not require updating the QA lists; update the PR description (or add a short commit/PR note) to explicitly state that "QA list updates are unnecessary" because the changes are confined to tests/unittest (per the guideline to state whether QA list updates are unnecessary or optional), referencing this scope and the fact no tests/integration/test_lists/qa/* files were modified.
🤖 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.
Outside diff comments:
In `@tensorrt_llm/_torch/visual_gen/executor.py`:
- Around line 309-323: Move the per-request CUDA peak-memory reset to the start
of request handling so stale peak-memory values cannot be logged on early
failures: call self._reset_cuda_peak_memory_stats() immediately after
determining log_cuda_memory with self._cuda_memory_logging_enabled() (before
calling self._merge_defaults or any code that can raise), and ensure the same
pattern is applied for the other CUDA-logging sites in this file (the blocks
that reference self._reset_cuda_peak_memory_stats() later) so logging is gated
by a per-request reset rather than relying on post-success cleanup.
---
Duplicate comments:
In `@tests/unittest/_torch/visual_gen/test_offloading.py`:
- Around line 153-156: The regex in the pytest.raises call (pytest.raises(...,
match="Failed to copy offload tensor 'group.weight'")) treats the dot in
"group.weight" as a wildcard; update the match to escape the dot so it matches a
literal period (e.g., use a raw string with the dot escaped or use re.escape on
"group.weight") in the pytest.raises call in test_offloading.py to ensure the
assertion only matches the exact message.
---
Nitpick comments:
In `@tests/unittest/_torch/visual_gen/test_offloading.py`:
- Around line 1-268: The PR touches only unit test files under tests/unittest
and therefore does not require updating the QA lists; update the PR description
(or add a short commit/PR note) to explicitly state that "QA list updates are
unnecessary" because the changes are confined to tests/unittest (per the
guideline to state whether QA list updates are unnecessary or optional),
referencing this scope and the fact no tests/integration/test_lists/qa/* files
were modified.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1e658a24-a4fc-4f38-bedc-a2b2872a2bd5
⛔ Files ignored due to path filters (2)
examples/visual_gen/output_wan_fp8_offload.aviis excluded by!**/*.aviexamples/visual_gen/output_wan_fp8_offload_quality.aviis excluded by!**/*.avi
📒 Files selected for processing (9)
examples/visual_gen/README.mdexamples/visual_gen/visual_gen_wan_t2v.pytensorrt_llm/_torch/visual_gen/config.pytensorrt_llm/_torch/visual_gen/executor.pytensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.pytensorrt_llm/_torch/visual_gen/offloading.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/_torch/visual_gen/pipeline_loader.pytests/unittest/_torch/visual_gen/test_offloading.py
chang-l
approved these changes
Jul 29, 2026
LarryXFly
approved these changes
Jul 30, 2026
Contributor
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62805 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62805 [ run ] completed with state
|
Contributor
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62849 [ run ] triggered by Bot. Commit: |
Contributor
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62849 [ run ] completed with state
|
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
arysef
approved these changes
Jul 31, 2026
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Collaborator
|
/bot run --disable-fail-fast |
Contributor
|
Maybe already obsolete, but:
The chain:
A/B confirms it:
Fix The semantic change is correct — under offloading the transformer can sit on CPU, so pipeline = object.__new__(ltx2_two_stages.LTX2TwoStagesPipeline)
pipeline._device = torch.device("cuda") # __init__ is bypassed; device property reads this |
Collaborator
|
PR_Github #63078 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63078 [ run ] completed with state
|
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63134 [ run ] triggered by Bot. Commit: |
Collaborator
|
@rahul-steiger-nv Please address the deterministic LTX2 unit-test regression identified in #14095 (comment) before another CI retry. The PR correctly changes |
Collaborator
|
PR_Github #63134 [ run ] completed with state
|
Signed-off-by: Rahul Steiger <rsteiger@nvidia.com>
Collaborator
Author
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63301 [ run ] triggered by Bot. Commit: |
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
11 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.
Description
This PR adds initial offloading support for
tensorrt_llm/_torch/visual_gen.The goal is to reduce peak GPU memory usage for memory-heavy visual generation pipelines by allowing selected model components, auxiliary modules, or guardrails to be offloaded when they are not actively needed during generation.
This is intended to help support larger visual generation workloads, longer video generation, multi-view generation, and lower-VRAM GPUs.
This PR is an initial implementation for the feature request in #13893.
Multi-GPU support and quantization compatibility are still being validated. There may also be some design changes after further discussion to make the implementation more general and easier to maintain upstream.
Test Coverage
Current testing:
Planned / in-progress testing:
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
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.Summary by CodeRabbit
Release Notes
New Features
--enable_offloading,--offload_stages, and--enable_cuda_memory_logging.Documentation & Examples
Performance results
Measured with
Wan-AI/Wan2.2-T2V-A14B-Diffusersat 480x832 resolution and 33 frames.In this single-run measurement, CPU offloading reduced peak CUDA memory by 36.92 GiB, from 68.54 GiB to 31.62 GiB, which is about a 54% reduction. End-to-end pipeline time remained effectively flat, with only 0.32s additional runtime, or about 0.4% overhead, excluding the one-time offloading setup cost.
This enables Wan-AI/Wan2.2-T2V-A14B-Diffusers to run comfortably on L40 and A6000-class 48 GB GPUs. RTX 5090 support may be feasible with careful tuning, but 32 GB remains a marginal target due to allocator fragmentation and limited headroom at higher resolutions or frame counts.