From 9c29119f217cb2be259657cd67934eb4c8ea41bd Mon Sep 17 00:00:00 2001 From: Jie Li Date: Fri, 5 Dec 2025 16:24:56 +0800 Subject: [PATCH] [https://nvbugs/5519544][feat] Adaptive number of thread workers when loading model concurrently Signed-off-by: Jie Li --- tensorrt_llm/_torch/models/modeling_utils.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index c17eacbefadf..7875d747a14f 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -795,6 +795,11 @@ def run_concurrently(func, reduce_func: an optional function to reduce the results. pbar: an optional tqdm progress bar. """ + if num_workers is None: + tp_size = int(os.getenv('TP_SIZE', '1')) + # default: 1 worker per TP rank, max 8 + num_workers = min(8, max(1, os.cpu_count() // tp_size)) + from concurrent import futures with futures.ThreadPoolExecutor(max_workers=num_workers) as executor: # Submit all tasks