[None][test] Reuse MPI sessions and cache HF weights across grouped multi-GPU LLM tests

tensorrt_llm/_torch/distributed/communicator.py

@@ -1246,6 +1246,18 @@ def init_pp_comm(mapping):
     global _pp_comm
     if mpi_disabled():
         _pp_comm = PPCommTorch(mapping)
+    elif isinstance(_pp_comm, PPCommNCCL) and \
+            _pp_comm.mapping.world_size == mapping.world_size:
+        # Reuse the existing world NCCL communicator across LLM instances that
+        # share the same worker processes (e.g. a reused MpiPoolSession). The
+        # underlying comm depends only on (world_size, rank) -- it is a world
+        # communicator, independent of the pp/tp/ep layout -- so only the
+        # routing mapping needs refreshing. Recreating it would drop the old
+        # comm and trigger a collective ncclCommDestroy at an unsynchronized
+        # point during the next model build, which can deadlock on reused
+        # workers. Single-LLM (production) runs are unaffected: _pp_comm starts
+        # as None, so the first call still constructs a fresh PPCommNCCL.
+        _pp_comm.mapping = mapping
     else:
         _pp_comm = PPCommNCCL(mapping)
     init_helix_cp_comm(mapping)

tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py

@@ -15,6 +15,8 @@
 import glob
 import multiprocessing
 import os
+import threading
+from collections import OrderedDict
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, List
 
@@ -33,6 +35,12 @@
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
 
+_WEIGHT_CACHE_ENV = "TRTLLM_HF_WEIGHT_CACHE"
+_WEIGHT_CACHE_MAX_ENTRIES_ENV = "TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES"
+_DEFAULT_WEIGHT_CACHE_MAX_ENTRIES = 2
+_WEIGHT_CACHE_LOCK = threading.Lock()
+_WEIGHT_CACHE: OrderedDict[tuple, dict[str, Any]] = OrderedDict()
+
 
 @register_checkpoint_weight_loader("MX")
 @register_checkpoint_weight_loader("mistral")
@@ -43,6 +51,85 @@ class HfWeightLoader(BaseWeightLoader):
     Loads weights from SafeTensors/bin/pth files.
     """
 
+    @staticmethod
+    def _is_weight_cache_enabled() -> bool:
+        return os.environ.get(_WEIGHT_CACHE_ENV,
+                              "0").lower() in ("1", "true", "yes", "on")
+
+    @staticmethod
+    def _weight_cache_max_entries() -> int:
+        try:
+            return max(
+                0,
+                int(
+                    os.environ.get(_WEIGHT_CACHE_MAX_ENTRIES_ENV,
+                                   _DEFAULT_WEIGHT_CACHE_MAX_ENTRIES)))
+        except ValueError:
+            logger.warning(
+                f"Invalid {_WEIGHT_CACHE_MAX_ENTRIES_ENV} value; disabling HF weight cache."
+            )
+            return 0
+
+    @staticmethod
+    def _weight_files_cache_key(weight_files: List[str],
+                                use_consolidated: bool) -> tuple:
+        file_fingerprint = []
+        for file_name in sorted(weight_files):
+            stat = os.stat(file_name)
+            file_fingerprint.append(
+                (os.path.abspath(file_name), stat.st_size, stat.st_mtime_ns))
+        return (tuple(file_fingerprint), use_consolidated)
+
+    @classmethod
+    def _clear_weight_cache(cls) -> None:
+        with _WEIGHT_CACHE_LOCK:
+            _WEIGHT_CACHE.clear()
+
+    @staticmethod
+    def _evict_to_make_room() -> None:
+        """Evict LRU entries on a miss BEFORE the new load, so CPU never holds
+        the old (cached) and new (loading) weights at once (a ~2x peak)."""
+        max_entries = HfWeightLoader._weight_cache_max_entries()
+        if max_entries <= 0:
+            return
+        with _WEIGHT_CACHE_LOCK:
+            while len(_WEIGHT_CACHE) >= max_entries:
+                _WEIGHT_CACHE.popitem(last=False)
+
+    @staticmethod
+    def _wrap_cached_weights(weights: dict[str, Any]) -> ConsumableWeightsDict:
+        # Return a fresh dict wrapper because model loaders call mark_consumed().
+        # Tensor values are intentionally shared: this cache targets read-only
+        # raw checkpoint tensors, not per-config materialized module weights.
+        return ConsumableWeightsDict(dict(weights))
+
+    @staticmethod
+    def _cache_loaded_weights(cache_key: tuple,
+                              loaded_weights: dict[str, Any]) -> None:
+        max_entries = HfWeightLoader._weight_cache_max_entries()
+        if max_entries <= 0:
+            return
+
+        if isinstance(loaded_weights, ConsumableWeightsDict):
+            cache_weights = dict(loaded_weights.items())
+        else:
+            cache_weights = dict(loaded_weights)
+
+        with _WEIGHT_CACHE_LOCK:
+            _WEIGHT_CACHE[cache_key] = cache_weights
+            _WEIGHT_CACHE.move_to_end(cache_key)
+            while len(_WEIGHT_CACHE) > max_entries:
+                _WEIGHT_CACHE.popitem(last=False)
+
+    @staticmethod
+    def _get_cached_weights(cache_key: tuple) -> ConsumableWeightsDict | None:
+        with _WEIGHT_CACHE_LOCK:
+            weights = _WEIGHT_CACHE.get(cache_key)
+            if weights is None:
+                return None
+            _WEIGHT_CACHE.move_to_end(cache_key)
+            return HfWeightLoader._wrap_cached_weights(weights)
+
     @staticmethod
     def _get_local_available_host_memory() -> int:
         """Determine the minimum available memory observed on the local node
