[TRTLLM-12647][feat] Parallelize LTX-2 LoRA weight loading (#13911)

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Author: yibinl-nvidia

tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py

@@ -5,6 +5,7 @@
 import json
 import math
 import time
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from typing import Any, Callable, Dict, List, Optional, Union
 
@@ -49,6 +50,7 @@
 _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")
 
 
 # ---------------------------------------------------------------------------
@@ -99,6 +101,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,
@@ -117,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] = {}
@@ -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,49 @@ 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
+        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(
+                _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,
+            )
+        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,
+    ) -> None:
         # --- Spatial upsampler ---
         if spatial_upsampler_path:
             logger.info(f"Loading spatial upsampler from {spatial_upsampler_path}...")
@@ -1018,12 +1076,10 @@ 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,
-            )
+            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"
             )

tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py

@@ -960,6 +960,79 @@ 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(
+            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.pipeline_config = SimpleNamespace(
+            torch_dtype=torch.float32,
+            extra_attrs={
+                "distilled_lora_path": "/fake/lora.safetensors",
+            },
+        )
+        pipeline.model_config = pipeline.pipeline_config
+        pipeline.transformer = FakeTransformer()
+
+        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)."""