From 094031564b75f087905030c22d2d13825b047765 Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Fri, 8 May 2026 01:28:51 +0000 Subject: [PATCH 1/3] parallel LoRA weight loading Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- .../models/ltx2/pipeline_ltx2_two_stages.py | 133 +++++++++++++----- 1 file changed, 100 insertions(+), 33 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index f7eefe59f3cb..065a698b1b70 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -4,7 +4,9 @@ import json import math +import os import time +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional, Union @@ -49,6 +51,8 @@ _FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) # Baseline BF16 peak memory ~75 GiB, saving BF16 weights snapshot total ~108 GiB. _BF16_WEIGHTS_SNAPSHOT_FREE_MEMORY_THRESHOLD_GIB = 115.0 +_QKV_SUFFIXES = (".to_q", ".to_k", ".to_v") +_DISABLE_OVERLAP_LORA_LOAD_ENV = "TRTLLM_LTX2_DISABLE_OVERLAP_LORA_LOAD" # --------------------------------------------------------------------------- @@ -99,6 +103,36 @@ def _should_save_bf16_weights( return save_state +def _map_lora_param_name(base_name: str, strip_prefixes: List[str]) -> str: + param_name = base_name + for prefix in strip_prefixes: + if param_name.startswith(prefix): + param_name = param_name[len(prefix) :] + break + + # Apply the same key remapping as LTXModel.load_weights() so LoRA delta + # keys match TRT-LLM parameter names. + for ff_prefix in (".ff.", ".audio_ff."): + if ff_prefix + "net.0.proj" in param_name: + param_name = param_name.replace(ff_prefix + "net.0.proj", ff_prefix + "up_proj") + elif ff_prefix + "net.2" in param_name: + param_name = param_name.replace(ff_prefix + "net.2", ff_prefix + "down_proj") + param_name = param_name.replace(".q_norm.", ".norm_q.") + param_name = param_name.replace(".k_norm.", ".norm_k.") + return param_name + + +def _has_lora_target(param_name: str, model_params: set[str]) -> bool: + if param_name in model_params or f"{param_name}.weight" in model_params: + return True + + for suffix in _QKV_SUFFIXES: + if param_name.endswith(suffix): + attn_prefix = param_name[: -len(suffix)] + return f"{attn_prefix}.qkv_proj.weight" in model_params + return False + + def _load_lora_deltas( lora_path: str, transformer: torch.nn.Module, @@ -163,39 +197,28 @@ def _strip(key, suffixes): if suffix: strip_prefixes.append(suffix) + model_params = {n for n, _ in transformer.named_parameters()} deltas: Dict[str, torch.Tensor] = {} + skipped_non_targets = 0 for base_name in down_keys: if base_name not in up_keys: continue + + param_name = _map_lora_param_name(base_name, strip_prefixes) + if not _has_lora_target(param_name, model_params): + skipped_non_targets += 1 + continue + A = down_keys[base_name] # (rank, in_features) B = up_keys[base_name] # (out_features, rank) rank = A.shape[0] alpha = alpha_dict.get(base_name, float(rank)) scale = strength * alpha / rank - - param_name = base_name - for prefix in strip_prefixes: - if param_name.startswith(prefix): - param_name = param_name[len(prefix) :] - break - - # Apply the same key remapping as LTXModel.load_weights() so - # that LoRA delta keys match TRT-LLM parameter names. - for ff_prefix in (".ff.", ".audio_ff."): - if ff_prefix + "net.0.proj" in param_name: - param_name = param_name.replace(ff_prefix + "net.0.proj", ff_prefix + "up_proj") - elif ff_prefix + "net.2" in param_name: - param_name = param_name.replace(ff_prefix + "net.2", ff_prefix + "down_proj") - param_name = param_name.replace(".q_norm.", ".norm_q.") - param_name = param_name.replace(".k_norm.", ".norm_k.") - delta = (B.float() @ A.float()) * scale deltas[param_name] = delta # Fuse separate to_q / to_k / to_v deltas into a single qkv_proj # delta when the transformer uses QKV fusion (FUSE_QKV mode). - model_params = {n for n, _ in transformer.named_parameters()} - _QKV_SUFFIXES = (".to_q", ".to_k", ".to_v") fused_keys: set = set() qkv_groups: Dict[str, List[str]] = {} for key in list(deltas.keys()): @@ -223,7 +246,10 @@ def _strip(key, suffixes): for key in fused_keys: del deltas[key] - logger.info(f"Loaded {len(deltas)} LoRA deltas from {lora_path} (strength={strength})") + logger.info( + f"Loaded {len(deltas)} LoRA deltas from {lora_path} " + f"(strength={strength}, skipped_non_targets={skipped_non_targets})" + ) return deltas @@ -967,17 +993,52 @@ def load_standard_components( # The BF16 snapshot threshold is a whole-pipeline capacity gate, so # record it before loading model/runtime components that consume HBM. self._bf16_snapshot_preload_free_gib = _get_free_gpu_memory_gib(device=device) - super().load_standard_components( - checkpoint_dir, - device, - skip_components=skip_components, - **kwargs, - ) - dtype = self.pipeline_config.torch_dtype spatial_upsampler_path = self.pipeline_config.extra_attrs.get("spatial_upsampler_path", "") distilled_lora_path = self.pipeline_config.extra_attrs.get("distilled_lora_path", "") + lora_executor = None + lora_future = None + disable_overlap_lora_load = os.getenv(_DISABLE_OVERLAP_LORA_LOAD_ENV, "0") == "1" + if distilled_lora_path and not disable_overlap_lora_load: + logger.info(f"Starting distilled LoRA pre-compute from {distilled_lora_path}...") + lora_executor = ThreadPoolExecutor(max_workers=1) + lora_future = lora_executor.submit( + _load_lora_deltas, + distilled_lora_path, + self.transformer, + self._TRANSFORMER_PREFIX, + ) + + try: + super().load_standard_components( + checkpoint_dir, + device, + skip_components=skip_components, + **kwargs, + ) + + self._load_two_stage_components( + device, + dtype, + spatial_upsampler_path, + distilled_lora_path, + lora_future, + disable_overlap_lora_load, + ) + finally: + if lora_executor is not None: + lora_executor.shutdown(wait=True) + + def _load_two_stage_components( + self, + device: torch.device, + dtype: torch.dtype, + spatial_upsampler_path: str, + distilled_lora_path: str, + lora_future, + disable_overlap_lora_load: bool, + ) -> None: # --- Spatial upsampler --- if spatial_upsampler_path: logger.info(f"Loading spatial upsampler from {spatial_upsampler_path}...") @@ -1018,12 +1079,18 @@ def load_standard_components( self._distilled_lora_deltas: Dict[str, torch.Tensor] = {} self._distilled_lora_weight_cache: Optional[_PersistentLoRAWeightCache] = None if distilled_lora_path: - logger.info(f"Loading distilled LoRA from {distilled_lora_path}...") - self._distilled_lora_deltas = _load_lora_deltas( - distilled_lora_path, - self.transformer, - transformer_prefix=self._TRANSFORMER_PREFIX, - ) + if disable_overlap_lora_load: + logger.info(f"Loading distilled LoRA from {distilled_lora_path}...") + self._distilled_lora_deltas = _load_lora_deltas( + distilled_lora_path, + self.transformer, + self._TRANSFORMER_PREFIX, + ) + else: + logger.info("Waiting for distilled LoRA pre-compute...") + if lora_future is None: + raise RuntimeError("Distilled LoRA pre-compute was not started.") + self._distilled_lora_deltas = lora_future.result() logger.info( f"Distilled LoRA ready: {len(self._distilled_lora_deltas)} parameter deltas" ) From af5a827491163f918ba025a5c9b63cfc2ea9e07c Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Fri, 8 May 2026 18:05:32 +0000 Subject: [PATCH 2/3] add tests and remove env variable Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- .../models/ltx2/pipeline_ltx2_two_stages.py | 23 ++----- .../_torch/visual_gen/test_ltx2_pipeline.py | 60 +++++++++++++++++++ 2 files changed, 65 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index 065a698b1b70..841f8e0e90c3 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -4,7 +4,6 @@ import json import math -import os import time from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field @@ -52,7 +51,6 @@ # Baseline BF16 peak memory ~75 GiB, saving BF16 weights snapshot total ~108 GiB. _BF16_WEIGHTS_SNAPSHOT_FREE_MEMORY_THRESHOLD_GIB = 115.0 _QKV_SUFFIXES = (".to_q", ".to_k", ".to_v") -_DISABLE_OVERLAP_LORA_LOAD_ENV = "TRTLLM_LTX2_DISABLE_OVERLAP_LORA_LOAD" # --------------------------------------------------------------------------- @@ -999,8 +997,7 @@ def load_standard_components( lora_executor = None lora_future = None - disable_overlap_lora_load = os.getenv(_DISABLE_OVERLAP_LORA_LOAD_ENV, "0") == "1" - if distilled_lora_path and not disable_overlap_lora_load: + if distilled_lora_path: logger.info(f"Starting distilled LoRA pre-compute from {distilled_lora_path}...") lora_executor = ThreadPoolExecutor(max_workers=1) lora_future = lora_executor.submit( @@ -1024,7 +1021,6 @@ def load_standard_components( spatial_upsampler_path, distilled_lora_path, lora_future, - disable_overlap_lora_load, ) finally: if lora_executor is not None: @@ -1037,7 +1033,6 @@ def _load_two_stage_components( spatial_upsampler_path: str, distilled_lora_path: str, lora_future, - disable_overlap_lora_load: bool, ) -> None: # --- Spatial upsampler --- if spatial_upsampler_path: @@ -1079,18 +1074,10 @@ def _load_two_stage_components( self._distilled_lora_deltas: Dict[str, torch.Tensor] = {} self._distilled_lora_weight_cache: Optional[_PersistentLoRAWeightCache] = None if distilled_lora_path: - if disable_overlap_lora_load: - logger.info(f"Loading distilled LoRA from {distilled_lora_path}...") - self._distilled_lora_deltas = _load_lora_deltas( - distilled_lora_path, - self.transformer, - self._TRANSFORMER_PREFIX, - ) - else: - logger.info("Waiting for distilled LoRA pre-compute...") - if lora_future is None: - raise RuntimeError("Distilled LoRA pre-compute was not started.") - self._distilled_lora_deltas = lora_future.result() + logger.info("Waiting for distilled LoRA pre-compute...") + if lora_future is None: + raise RuntimeError("Distilled LoRA pre-compute was not started.") + self._distilled_lora_deltas = lora_future.result() logger.info( f"Distilled LoRA ready: {len(self._distilled_lora_deltas)} parameter deltas" ) diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index 657a22de3c4d..e03a7314ef4c 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py @@ -960,6 +960,66 @@ def test_cache_dit_config_prevents_promotion_with_explicit_auxiliary_paths( assert LTX2Pipeline.resolve_variant(config) is LTX2Pipeline +class TestTwoStageLoRALoadingOverlap: + """Test two-stage LoRA load orchestration without loading checkpoints.""" + + def test_load_standard_components_overlaps_lora_with_base_components(self, monkeypatch): + from types import SimpleNamespace + + from tensorrt_llm._torch.visual_gen.models.ltx2 import pipeline_ltx2_two_stages + from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2Pipeline + from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2_two_stages import ( + LTX2TwoStagesPipeline, + ) + + events = [] + lora_deltas = {"linear.weight": torch.ones(1)} + + class FakeFuture: + def result(self): + events.append("lora_result") + return lora_deltas + + class FakeExecutor: + def __init__(self, max_workers): + assert max_workers == 1 + + def submit(self, fn, *args): + events.append("lora_submit") + assert fn is pipeline_ltx2_two_stages._load_lora_deltas + assert args == ( + "/fake/lora.safetensors", + pipeline.transformer, + pipeline._TRANSFORMER_PREFIX, + ) + return FakeFuture() + + def shutdown(self, wait=True): + events.append(f"shutdown:{wait}") + + def fake_base_load(self, checkpoint_dir, device, skip_components=None, **kwargs): + events.append("base_load") + assert checkpoint_dir == "/fake/checkpoint" + assert device == torch.device("cpu") + + monkeypatch.setattr(pipeline_ltx2_two_stages, "ThreadPoolExecutor", FakeExecutor) + monkeypatch.setattr(LTX2Pipeline, "load_standard_components", fake_base_load) + + pipeline = LTX2TwoStagesPipeline.__new__(LTX2TwoStagesPipeline) + pipeline.model_config = SimpleNamespace( + torch_dtype=torch.float32, + extra_attrs={ + "distilled_lora_path": "/fake/lora.safetensors", + }, + ) + pipeline.transformer = object() + + pipeline.load_standard_components("/fake/checkpoint", torch.device("cpu")) + + assert events == ["lora_submit", "base_load", "lora_result", "shutdown:True"] + assert pipeline._distilled_lora_deltas is lora_deltas + + class TestTwoStageUpsamplerBuildingBlocks: """Test upsampler components without checkpoints (random weights).""" From 558a64888b9db1c12ef9b800c674c9940b181f69 Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Thu, 2 Jul 2026 03:58:13 +0000 Subject: [PATCH 3/3] Fix LTX-2 async LoRA loading test Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> --- .../models/ltx2/pipeline_ltx2_two_stages.py | 4 +++- .../_torch/visual_gen/test_ltx2_pipeline.py | 17 +++++++++++++++-- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index 841f8e0e90c3..b0b661fd7b56 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -149,7 +149,9 @@ def _load_lora_deltas( sft_paths = _find_safetensors_files(lora_path) if not sft_paths: raise ValueError(f"No safetensors files found at {lora_path}") - _prefetch_ltx2_safetensors_files(sft_paths) + # This helper can run in a background thread while base components load. + # Avoid distributed prefetch collectives here; every rank must enter those + # from the same foreground load sequence to avoid hangs. raw: Dict[str, torch.Tensor] = {} alpha_dict: Dict[str, float] = {} diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index e03a7314ef4c..d2bbf46d32a5 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py @@ -1003,16 +1003,29 @@ def fake_base_load(self, checkpoint_dir, device, skip_components=None, **kwargs) assert device == torch.device("cpu") monkeypatch.setattr(pipeline_ltx2_two_stages, "ThreadPoolExecutor", FakeExecutor) + monkeypatch.setattr( + pipeline_ltx2_two_stages, + "_get_free_gpu_memory_gib", + lambda device=None: None, + ) monkeypatch.setattr(LTX2Pipeline, "load_standard_components", fake_base_load) + class FakeTransformer: + def named_modules(self): + return [] + + def named_parameters(self): + return [] + pipeline = LTX2TwoStagesPipeline.__new__(LTX2TwoStagesPipeline) - pipeline.model_config = SimpleNamespace( + pipeline.pipeline_config = SimpleNamespace( torch_dtype=torch.float32, extra_attrs={ "distilled_lora_path": "/fake/lora.safetensors", }, ) - pipeline.transformer = object() + pipeline.model_config = pipeline.pipeline_config + pipeline.transformer = FakeTransformer() pipeline.load_standard_components("/fake/checkpoint", torch.device("cpu"))