@@ -77,6 +164,28 @@ def load_weights(self,
         if len(filtered_weight_files) > 0:
             weight_files = filtered_weight_files
         if weight_files:
+            # Only fingerprint the files (per-file os.stat) when the cache is
+            # enabled; the default path must not pay for a key it never uses.
+            cache_key = self._weight_files_cache_key(
+                weight_files,
+                use_consolidated) if self._is_weight_cache_enabled() else None
+            if cache_key is not None:
+                cached_weights = self._get_cached_weights(cache_key)
+                if cached_weights is not None:
+                    logger.info(
+                        f"Reusing cached HF checkpoint weights from {checkpoint_dir}."
+                    )
+                    # The safetensors path always has a local barrier before
+                    # deserialization. Cache hits must participate as well:
+                    # rank-local caches can diverge after eviction, and a hit
+                    # on one rank must not skip a barrier that a miss on
+                    # another rank is about to enter.
+                    local_mpi_barrier()
+                    return cached_weights
+                # Cache miss: evict now so the upcoming load doesn't transiently
+                # hold the old (cached) and new (loading) weights together.
+                self._evict_to_make_room()
+
             # Prefetch the weight files to CPU memory if the size is less than 90% of the available memory.
             # This is a heuristic to avoid prefetching files that are too large and causing file cache thrashing.
             prefetch_size = sum(os.path.getsize(file) for file in weight_files)
@@ -98,18 +207,38 @@ def load_weights(self,
             # skipped. Ranks that didn't prefetch reach the barrier immediately.
             local_mpi_barrier()
 
-            return self._load_weights_in_parallel(
+            weights = self._load_weights_in_parallel(
                 weight_files, self._load_safetensors_file,
                 "Loading safetensors weights in parallel")
+            if cache_key is not None:
+                self._cache_loaded_weights(cache_key, weights)
+            return weights
 
         weight_files = glob.glob(f"{checkpoint_dir}/*.bin")
         if not weight_files:
             weight_files = glob.glob(f"{checkpoint_dir}/*.pth")
 
         if weight_files:
-            return self._load_weights_in_parallel(
+            cache_key = self._weight_files_cache_key(
+                weight_files,
+                use_consolidated) if self._is_weight_cache_enabled() else None
+            if cache_key is not None:
+                cached_weights = self._get_cached_weights(cache_key)
+                if cached_weights is not None:
+                    logger.info(
+                        f"Reusing cached HF checkpoint weights from {checkpoint_dir}."
+                    )
+                    return cached_weights
+                # Cache miss: evict now so the upcoming load doesn't transiently
+                # hold the old (cached) and new (loading) weights together.
+                self._evict_to_make_room()
+
+            weights = self._load_weights_in_parallel(
                 weight_files, self._load_bin_or_path_file,
                 "Loading bin weights in parallel")
+            if cache_key is not None:
+                self._cache_loaded_weights(cache_key, weights)
+            return weights
 
         raise RuntimeError(f"No weight files found in {checkpoint_dir}.")
 

tensorrt_llm/executor/executor.py

@@ -639,7 +639,7 @@ def create(
             return GenerationExecutor._create_ipc_executor(
                 worker_kwargs,
                 model_world_size=model_world_size,
-                mpi_session=None,  # use mpi4py
+                mpi_session=mpi_session,
                 postproc_worker_config=postproc_worker_config,
                 is_llm_executor=is_llm_executor,
                 use_worker=False)

tensorrt_llm/executor/proxy.py

@@ -98,6 +98,7 @@ def __init__(
         self.worker_cls = worker_cls
 
         mpi_process_pre_spawned: bool = get_spawn_proxy_process_env()
+        self._owns_mpi_session = mpi_session is None
 
         if mpi_session is None:
             if mpi_process_pre_spawned:
@@ -527,7 +528,8 @@ def shutdown(self):
             self._resource_governor_queue.close()
 
         self.workers_started = False
-        self.mpi_session.shutdown()
+        if self._owns_mpi_session:
+            self.mpi_session.shutdown()
 
         # Process the errors in-case error during shutting down the threads
         self._handle_background_error()

tensorrt_llm/executor/rpc_proxy.py

@@ -64,6 +64,7 @@ def __init__(
         )
 
         self.model_world_size = model_world_size
+        self._owns_mpi_session = mpi_session is None
         self._create_mpi_session(model_world_size, mpi_session)
 
         # Inject the generated HMAC key into worker_kwargs for workers
@@ -275,7 +276,7 @@ def shutdown(self):
 
         # 3. shutdown the mpi session, this should wait until all the PyExecutor
         # processes are shutdown
-        if self.mpi_session is not None:
+        if self.mpi_session is not None and self._owns_mpi_session:
             logger_debug(f"Shutting down mpi session", color="yellow")
             self.mpi_session.shutdown()
             logger_debug(f"Mpi session shutdown", color="yellow")

tensorrt_llm/llmapi/_grouped_test_utils.py

@@ -0,0 +1,132 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Shared helpers for grouped multi-GPU tests that reuse one ``MpiPoolSession``
+across many ``LLM`` instances.
+
+This lives in-package (rather than under ``tests/``) because the unittest and
+integration test roots have separate ``sys.path`` bases and cannot import a
+single shared module otherwise. It is test-infrastructure only and must not be
+used on production inference paths. Heavy imports are done lazily so importing
+this module never triggers a circular import during package initialization.
+"""
+
+import os
+from contextlib import contextmanager
+from typing import Optional
+
+
+def restore_env_var(name: str, value: Optional[str]) -> None:
+    """Restore an env var to a previously captured value (or remove it)."""
+    if value is None:
+        os.environ.pop(name, None)
+    else:
+        os.environ[name] = value
+
+
+@contextmanager
+def hf_weight_cache_env(max_entries: str = "1"):
+    """Enable the HF weight cache via env vars, restoring prior values on exit.
+
+    IMPORTANT (load-bearing ordering): enter this BEFORE spawning the shared
+    ``MpiPoolSession``. ``MpiPoolSession`` snapshots ``TRTLLM*``/``TLLM*`` env at
+    spawn time and passes that copy to the workers (which are the processes that
+    actually load weights); enabling the cache AFTER the pool spawns leaves it
+    silently disabled in the workers. Back a module-scoped ``hf_weight_cache``
+    fixture with this and have the session fixture depend on that fixture so the
+    env is exported first.
+    """
+    prev = os.environ.get("TRTLLM_HF_WEIGHT_CACHE")
+    prev_entries = os.environ.get("TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES")
+    os.environ["TRTLLM_HF_WEIGHT_CACHE"] = "1"
+    os.environ["TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES"] = max_entries
+    try:
+        yield
+    finally:
+        restore_env_var("TRTLLM_HF_WEIGHT_CACHE", prev)
+        restore_env_var("TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES", prev_entries)
+
+
+def shared_mpi_session(n_workers: int):
+    """Generator yielding a shared ``MpiPoolSession`` (or ``None`` when MPI is
+    disabled), shutting it down on teardown. Intended to back a module-scoped
+    pytest fixture so a pool of workers is reused across many LLMs."""
+    from tensorrt_llm._utils import mpi_disabled
+
+    if mpi_disabled():
+        yield None
+        return
+
+    from tensorrt_llm.llmapi.mpi_session import MpiPoolSession
+
+    mpi_session = MpiPoolSession(n_workers=n_workers)
+    try:
+        yield mpi_session
+    finally:
+        mpi_session.shutdown()
+
+
+def mpi_session_kwargs(mpi_session) -> dict:
+    """LLM kwargs that inject a shared MPI session, or ``{}`` when there is
+    none (e.g. MPI disabled). Lets harness-based tests forward the session."""
+    return {"_mpi_session": mpi_session} if mpi_session is not None else {}
+
+
+def reset_worker_torch_compile_state() -> None:
+    """Reset per-worker torch.compile / Dynamo state (runs inside each worker).
+
+    Dynamo's recompile counter is process-global and per-code-object. When
+    worker processes are reused across LLMs (shared ``MpiPoolSession``), each
+    ``torch_compile`` case recompiles the same ``model.forward`` code object
+    under new guards; the count accumulates and eventually trips
+    ``recompile_limit`` (16), which is a HARD failure under ``fullgraph=True``
+    (``FailOnRecompileLimitHit``) and aborts the whole MPI job. Resetting
+    between cases makes each LLM start from a clean compile cache, like a fresh
+    process. Submit this to each worker via ``mpi_session.submit_sync(...)``.
+    """
+    import torch
+
+    torch._dynamo.reset()
+
+
+def reset_shared_session_torch_compile_state(make_llm) -> None:
+    """Reset per-worker torch.compile state on the shared session behind a
+    ``make_shared_llm`` factory (no-op if there is no shared session).
+
+    Call from a grouped test's per-case teardown so recompile counts don't
+    accumulate across cases on reused workers (see
+    ``reset_worker_torch_compile_state`` for the failure it prevents).
+    """
+    mpi_session = getattr(make_llm, "mpi_session", None)
+    if mpi_session is not None:
+        mpi_session.submit_sync(reset_worker_torch_compile_state)
+
+
+def clear_worker_weight_cache() -> None:
+    """Drop the per-worker HF raw-weight cache (runs inside each worker).
+
+    The cache is a process-global keyed by checkpoint file fingerprints, so it
+    otherwise lives until the worker process exits. Submit this to each worker
+    via ``mpi_session.submit_sync(...)`` on group teardown to invalidate it
+    explicitly instead of relying on process death.
+    """
+    from tensorrt_llm._torch.models.checkpoints import HfWeightLoader
+
+    HfWeightLoader._clear_weight_cache()
+
+
+def make_shared_llm(mpi_session):
+    """Return an ``LLM`` factory that transparently injects a shared MPI session.
+
+    Tests build the LLM by calling the factory exactly like ``LLM(...)``; the
+    shared session (if any) is passed through without the test having to know it
+    exists. Falls back to a private per-LLM session when ``mpi_session`` is None.
+    The factory exposes ``.mpi_session`` so callers can reset per-worker compile
+    state between cases without threading the session separately.
+    """
+    from tensorrt_llm import LLM
+
+    def shared_llm(*args, **kwargs):
+        return LLM(*args, **kwargs, **mpi_session_kwargs(mpi_session))
+
+    shared_llm.mpi_session = mpi_session
+    return shared_llm