diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 6fbbf16a19df..289795c63820 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/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) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index 264104c6f2e1..b135fa568903 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/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}.") diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index fdd3dff44a48..d11b12393dc6 100644 --- a/tensorrt_llm/executor/executor.py +++ b/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) diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index ffe793bc5b48..027aa7cae84c 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/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() diff --git a/tensorrt_llm/executor/rpc_proxy.py b/tensorrt_llm/executor/rpc_proxy.py index 7356a9d86d5f..d72259c3e02e 100644 --- a/tensorrt_llm/executor/rpc_proxy.py +++ b/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") diff --git a/tensorrt_llm/llmapi/_grouped_test_utils.py b/tensorrt_llm/llmapi/_grouped_test_utils.py new file mode 100644 index 000000000000..3bc12a22def8 --- /dev/null +++ b/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 diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index b637c1738dc8..ab5d7c446fda 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -313,6 +313,10 @@ def __init__(self, logger_debug(f"LLM.args.mpi_session: {self.args.mpi_session}\n", "yellow") self.mpi_session = self.args.mpi_session + # Keep the live session on LLM only. LLM args are passed to model-build + # tasks and executor workers, and MpiSession objects are not pickleable. + self.args.mpi_session = None + self._owns_mpi_session = self.mpi_session is None # Build this LLM's post-processing hook for the in-proxy detok path (each # postproc worker builds its own). Resolving here fails fast on a bad @@ -338,10 +342,12 @@ def __init__(self, logger_debug("LLM create MpiPoolSession\n", "yellow") self.mpi_session = MpiPoolSession( n_workers=self.args.parallel_config.world_size) + self._owns_mpi_session = True else: logger_debug("LLM create MpiCommSession\n", "yellow") self.mpi_session = create_mpi_comm_session( self.args.parallel_config.world_size) + self._owns_mpi_session = True try: # Due to the Executor can only accept a engine path, we need to save the engine to a directory @@ -364,7 +370,8 @@ def __init__(self, self._build_model() except Exception: - if self.mpi_session is not None: + if (self.mpi_session is not None + and getattr(self, "_owns_mpi_session", True)): self.mpi_session.shutdown() raise @@ -1472,7 +1479,8 @@ def shutdown(self) -> None: self._encoder_executor.shutdown() self._encoder_executor = None - if hasattr(self, 'mpi_session') and self.mpi_session is not None: + if (hasattr(self, 'mpi_session') and self.mpi_session is not None + and getattr(self, "_owns_mpi_session", True)): self.mpi_session.shutdown() self.mpi_session = None diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 641ed8c11ca4..360ab3b56c5a 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -74,6 +74,53 @@ def patched_start_mpi_pool(self): patched_start_mpi_pool) +from tensorrt_llm.llmapi._grouped_test_utils import \ + clear_worker_weight_cache as _clear_worker_weight_cache # noqa: E402 +from tensorrt_llm.llmapi._grouped_test_utils import \ + hf_weight_cache_env as _hf_weight_cache_env +from tensorrt_llm.llmapi._grouped_test_utils import \ + make_shared_llm as _make_shared_llm +from tensorrt_llm.llmapi._grouped_test_utils import \ + reset_shared_session_torch_compile_state as \ + _reset_shared_session_torch_compile_state +from tensorrt_llm.llmapi._grouped_test_utils import \ + shared_mpi_session as _shared_mpi_session + + +@pytest.fixture(scope="module") +def hf_weight_cache(): + with _hf_weight_cache_env(): + yield + + +@pytest.fixture(scope="module") +def shared_mpi_session_4gpu(hf_weight_cache): + # Ordering is load-bearing: MpiPoolSession snapshots TRTLLM*/TLLM* env at + # spawn time and passes that copy to the workers. Any env exported AFTER the + # pool spawns never reaches them. Depending on hf_weight_cache guarantees + # TRTLLM_HF_WEIGHT_CACHE is exported before MpiPoolSession spawns, otherwise + # the weight cache would be silently disabled in the workers (which are the + # processes that actually load weights) and only session reuse would speed + # things up. The assertion pins this contract so a future fixture reorder + # fails loudly instead of silently regressing the optimization. + assert os.environ.get("TRTLLM_HF_WEIGHT_CACHE") == "1", ( + "hf_weight_cache must be set up before the shared MPI pool spawns so " + "the workers inherit TRTLLM_HF_WEIGHT_CACHE at spawn time.") + yield from _shared_mpi_session(4) + + +@pytest.fixture(scope="module") +def shared_llm_4gpu(shared_mpi_session_4gpu): + try: + yield _make_shared_llm(shared_mpi_session_4gpu) + finally: + # Explicitly invalidate the per-worker HF weight cache while the shared + # session is still alive (this fixture tears down before + # shared_mpi_session_4gpu), instead of relying on worker process exit. + if shared_mpi_session_4gpu is not None: + shared_mpi_session_4gpu.submit_sync(_clear_worker_weight_cache) + + def _get_default_torch_compile_config(torch_compile): return TorchCompileConfig(enable_fullgraph=True, enable_piecewise_cuda_graph=True, @@ -1337,6 +1384,64 @@ def test_auto_dtype_vswa_chunked_prefill_reuse(self): class TestDeepSeekV3Lite(LlmapiAccuracyTestHarness): MODEL_NAME = "deepseek-ai/DeepSeek-V3-Lite" MODEL_PATH = f"{llm_models_root()}/DeepSeek-V3-Lite/bf16" + # (id, moe_backend, mtp_nextn, tp_size, pp_size, ep_size, torch_compile, + # low_precision_combine). fp8kv / attention_dp / cuda_graph / + # overlap_scheduler are True for every case (see _NVFP4_4GPU_PREMERGE_CASES). + _NVFP4_4GPU_PREMERGE_MATRIX = ( + ("cutlass_mtp0_tp4_compile_off", "CUTLASS", 0, 4, 1, 1, False, False), + ("cutlass_mtp0_tp4_compile_on", "CUTLASS", 0, 4, 1, 1, True, False), + ("cutlass_mtp0_ep4_compile_on", "CUTLASS", 0, 4, 1, 4, True, False), + ("cutlass_mtp0_tp2pp2_compile_off", "CUTLASS", 0, 2, 2, 1, False, + False), + ("cutlass_mtp0_tp2pp2_compile_on", "CUTLASS", 0, 2, 2, 1, True, False), + ("cutlass_mtp2_tp4_compile_off", "CUTLASS", 2, 4, 1, 1, False, False), + ("trtllm_mtp2_ep4_compile_off", "TRTLLM", 2, 4, 1, 4, False, False), + ("cutlass_mtp2_pp4_compile_off", "CUTLASS", 2, 1, 4, 1, False, False), + ("cutlass_mtp0_pp4_compile_off", "CUTLASS", 0, 1, 4, 1, False, False), + ("cutlass_mtp0_pp4_compile_on", "CUTLASS", 0, 1, 4, 1, True, False), + ("trtllm_mtp0_tp4_compile_off", "TRTLLM", 0, 4, 1, 1, False, False), + ("trtllm_mtp0_ep4_compile_off", "TRTLLM", 0, 4, 1, 4, False, False), + ("cutlass_mtp0_tp4_lpc_compile_off", "CUTLASS", 0, 4, 1, 1, False, + True), + # CUTEDSL coverage (skipped off SM100/103 by _run_nvfp4_4gpus_case); + # preserves the pre-merge CUTEDSL cases upstream carried individually. + ("cutedsl_mtp0_tp2pp2_compile_off", "CUTEDSL", 0, 2, 2, 1, False, False + ), + ("cutedsl_mtp2_tp4_compile_off", "CUTEDSL", 2, 4, 1, 1, False, False), + ) + # Each case is pure _run_nvfp4_4gpus_case kwargs; the id (row[0]) stays in + # the matrix and is used only for the parametrize ids. + _NVFP4_4GPU_PREMERGE_CASES = tuple( + dict(moe_backend=row[1], + mtp_nextn=row[2], + tp_size=row[3], + pp_size=row[4], + ep_size=row[5], + torch_compile=row[6], + low_precision_combine=row[7], + fp8kv=True, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True) for row in _NVFP4_4GPU_PREMERGE_MATRIX) + + # bf16 pre-merge cases share the SAME 4-GPU MPI session as the nvfp4 group + # (both are 4-worker DeepSeek-V3-Lite, so world_size matches). mtp_nextn / + # attention_dp / cuda_graph / overlap_scheduler are True for every case. + # (id, tp_size, pp_size, ep_size, torch_compile). + _BF16_4GPU_PREMERGE_MATRIX = ( + ("tp4_compile_on", 4, 1, 1, True), + ("ep4_compile_off", 4, 1, 4, False), + ("tp2pp2_compile_off", 2, 2, 1, False), + ) + _BF16_4GPU_PREMERGE_CASES = tuple( + dict(tp_size=row[1], + pp_size=row[2], + ep_size=row[3], + torch_compile=row[4], + mtp_nextn=2, + attention_dp=True, + cuda_graph=True, + overlap_scheduler=True) for row in _BF16_4GPU_PREMERGE_MATRIX) @pytest.mark.skip_less_device_memory(60000) @parametrize_with_ids("v2_kv_cache", [True, False]) @@ -1515,6 +1620,26 @@ def test_bfloat16_4gpus_kv_cache_aware_routing(self, mtp_nextn): def test_bfloat16_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, attention_dp, cuda_graph, overlap_scheduler, torch_compile): + self._run_bfloat16_4gpus_case(tp_size=tp_size, + pp_size=pp_size, + ep_size=ep_size, + mtp_nextn=mtp_nextn, + attention_dp=attention_dp, + cuda_graph=cuda_graph, + overlap_scheduler=overlap_scheduler, + torch_compile=torch_compile) + + def _run_bfloat16_4gpus_case(self, + *, + tp_size, + pp_size, + ep_size, + mtp_nextn, + attention_dp, + cuda_graph, + overlap_scheduler, + torch_compile, + make_llm=LLM): if pp_size > 1 and mtp_nextn > 0: num_hidden_layers = 30 pp_partition = [num_hidden_layers // pp_size + 1] * pp_size @@ -1533,17 +1658,23 @@ def test_bfloat16_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, mtp_config = None if mtp_nextn > 0: mtp_config = MTPDecodingConfig(max_draft_len=mtp_nextn) - with LLM(self.MODEL_PATH, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - pp_partition=pp_partition, - moe_expert_parallel_size=ep_size, - kv_cache_config=kv_cache_config, - **pytorch_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config) as llm: - task = GSM8K(self.MODEL_NAME) - task.evaluate(llm) + try: + with make_llm(self.MODEL_PATH, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + pp_partition=pp_partition, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + **pytorch_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config) as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + finally: + # On a reused MPI session, reset each worker's torch.compile/Dynamo + # state so recompile counts don't accumulate across cases and trip + # the fullgraph recompile limit (FailOnRecompileLimitHit). + _reset_shared_session_torch_compile_state(make_llm) @pytest.mark.skip_less_device(4) @parametrize_with_ids("mtp_nextn", @@ -2207,6 +2338,67 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, overlap_scheduler, low_precision_combine, tp_size, pp_size, ep_size, torch_compile, mtp_nextn, moe_backend): + self._run_nvfp4_4gpus_case( + fp8kv=fp8kv, + attention_dp=attention_dp, + cuda_graph=cuda_graph, + overlap_scheduler=overlap_scheduler, + low_precision_combine=low_precision_combine, + tp_size=tp_size, + pp_size=pp_size, + ep_size=ep_size, + torch_compile=torch_compile, + mtp_nextn=mtp_nextn, + moe_backend=moe_backend, + ) + + @pytest.mark.skip_less_device(4) + @skip_pre_blackwell + # The shared module-scoped MPI session is intentionally kept alive across the + # parametrized cases (that is the whole point of reuse), so mpi4py's + # `_manager_spawn` management thread outlives each test. Disable the + # threadleak check here, consistent with other session/proxy-based tests. + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize( + "case", + _NVFP4_4GPU_PREMERGE_CASES, + ids=[row[0] for row in _NVFP4_4GPU_PREMERGE_MATRIX], + ) + def test_nvfp4_4gpus_premerge_grouped(self, case, shared_llm_4gpu): + # shared_llm_4gpu injects the shared 4-GPU MPI session and (transitively) + # depends on hf_weight_cache, so the weight-cache env is exported before + # the pool spawns (see the shared_mpi_session_4gpu fixture for details). + self._run_nvfp4_4gpus_case(make_llm=shared_llm_4gpu, **case) + + @pytest.mark.skip_less_device(4) + @skip_pre_hopper + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize( + "case", + _BF16_4GPU_PREMERGE_CASES, + ids=[row[0] for row in _BF16_4GPU_PREMERGE_MATRIX], + ) + def test_bfloat16_4gpus_premerge_grouped(self, case, shared_llm_4gpu): + # Reuses the same module-scoped shared_llm_4gpu session as the nvfp4 + # group. The HF weight cache is keyed by checkpoint file fingerprint, so + # the bf16 model gets its own entry; with max_entries=1 the LRU evicts + # the other model on the first miss, bounding CPU cache to one model. + self._run_bfloat16_4gpus_case(make_llm=shared_llm_4gpu, **case) + + def _run_nvfp4_4gpus_case(self, + *, + fp8kv, + attention_dp, + cuda_graph, + overlap_scheduler, + low_precision_combine, + tp_size, + pp_size, + ep_size, + torch_compile, + mtp_nextn, + moe_backend, + make_llm=LLM): sm_version = get_sm_version() if moe_backend == "TRTLLM" and sm_version in (120, 121): pytest.skip(f"{moe_backend} backend does not support SM 120 or 121") @@ -2233,20 +2425,26 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, if fp8kv: kv_cache_config.dtype = "fp8" - with LLM( - f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - kv_cache_config=kv_cache_config, - **pytorch_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config, - ) as llm: - assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 + try: + with make_llm( + f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + **pytorch_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - task = GSM8K(self.MODEL_NAME) - task.evaluate(llm) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + finally: + # On a reused MPI session, reset each worker's torch.compile/Dynamo + # state so recompile counts don't accumulate across cases and trip + # the fullgraph recompile limit (FailOnRecompileLimitHit). + _reset_shared_session_torch_compile_state(make_llm) @parametrize_with_ids( "fp8kv,attention_dp,cuda_graph,overlap_scheduler", diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index d65e8ba0cd70..bd9b323dfb7a 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -15,7 +15,8 @@ l0_dgx_h100: auto_trigger: others orchestrator: mpi tests: - - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2" + - unittest/llmapi/test_llm_multi_gpu_pytorch.py::TestLlmMultiGpu2gpu::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 + - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2" -k "not test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1" - unittest/llmapi/test_additional_model_outputs.py -m "gpu2" - unittest/_torch/multi_gpu -m "not post_merge" TIMEOUT (90) - unittest/_torch/modeling/test_modeling_pixtral.py::test_tensor_parallelism diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index f12e16220744..bdd0040a1b9d 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -20,15 +20,27 @@ l0_gb200_multi_gpus: - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - # ---- curated pre-merge smoke set: diversify moe_backend across CUTLASS/TRTLLM/CUTEDSL ---- - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=0-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=0-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=2-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=2-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=True-torch_compile=False] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_premerge_grouped[tp4_compile_on] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_premerge_grouped[ep4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_premerge_grouped[tp2pp2_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_kv_cache_aware_routing[mtp_nextn=0] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_kv_cache_aware_routing[mtp_nextn=2] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_tp4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_tp4_compile_on] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_ep4_compile_on] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_tp2pp2_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_tp2pp2_compile_on] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp2_tp4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[trtllm_mtp2_ep4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp2_pp4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_pp4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_pp4_compile_on] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[trtllm_mtp0_tp4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[trtllm_mtp0_ep4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutlass_mtp0_tp4_lpc_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutedsl_mtp0_tp2pp2_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_premerge_grouped[cutedsl_mtp2_tp4_compile_off] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=True-moe_backend=WIDEEP] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=True-moe_backend=TRTLLM] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_online_eplb[moe_backend=TRTLLM] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_block_reuse[TEP4] diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index db0dc31b5223..d0c1f2a6fc74 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -488,7 +488,7 @@ unittest/executor/test_rpc.py::TestRpcCorrectness::test_incremental_task_async S unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741) unittest/executor/test_rpc_worker.py SKIP (https://nvbugs/5605741) unittest/llmapi/test_llm_multi_gpu.py -m "gpu4 and part0" SKIP (https://nvbugs/5348958) -unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 SKIP (https://nvbugs/6109745) +unittest/llmapi/test_llm_multi_gpu_pytorch.py::TestLlmMultiGpu2gpu::test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1 SKIP (https://nvbugs/6109745) unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[None] SKIP (https://nvbugs/6162504) unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[cuda_graph_config0] SKIP (https://nvbugs/6162504) unittest/llmapi/test_memory_profiling.py::test_profile_kvcache SKIP (https://nvbugs/5580781) diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py index fbe8a3dbf481..8ec25c3a4b55 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py @@ -3,6 +3,7 @@ import pytest from tensorrt_llm._torch.models.checkpoints import HfWeightLoader +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict from tensorrt_llm.mapping import Mapping @@ -97,3 +98,111 @@ def test_load_weights_ignores_consolidated_ckpt_when_sharded_ckpt_exists( load_weights_in_parallel.assert_called_once() loaded_weight_files = load_weights_in_parallel.call_args[0][0] assert set(loaded_weight_files) == expected_safetensor_filenames + + +def test_weight_cache_reuses_raw_weights_with_fresh_consumable_wrapper(tmp_path, monkeypatch): + monkeypatch.setenv("TRTLLM_HF_WEIGHT_CACHE", "1") + HfWeightLoader._clear_weight_cache() + + checkpoint_dir = tmp_path / "foo" + checkpoint_dir.mkdir() + (checkpoint_dir / "model.safetensors").touch() + + raw_weight = object() + loader = HfWeightLoader() + + try: + with ( + mock.patch.object( + loader, + "_load_weights_in_parallel", + return_value=ConsumableWeightsDict({"foo.weight": raw_weight}), + ) as load_weights_in_parallel, + mock.patch.object(loader, "prefetch_files"), + ): + first = loader.load_weights(str(checkpoint_dir), mapping=Mapping()) + assert first["foo.weight"] is raw_weight + assert first.mark_consumed("foo") == 1 + assert len(first) == 0 + + second = loader.load_weights(str(checkpoint_dir), mapping=Mapping()) + + load_weights_in_parallel.assert_called_once() + assert second["foo.weight"] is raw_weight + finally: + HfWeightLoader._clear_weight_cache() + + +def test_weight_cache_evicts_before_load_on_miss(tmp_path, monkeypatch): + # On a cross-model miss with a full cache (max_entries=1), the old entry + # must be evicted BEFORE the new weights load, so CPU never holds both the + # old (cached) and new (loading) weights at once (no transient 2x peak). + from tensorrt_llm._torch.models.checkpoints.hf import weight_loader as wl + + monkeypatch.setenv("TRTLLM_HF_WEIGHT_CACHE", "1") + monkeypatch.setenv("TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES", "1") + HfWeightLoader._clear_weight_cache() + + dir_a = tmp_path / "a" + dir_a.mkdir() + (dir_a / "model.safetensors").touch() + dir_b = tmp_path / "b" + dir_b.mkdir() + (dir_b / "model.safetensors").touch() + + loader = HfWeightLoader() + try: + with mock.patch.object(loader, "prefetch_files"): + with mock.patch.object( + loader, + "_load_weights_in_parallel", + return_value=ConsumableWeightsDict({"foo.weight": object()}), + ): + loader.load_weights(str(dir_a), mapping=Mapping()) + assert len(wl._WEIGHT_CACHE) == 1 + + def assert_room_freed_before_load(*args, **kwargs): + # The old (A) entry must already be gone by the time B loads. + assert len(wl._WEIGHT_CACHE) == 0 + return ConsumableWeightsDict({"foo.weight": object()}) + + with mock.patch.object( + loader, + "_load_weights_in_parallel", + side_effect=assert_room_freed_before_load, + ): + loader.load_weights(str(dir_b), mapping=Mapping()) + + assert len(wl._WEIGHT_CACHE) == 1 + finally: + HfWeightLoader._clear_weight_cache() + + +def test_weight_cache_disabled_by_default(tmp_path, monkeypatch): + monkeypatch.delenv("TRTLLM_HF_WEIGHT_CACHE", raising=False) + HfWeightLoader._clear_weight_cache() + + checkpoint_dir = tmp_path / "foo" + checkpoint_dir.mkdir() + (checkpoint_dir / "model.safetensors").touch() + + loader = HfWeightLoader() + + try: + with ( + mock.patch.object( + loader, + "_load_weights_in_parallel", + side_effect=[ + ConsumableWeightsDict({"foo.weight": object()}), + ConsumableWeightsDict({"foo.weight": object()}), + ], + ) as load_weights_in_parallel, + mock.patch.object(loader, "prefetch_files"), + ): + loader.load_weights(str(checkpoint_dir), mapping=Mapping()) + loader.load_weights(str(checkpoint_dir), mapping=Mapping()) + + assert load_weights_in_parallel.call_count == 2 + finally: + HfWeightLoader._clear_weight_cache() diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index fb174a99152d..1579261922a2 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -1827,7 +1827,8 @@ def llm_return_logprobs_test_harness(prompt_logprobs: Optional[int], return_generation_logits: bool, tp_size=1, streaming=False, - backend=None): + backend=None, + **extra_llm_kwargs): LLM_CLASS = LLM llm_args_extra = {} kv_cache_args_extra = {} @@ -1849,6 +1850,7 @@ def llm_return_logprobs_test_harness(prompt_logprobs: Optional[int], tensor_parallel_size=tp_size, gather_generation_logits=True, **llm_args_extra, + **extra_llm_kwargs, ) prompts = ["A B C D E F G H I J K"] @@ -1920,6 +1922,7 @@ async def main(): await asyncio.gather(*tasks) asyncio.run(main()) + llm.shutdown() @force_ampere @@ -2205,7 +2208,8 @@ def llm_get_stats_test_harness(tp_size: int = 1, pytorch_backend: bool = False, use_overlap: bool = False, enable_chunked_prefill: bool = False, - enable_iter_req_stats: bool = False): + enable_iter_req_stats: bool = False, + **extra_llm_kwargs): if return_context_logits and pytorch_backend: pytest.skip("pytorch backend does not support context logits") @@ -2252,7 +2256,8 @@ def llm_get_stats_test_harness(tp_size: int = 1, kv_cache_config=global_kvcache_config, tensor_parallel_size=tp_size, pipeline_parallel_size=pp_size, - **llm_args_extra) as llm: + **llm_args_extra, + **extra_llm_kwargs) as llm: max_tokens = 5 sampling_params = SamplingParams(max_tokens=max_tokens, @@ -2354,7 +2359,8 @@ def llm_get_stats_async_test_harness(tp_size: int = 1, pytorch_backend: bool = False, use_overlap: bool = False, enable_chunked_prefill: bool = False, - enable_iter_req_stats: bool = False): + enable_iter_req_stats: bool = False, + **extra_llm_kwargs): if return_context_logits and pytorch_backend: pytest.skip("pytorch backend does not support context logits") @@ -2395,7 +2401,8 @@ def llm_get_stats_async_test_harness(tp_size: int = 1, kv_cache_config=global_kvcache_config, tensor_parallel_size=tp_size, pipeline_parallel_size=pp_size, - **llm_args_extra) as llm: + **llm_args_extra, + **extra_llm_kwargs) as llm: max_tokens = 6 sampling_params = SamplingParams(max_tokens=max_tokens, @@ -2491,7 +2498,9 @@ def success_path(): success_path() -def _test_llm_capture_request_error(pytorch_backend: bool, tp_size: int = 1): +def _test_llm_capture_request_error(pytorch_backend: bool, + tp_size: int = 1, + **extra_llm_kwargs): llm_args_extra = {} if pytorch_backend: LLM_CLASS = LLM_torch @@ -2507,12 +2516,16 @@ def _test_llm_capture_request_error(pytorch_backend: bool, tp_size: int = 1): model=llama_model_path, tensor_parallel_size=tp_size, **llm_args_extra, + **extra_llm_kwargs, ) prompt = 'A ' * 65 # the minimum max_num_tokens is 64 # Both backends now consistently raise RequestError for max_num_tokens validation - with pytest.raises(RequestError): - llm.generate(prompt) + try: + with pytest.raises(RequestError): + llm.generate(prompt) + finally: + llm.shutdown() def test_llm_capture_request_error(): diff --git a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py index 0e15b38b8b95..4a194a6e820c 100644 --- a/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py +++ b/tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py @@ -4,6 +4,12 @@ from tensorrt_llm import LLM from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.llmapi._grouped_test_utils import \ + make_shared_llm as _make_shared_llm +from tensorrt_llm.llmapi._grouped_test_utils import \ + mpi_session_kwargs as _mpi_session_kwargs +from tensorrt_llm.llmapi._grouped_test_utils import \ + shared_mpi_session as _shared_mpi_session from tensorrt_llm.lora_helper import LoraConfig from tensorrt_llm.sampling_params import SamplingParams @@ -21,68 +27,11 @@ global_kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4) -@pytest.mark.gpu2 -def test_llm_capture_request_error(): - _test_llm_capture_request_error(pytorch_backend=True, tp_size=2) - - -@pytest.mark.gpu4 -def test_tinyllama_logits_processor_tp2pp2(): - tinyllama_logits_processor_test_harness(backend="pytorch", - tensor_parallel_size=2, - pipeline_parallel_size=2) - - -@pytest.mark.gpu2 -@pytest.mark.part0 -@pytest.mark.parametrize("tp_size, pp_size", [(1, 2), (2, 1)]) -def test_tinyllama_logits_processor_2gpu(tp_size: int, pp_size: int): - tinyllama_logits_processor_test_harness(backend="pytorch", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size) - - -@pytest.mark.gpu2 -def test_llama_7b_lora_tp2(): - llama_7b_lora_from_dir_test_harness(tensor_parallel_size=2, - kv_cache_config=global_kv_cache_config) - - -@pytest.mark.gpu4 -@skip_ray # https://nvbugs/5682551 -@test_lora_with_and_without_cuda_graph -def test_llama_7b_multi_lora_tp4(cuda_graph_config): - # For LoRA checkpoints without finetuned embedding and lm_head, we can either: - # (1) specify lora_target_modules, or - # (2) provide a lora_dir to infer the lora_target_modules. - lora_config = LoraConfig(lora_target_modules=['attn_q', 'attn_k', 'attn_v'], - max_lora_rank=8, - max_loras=1, - max_cpu_loras=8) - check_llama_7b_multi_lora_from_request_test_harness( - LLM, - lora_config=lora_config, - tensor_parallel_size=4, - kv_cache_config=global_kv_cache_config, - cuda_graph_config=cuda_graph_config) - - -@skip_ray # https://nvbugs/5727075 -@pytest.mark.gpu2 -@test_lora_with_and_without_cuda_graph -def test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1( - cuda_graph_config) -> None: - check_phi3_lora_fused_modules_output_tp2_identical_to_tp1( - LLM, cuda_graph_config=cuda_graph_config) - - -@skip_ray -@pytest.mark.gpu2 -def test_llm_rpc_tp2(): - with LLM(model=llama_model_path, - kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), - orchestrator_type="rpc", - tensor_parallel_size=2) as llm: +def _run_llm_rpc_tp2(make_llm=LLM): + with make_llm(model=llama_model_path, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), + orchestrator_type="rpc", + tensor_parallel_size=2) as llm: assert isinstance(llm._executor, GenerationExecutorRpcProxy) res = llm.generate("Tell me a joke", @@ -94,14 +43,11 @@ def test_llm_rpc_tp2(): assert len(res.outputs[0].token_ids) == 10 -@skip_ray -@pytest.mark.gpu2 -@pytest.mark.asyncio -async def test_llm_rpc_streaming_tp2(): - with LLM(model=llama_model_path, - kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), - orchestrator_type="rpc", - tensor_parallel_size=2) as llm: +async def _run_llm_rpc_streaming_tp2(make_llm=LLM): + with make_llm(model=llama_model_path, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), + orchestrator_type="rpc", + tensor_parallel_size=2) as llm: assert isinstance(llm._executor, GenerationExecutorRpcProxy) async for output in llm.generate_async("Tell me a joke", @@ -110,81 +56,178 @@ async def test_llm_rpc_streaming_tp2(): print(f"get result: {output}") -@skip_ray -@pytest.mark.gpu2 -@pytest.mark.parametrize( - "prompt_logprobs, logprobs, return_context_logits, return_generation_logits", - [ - (None, 1, False, - False), # generation logprobs only (top-1, PyTorch limit) - ]) -def test_llm_return_logprobs_streaming_tp2(prompt_logprobs, logprobs, - return_context_logits, - return_generation_logits): - llm_return_logprobs_test_harness(prompt_logprobs, - logprobs, - return_context_logits, - return_generation_logits, - streaming=True, - backend="pytorch", - tp_size=2) - - -@skip_ray -@pytest.mark.gpu2 -@pytest.mark.parametrize( - "return_context_logits, enable_chunked_prefill, enable_iter_req_stats", - [ - (False, False, True), - (False, True, True), - ], -) -def test_llm_get_stats_pp2(return_context_logits, enable_chunked_prefill, - enable_iter_req_stats): - llm_get_stats_test_harness( - tp_size=1, - pp_size=2, - return_context_logits=return_context_logits, - pytorch_backend=True, - enable_chunked_prefill=enable_chunked_prefill, - enable_iter_req_stats=enable_iter_req_stats, - ) +class _MultiGpuLlmTests: + """Base for the multi-GPU LLM API test classes. + Subclasses set ``n_gpus``; the class-scoped fixtures below build one shared + MpiPoolSession (of that many workers) reused across the subclass's tests and + torn down when the subclass finishes -- so the 2-GPU and 4-GPU pools never + coexist. Fixtures are lazy, so tests that opt out (e.g. constructing a bare + ``LLM``) simply don't request them and pay no session cost. + """ + n_gpus: int -@skip_ray -@pytest.mark.gpu4 -@pytest.mark.parametrize( - "return_context_logits, enable_chunked_prefill, enable_iter_req_stats", - [ - (False, False, True), - (False, True, True), - ], -) -def test_llm_get_stats_pp4(return_context_logits, enable_chunked_prefill, - enable_iter_req_stats): - llm_get_stats_test_harness( - tp_size=1, - pp_size=4, - return_context_logits=return_context_logits, - pytorch_backend=True, - enable_chunked_prefill=enable_chunked_prefill, - enable_iter_req_stats=enable_iter_req_stats, - ) + @pytest.fixture(scope="class") + def mpi_session(self): + yield from _shared_mpi_session(self.n_gpus) + @pytest.fixture(scope="class") + def mpi_kwargs(self, mpi_session): + return _mpi_session_kwargs(mpi_session) -@skip_ray -@pytest.mark.gpu2 -def test_llm_get_stats_tp2(): - llm_get_stats_test_harness(tp_size=2, pytorch_backend=True) + @pytest.fixture(scope="class") + def shared_llm(self, mpi_session): + return _make_shared_llm(mpi_session) -@skip_ray @pytest.mark.gpu2 -def test_llm_get_stats_async_tp2(): - llm_get_stats_async_test_harness(tp_size=2, pytorch_backend=True) +class TestLlmMultiGpu2gpu(_MultiGpuLlmTests): + """2-GPU multi-GPU LLM API tests (shared session from _MultiGpuLlmTests).""" + n_gpus = 2 + + def test_llm_capture_request_error(self, mpi_kwargs): + _test_llm_capture_request_error(pytorch_backend=True, + tp_size=2, + **mpi_kwargs) + + @pytest.mark.part0 + @pytest.mark.parametrize("tp_size, pp_size", [(1, 2), (2, 1)]) + def test_tinyllama_logits_processor_2gpu(self, tp_size, pp_size, + mpi_kwargs): + tinyllama_logits_processor_test_harness(backend="pytorch", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + **mpi_kwargs) + + def test_llama_7b_lora_tp2(self, mpi_kwargs): + llama_7b_lora_from_dir_test_harness( + tensor_parallel_size=2, + kv_cache_config=global_kv_cache_config, + **mpi_kwargs) + + @skip_ray # https://nvbugs/5727075 + @test_lora_with_and_without_cuda_graph + def test_phi3_lora_fused_modules_output_on_tp2_identical_to_tp1( + self, cuda_graph_config) -> None: + check_phi3_lora_fused_modules_output_tp2_identical_to_tp1( + LLM, cuda_graph_config=cuda_graph_config) + + @skip_ray + def test_llm_rpc_tp2(self, shared_llm): + _run_llm_rpc_tp2(shared_llm) + + @skip_ray + @pytest.mark.asyncio + async def test_llm_rpc_streaming_tp2(self, shared_llm): + await _run_llm_rpc_streaming_tp2(shared_llm) + + @skip_ray + @pytest.mark.parametrize( + "prompt_logprobs, logprobs, return_context_logits, return_generation_logits", + [ + (None, 1, False, + False), # generation logprobs only (top-1, PyTorch limit) + ]) + def test_llm_return_logprobs_streaming_tp2(self, prompt_logprobs, logprobs, + return_context_logits, + return_generation_logits, + mpi_kwargs): + llm_return_logprobs_test_harness(prompt_logprobs, + logprobs, + return_context_logits, + return_generation_logits, + streaming=True, + backend="pytorch", + tp_size=2, + **mpi_kwargs) + + @skip_ray + @pytest.mark.parametrize( + "return_context_logits, enable_chunked_prefill, enable_iter_req_stats", + [ + (False, False, True), + (False, True, True), + ], + ) + def test_llm_get_stats_pp2(self, return_context_logits, + enable_chunked_prefill, enable_iter_req_stats, + mpi_kwargs): + llm_get_stats_test_harness( + tp_size=1, + pp_size=2, + return_context_logits=return_context_logits, + pytorch_backend=True, + enable_chunked_prefill=enable_chunked_prefill, + enable_iter_req_stats=enable_iter_req_stats, + **mpi_kwargs, + ) + + @skip_ray + def test_llm_get_stats_tp2(self, mpi_kwargs): + llm_get_stats_test_harness(tp_size=2, + pytorch_backend=True, + **mpi_kwargs) + + @skip_ray + def test_llm_get_stats_async_tp2(self, mpi_kwargs): + llm_get_stats_async_test_harness(tp_size=2, + pytorch_backend=True, + **mpi_kwargs) + + @skip_ray + def test_llm_get_stats_async_pp2(self, mpi_kwargs): + llm_get_stats_async_test_harness(pp_size=2, + pytorch_backend=True, + **mpi_kwargs) -@skip_ray -@pytest.mark.gpu2 -def test_llm_get_stats_async_pp2(): - llm_get_stats_async_test_harness(pp_size=2, pytorch_backend=True) +@pytest.mark.gpu4 +class TestLlmMultiGpu4gpu(_MultiGpuLlmTests): + """4-GPU multi-GPU LLM API tests (shared session from _MultiGpuLlmTests).""" + n_gpus = 4 + + def test_tinyllama_logits_processor_tp2pp2(self, mpi_kwargs): + tinyllama_logits_processor_test_harness(backend="pytorch", + tensor_parallel_size=2, + pipeline_parallel_size=2, + **mpi_kwargs) + + @skip_ray # https://nvbugs/5682551 + @test_lora_with_and_without_cuda_graph + def test_llama_7b_multi_lora_tp4(self, cuda_graph_config, mpi_kwargs): + # For LoRA checkpoints without finetuned embedding and lm_head, we can + # either: (1) specify lora_target_modules, or (2) provide a lora_dir to + # infer the lora_target_modules. + lora_config = LoraConfig( + lora_target_modules=['attn_q', 'attn_k', 'attn_v'], + max_lora_rank=8, + max_loras=1, + max_cpu_loras=8) + check_llama_7b_multi_lora_from_request_test_harness( + LLM, + lora_config=lora_config, + tensor_parallel_size=4, + kv_cache_config=global_kv_cache_config, + cuda_graph_config=cuda_graph_config, + **mpi_kwargs) + + @skip_ray + @pytest.mark.parametrize( + "return_context_logits, enable_chunked_prefill, enable_iter_req_stats", + [ + (False, False, True), + (False, True, True), + ], + ) + def test_llm_get_stats_pp4(self, return_context_logits, + enable_chunked_prefill, enable_iter_req_stats, + mpi_kwargs): + llm_get_stats_test_harness( + tp_size=1, + pp_size=4, + return_context_logits=return_context_logits, + pytorch_backend=True, + enable_chunked_prefill=enable_chunked_prefill, + enable_iter_req_stats=enable_iter_req_stats, + **mpi_kwargs, + )