[TRTLLM-12647][feat] Parallelize LTX-2 LoRA weight loading#13911
Conversation
📝 WalkthroughWalkthroughThis PR refactors LoRA delta loading in the LTX2 visual generation pipeline by introducing parameter mapping helpers and enabling optional background pre-computation overlapped with component loading via ChangesLoRA Delta Loading Optimization with Threading
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 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.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py (1)
671-679: 💤 Low valueAdd type annotation for
lora_futureparameter.The
lora_futureparameter lacks a type annotation, which violates the coding guideline requiring type annotations for all function arguments. As per coding guidelines, Python code should use type annotations for all function arguments and return types.Suggested fix
+from concurrent.futures import Future, ThreadPoolExecutor ... def _load_two_stage_components( self, device: torch.device, dtype: torch.dtype, spatial_upsampler_path: str, distilled_lora_path: str, - lora_future, + lora_future: Optional[Future[Dict[str, torch.Tensor]]], disable_overlap_lora_load: bool, ) -> 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/models/ltx2/pipeline_ltx2_two_stages.py` around lines 671 - 679, The parameter lora_future in _load_two_stage_components lacks a type annotation; update the function signature to annotate it (e.g., use concurrent.futures.Future or asyncio.Future depending on implementation) like lora_future: concurrent.futures.Future, add the corresponding import (from concurrent.futures import Future or import concurrent.futures) at the top of the module, and adjust to Optional[Future] if the value can be None—ensure the annotation matches the actual future type used elsewhere in this file.
🤖 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/models/ltx2/pipeline_ltx2_two_stages.py`:
- Around line 671-679: The parameter lora_future in _load_two_stage_components
lacks a type annotation; update the function signature to annotate it (e.g., use
concurrent.futures.Future or asyncio.Future depending on implementation) like
lora_future: concurrent.futures.Future, add the corresponding import (from
concurrent.futures import Future or import concurrent.futures) at the top of the
module, and adjust to Optional[Future] if the value can be None—ensure the
annotation matches the actual future type used elsewhere in this file.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e4e89d16-9958-418e-bdca-0485cd5a82a6
📒 Files selected for processing (1)
tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py
|
/bot run |
|
PR_Github #47432 [ run ] triggered by Bot. Commit: |
|
PR_Github #47432 [ run ] completed with state |
|
/bot run |
|
PR_Github #47441 [ run ] triggered by Bot. Commit: |
|
PR_Github #47441 [ run ] completed with state
|
|
/bot run |
|
PR_Github #47794 [ run ] triggered by Bot. Commit: |
|
PR_Github #47794 [ run ] completed with state |
yibinl-nvidia
changed the title
yibinl-nvidia
changed the title
yibinl-nvidia
force-pushed
the
dev-yibinl-TRT-12527
branch
from
2c2f177 to
af700c8
Compare
|
/bot run |
|
PR_Github #48444 [ run ] triggered by Bot. Commit: |
|
PR_Github #48444 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #48459 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #48492 [ run ] triggered by Bot. Commit: |
|
PR_Github #48459 [ run ] completed with state
|
|
PR_Github #48492 [ run ] completed with state |
yibinl-nvidia
force-pushed
the
dev-yibinl-TRT-12527
branch
2 times, most recently
from
e417e26 to
22444f6
Compare
|
/bot run |
|
PR_Github #51639 [ run ] triggered by Bot. Commit: |
|
PR_Github #51721 [ run ] triggered by Bot. Commit: |
|
PR_Github #51721 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51743 [ run ] triggered by Bot. Commit: |
|
PR_Github #51743 [ run ] completed with state |
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
yibinl-nvidia
force-pushed
the
dev-yibinl-TRT-12527
branch
from
22444f6 to
af5a827
Compare
|
/bot run |
|
PR_Github #56396 [ run ] triggered by Bot. Commit: |
|
PR_Github #56396 [ run ] completed with state
|
|
/bot run |
|
PR_Github #56766 [ run ] triggered by Bot. Commit: |
|
PR_Github #56766 [ run ] completed with state
|
|
/bot run |
|
PR_Github #57022 [ run ] triggered by Bot. Commit: |
|
PR_Github #57022 [ run ] completed with state
|
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #57337 [ run ] triggered by Bot. Commit: |
|
PR_Github #57337 [ run ] completed with state
|
|
/bot run |
|
PR_Github #57439 [ run ] triggered by Bot. Commit: |
|
PR_Github #57439 [ run ] completed with state |
3 participants
Summary by CodeRabbit
Refactor
Performance
Description
Overlap LoRA weights loading with other components to speed up LTX-2 model loading time. This reduces model load time as much as 30%.
A typical model weight loading timestamp
15:50:52 LoRA starts
15:50:52 tokenizer starts
15:50:53 text encoder starts
15:50:57 video decoder starts
15:50:59 audio decoder / vocoder / text connectors
15:51:05 video encoder starts
15:51:08 spatial upsampler starts
15:51:10 LoRA deltas loaded
15:51:10 spatial upsampler loaded
15:51:10 Distilled LoRA ready
15:51:14 persistent cache ready
15:51:14 model loaded
about 18s LoRA weight overlapping with other components loading.
Test Coverage
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.