[TRTLLM-12527][feat] Parallelize multi-shard visual-gen checkpoint loading and pre-fetch checking

tensorrt_llm/_torch/visual_gen/checkpoints/prefetch.py

@@ -0,0 +1,165 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Host page-cache prefetch helpers for visual generation checkpoints."""
+
+import multiprocessing
+import os
+from concurrent.futures import ThreadPoolExecutor
+from typing import Iterable, Optional, Set
+
+import psutil
+import torch.distributed as dist
+
+from tensorrt_llm.logger import logger
+
+_PREFETCH_CHUNK_SIZE = 16 * 1024 * 1024
+
+
+def _dist_initialized() -> bool:
+    return dist.is_available() and dist.is_initialized()
+
+
+def _dist_barrier() -> None:
+    if _dist_initialized():
+        dist.barrier()
+
+
+def _get_int_env(name: str, default: int) -> int:
+    try:
+        return int(os.environ.get(name, ""))
+    except ValueError:
+        return default
+
+
+def _local_rank_and_size() -> tuple[int, int]:
+    if not _dist_initialized():
+        return 0, 1
+
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    local_rank = _get_int_env("LOCAL_RANK", rank)
+    local_size = _get_int_env("LOCAL_WORLD_SIZE", world_size)
+
+    if local_size < 1 or local_rank < 0 or local_rank >= local_size:
+        return rank, world_size
+    return local_rank, local_size
+
+
+def _get_local_available_host_memory() -> int:
+    """Return the minimum available host memory observed by local ranks.
+
+    The prefetch/skip decision must be the same for all ranks that synchronize
+    at the post-prefetch barrier. Use torch.distributed to collect per-rank
+    snapshots and reduce the local slice to its minimum.
+    """
+    available_memory = psutil.virtual_memory().available
+    if not _dist_initialized() or dist.get_world_size() == 1:
+        return available_memory
+
+    world_size = dist.get_world_size()
+    gathered_memory: list[int | None] = [None] * world_size
+    dist.all_gather_object(gathered_memory, int(available_memory))
+
+    local_rank, local_size = _local_rank_and_size()
+    local_start = dist.get_rank() - local_rank
+    local_end = min(local_start + local_size, world_size)
+    local_memory = [
+        memory for memory in gathered_memory[local_start:local_end] if memory is not None
+    ]
+    if not local_memory:
+        return available_memory
+    return min(local_memory)
+
+
+def _normalize_paths(
+    file_names: Iterable[str],
+    prefetched_paths: Optional[Set[str]],
+) -> list[str]:
+    paths: list[str] = []
+    seen: set[str] = set()
+    for file_name in file_names:
+        path = os.path.abspath(file_name)
+        if path in seen:
+            continue
+        seen.add(path)
+        if prefetched_paths is not None and path in prefetched_paths:
+            continue
+        paths.append(path)
+    return paths
+
+
+def _prefetch_file(file_name: str, description: str) -> None:
+    if not os.path.exists(file_name):
+        return
+
+    logger.info(f"Prefetching {description} file {file_name} to host page cache...")
+    with open(file_name, "rb") as f:
+        while f.read(_PREFETCH_CHUNK_SIZE):
+            pass
+    logger.info(f"Finished prefetching {description} file {file_name}.")
+
+
+def prefetch_files_to_host_cache(
+    file_names: Iterable[str],
+    *,
+    description: str,
+    prefetched_paths: Optional[Set[str]] = None,
+    ignore_errors: bool = False,
+) -> bool:
+    """Warm checkpoint files in host page cache across distributed local ranks.
+
+    Returns True only when all selected files were already prefetched or were
+    prefetched successfully. If prefetch is skipped or fails, returns False
+    when ignore_errors=True and raises otherwise.
+    """
+    paths = _normalize_paths(file_names, prefetched_paths)
+    success = False
+    try:
+        if not paths:
+            success = True
+            return True
+
+        prefetch_size = sum(os.path.getsize(path) for path in paths if os.path.exists(path))
+        available_memory = _get_local_available_host_memory()
+        if prefetch_size >= available_memory * 0.9:
+            logger.info(
+                f"Skipping {description} prefetch because files require "
+                f"{prefetch_size / (1024**3):.2f}GB and available host memory is "
+                f"{available_memory / (1024**3):.2f}GB."
+            )
+            return False
+
+        local_rank, local_size = _local_rank_and_size()
+        local_paths = paths[local_rank::local_size]
+        if local_paths:
+            logger.info(
+                f"Prefetching {prefetch_size / (1024**3):.2f}GB {description} "
+                "files across distributed local ranks."
+            )
+            max_workers = min(multiprocessing.cpu_count() * 2, 16, len(local_paths))
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                list(executor.map(lambda path: _prefetch_file(path, description), local_paths))
+
+        success = True
+        return True
+    except Exception as exc:
+        if not ignore_errors:
+            raise
+        logger.warning(f"{description} prefetch failed; continuing without prefetch: {exc}")
+        return False
+    finally:
+        if success and prefetched_paths is not None:
+            prefetched_paths.update(paths)
+        _dist_barrier()

tensorrt_llm/_torch/visual_gen/checkpoints/weight_loader.py

@@ -1,13 +1,15 @@
 """Weight loader for diffusion models."""
 
 import json
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 from typing import Any, Dict, List, Union
 
 import torch
 import tqdm
 
 from tensorrt_llm._torch.models.checkpoints.base_weight_loader import BaseWeightLoader
+from tensorrt_llm._torch.visual_gen.checkpoints.prefetch import prefetch_files_to_host_cache
 from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
@@ -18,7 +20,7 @@ class WeightLoader(BaseWeightLoader):
     Weight loader for diffusion models.
 
     Loads weights from safetensors/bin files, similar to HfWeightLoader
-    but simpler (no parallel loading optimization for now).
+    but tailored for diffusion checkpoint layouts.
 
     Supports loading multiple components (e.g., transformer and transformer_2):
         loader = WeightLoader(components=["transformer", "transformer_2"])
@@ -87,12 +89,13 @@ def load_weights(
             if not weight_files:
                 raise ValueError(f"No weight files found in {weight_dir}")
 
-            # Load all weights with progress bar
-            component_weights = {}
-            desc = f"Loading {component}" if is_pipeline else "Loading checkpoint"
-            for wf in tqdm.tqdm(weight_files, desc=desc):
-                component_weights.update(self._load_file(wf))
+            if all(wf.endswith(".safetensors") for wf in weight_files):
+                prefetch_files_to_host_cache(
+                    weight_files,
+                    description="visual-gen checkpoint",
+                )
 
+            component_weights = self._load_weight_files(weight_files, component, is_pipeline)
             all_weights[component] = component_weights
 
         # Return flat dict for single component (backward compatibility)
@@ -102,6 +105,32 @@ def load_weights(
         # Return nested dict for multiple components
         return all_weights
 
+    def _load_weight_files(
+        self, weight_files: List[str], component: str, is_pipeline: bool
+    ) -> Dict[str, Any]:
+        desc = f"Loading {component}" if is_pipeline else "Loading checkpoint"
+        if len(weight_files) <= 1:
+            component_weights = {}
+            for wf in tqdm.tqdm(weight_files, desc=desc):
+                component_weights.update(self._load_file(wf))
+            return component_weights
+
+        workers = min(4, len(weight_files))
+
+        logger.info(f"Loading {len(weight_files)} {component} shard files with {workers} workers")
+        component_weights = {}
+        with ThreadPoolExecutor(max_workers=workers) as executor:
+            futures = {executor.submit(self._load_file, wf): wf for wf in weight_files}
+            for future in tqdm.tqdm(as_completed(futures), total=len(futures), desc=desc):
+                wf = futures[future]
+                try:
+                    loaded = future.result()
+                except Exception as exc:
+                    raise RuntimeError(f"Failed to load weight file {wf}") from exc
+                component_weights.update(loaded)
+
+        return component_weights
+
     def _find_weight_files(self, weight_dir) -> List[str]:
         """Find safetensors or bin weight files.
 

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

