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..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 @@ -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" ) diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index 657a22de3c4d..d2bbf46d32a5 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/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)."""