[TRTLLM-12647][feat] Parallelize LTX-2 LoRA weight loading#13911
[TRTLLM-12647][feat] Parallelize LTX-2 LoRA weight loading#13911yibinl-nvidia wants to merge 3 commits into
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 | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. 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
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
2c2f177 to
af700c8
Compare
Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
|
/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 |
2 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%.
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.