@@ -8,7 +8,7 @@
 import os
 import time
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
 import safetensors.torch
 import torch
@@ -17,6 +17,7 @@
 
 from tensorrt_llm._torch.utils import make_weak_ref
 from tensorrt_llm._torch.visual_gen.cache.teacache import CacheContext
+from tensorrt_llm._torch.visual_gen.checkpoints.prefetch import prefetch_files_to_host_cache
 from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner, CUDAGraphRunnerConfig
 from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput
 from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, ExtraParamSchema
@@ -151,6 +152,23 @@ def _assert_resolution(height: int, width: int, *, is_two_stage: bool = False) -
         )
 
 
+_LTX2_PREFETCHED_SAFETENSORS: Set[str] = set()
+
+
+def _prefetch_ltx2_safetensors_files(file_names: List[str]) -> bool:
+    """Warm LTX-2 safetensors files in the host page cache.
+
+    For distributed runs, local ranks split the file list and synchronize before
+    weight loading so ranks do not duplicate the prefetch work on the same node.
+    """
+    return prefetch_files_to_host_cache(
+        file_names,
+        description="LTX-2 checkpoint",
+        prefetched_paths=_LTX2_PREFETCHED_SAFETENSORS,
+        ignore_errors=True,
+    )
+
+
 def _load_ltx2_transformer_weights(
     checkpoint_dir: str,
     prefix: str,
@@ -178,6 +196,8 @@ def _load_ltx2_transformer_weights(
     if not sft_paths:
         raise ValueError(f"No safetensors files found in {checkpoint_dir}")
 
+    _prefetch_ltx2_safetensors_files(sft_paths)
+
     exclude_prefixes = tuple(exclude_prefixes) if exclude_prefixes else ()
 
     weights: Dict[str, torch.Tensor] = {}
@@ -860,6 +880,7 @@ def load_standard_components(
         # --- Resolve native config ----------------------------------------
         native_config = self.model_config.extra_attrs.get("monolithic_safetensors_config")
         sft_paths = _find_safetensors_files(checkpoint_dir)
+        _prefetch_ltx2_safetensors_files(sft_paths)
 
         if native_config is None and sft_paths:
             native_config = _read_safetensors_config(sft_paths[0])

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

@@ -35,7 +35,12 @@
 )
 from .ltx2_core.upsampler import LatentUpsamplerConfigurator, upsample_video
 from .ltx2_core.video_vae import TilingConfig
-from .pipeline_ltx2 import LTX2Pipeline, _assert_resolution, _find_safetensors_files
+from .pipeline_ltx2 import (
+    LTX2Pipeline,
+    _assert_resolution,
+    _find_safetensors_files,
+    _prefetch_ltx2_safetensors_files,
+)
 
 STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
 _FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2)
@@ -64,6 +69,7 @@ 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)
 
     raw: Dict[str, torch.Tensor] = {}
     alpha_dict: Dict[str, float] = {}
@@ -627,6 +633,7 @@ def load_standard_components(
             sft_paths = _find_safetensors_files(spatial_upsampler_path)
             if not sft_paths:
                 raise ValueError(f"No safetensors files found at {spatial_upsampler_path}")
+            _prefetch_ltx2_safetensors_files(sft_paths)
 
             config: Dict[str, Any] = {}
             try: