diff --git a/tensorrt_llm/_torch/visual_gen/checkpoints/prefetch.py b/tensorrt_llm/_torch/visual_gen/checkpoints/prefetch.py new file mode 100644 index 000000000000..7e4f7a34e5a0 --- /dev/null +++ b/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() diff --git a/tensorrt_llm/_torch/visual_gen/checkpoints/weight_loader.py b/tensorrt_llm/_torch/visual_gen/checkpoints/weight_loader.py index f1b22553ba66..824120840131 100644 --- a/tensorrt_llm/_torch/visual_gen/checkpoints/weight_loader.py +++ b/tensorrt_llm/_torch/visual_gen/checkpoints/weight_loader.py @@ -1,6 +1,7 @@ """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 @@ -8,6 +9,7 @@ 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. diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index 31cfd0a4bd49..deb0a48f2155 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/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]) 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 1747df003619..0832a5da1adc 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 @@ -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: