[None][perf] Run the weight load on more torch CPU threads on request

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -5,6 +5,7 @@
 import inspect
 import os
 import sys
+import threading
 import traceback
 import warnings
 from contextlib import AbstractContextManager, contextmanager
@@ -87,6 +88,7 @@
                           "mixed_nvfp4_fp8", "mixed_nvfp4_float16",
                           "mixed_nvfp4_bfloat16")
 _SHADOW_WEIGHT_LOAD_PLAN_ENV = "TRTLLM_SHADOW_WEIGHT_LOAD_PLAN"
+_WEIGHT_LOAD_THREADS_ENV = "TLLM_WEIGHT_LOAD_THREADS"
 
 # Dense models do not consume the MoE backend or MoE mapping dimensions. Their
 # runtime topology remains bounded by the general and attention dimensions.
@@ -741,6 +743,58 @@ def _timed_checkpoint_weight_session(
             session.__exit__(None, None, None)
 
 
+def _weight_load_num_threads() -> Optional[int]:
+    """Threads ``TLLM_WEIGHT_LOAD_THREADS`` asks for, capped at usable CPUs."""
+    value = os.environ.get(_WEIGHT_LOAD_THREADS_ENV)
+    if value is None or not value.strip():
+        return None
+    try:
+        requested = int(value)
+    except ValueError as error:
+        raise ValueError(f"{_WEIGHT_LOAD_THREADS_ENV} must be a positive "
+                         f"integer, got {value!r}") from error
+    if requested <= 0:
+        raise ValueError(f"{_WEIGHT_LOAD_THREADS_ENV} must be a positive "
+                         f"integer, got {value!r}")
+    try:
+        available = len(os.sched_getaffinity(0))
+    except (AttributeError, OSError):
+        available = os.cpu_count() or 1
+    return min(requested, available)
+
+
+@contextmanager
+def _weight_load_threads() -> Iterator[None]:
+    """Run this thread's load on ``TLLM_WEIGHT_LOAD_THREADS`` torch threads.
+
+    Unset keeps the process's count. Only raises; restores on exit and on error.
+    """
+    requested = _weight_load_num_threads()
+    previous = torch.get_num_threads()
+    if requested is None or requested <= previous:
+        yield
+        return
+    # The native backend sets one count per process, once; a raise would stay.
+    if "ATen parallel backend: OpenMP" not in torch.__config__.parallel_info():
+        logger.warning(f"Ignoring {_WEIGHT_LOAD_THREADS_ENV}: only torch's "
+                       "OpenMP backend can restore the thread count")
+        yield
+        return
+    logger.info(f"Weight loading uses {requested} torch threads (was "
+                f"{previous}), as {_WEIGHT_LOAD_THREADS_ENV} asks")
+    try:
+        torch.set_num_threads(requested)
+        # New threads start from torch's process-wide count; set it back from
+        # another thread so this one keeps the raise and loader pools do not.
+        setter = threading.Thread(target=torch.set_num_threads,
+                                  args=(previous, ))
+        setter.start()
+        setter.join()
+        yield
+    finally:
+        torch.set_num_threads(previous)
+
+
 def _inspect_shadow_weight_load_plan(
         checkpoint_loader: Any, checkpoint_dir: str, weight_mapper: Any,
         **kwargs: Any) -> tuple[CheckpointCatalog, WeightLoadPlan] | None:
@@ -1497,13 +1551,14 @@ def init_meta_tensor(t: torch.Tensor):
 
                 weights_checkpoint_dir = (model.llm_checkpoint_dir if hasattr(
                     model, 'llm_checkpoint_dir') else checkpoint_dir)
-                weights_preloaded = self._materialize_checkpoint_weights(
-                    checkpoint_loader, weights_checkpoint_dir, model, config,
-                    load_weights_kwargs)
-
-                if loads_draft_weights:
-                    self._materialize_draft_checkpoint_weights(
-                        checkpoint_loader, model)
+                with _weight_load_threads():
+                    weights_preloaded = self._materialize_checkpoint_weights(
+                        checkpoint_loader, weights_checkpoint_dir, model,
+                        config, load_weights_kwargs)
+
+                    if loads_draft_weights:
+                        self._materialize_draft_checkpoint_weights(
+                            checkpoint_loader, model)
 
             elif load_format == LoadFormat.GMS:
                 # GPU Memory Service path: weight tensors live in a

tests/unittest/_torch/pyexecutor/test_model_loader_load_threads.py

@@ -0,0 +1,98 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""TLLM_WEIGHT_LOAD_THREADS: a real checkpoint loads to the same bytes on more torch threads."""
+
+import os
+
+import pytest
+import torch
+from transformers import LlamaConfig
+from transformers import LlamaForCausalLM as HFLlamaForCausalLM
+
+from tensorrt_llm._torch.models.modeling_llama import LlamaForCausalLM
+from tensorrt_llm._torch.modules.linear import Linear
+from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader, construct_checkpoint_loader
+from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
+from tensorrt_llm.mapping import Mapping
+
+
+def _load(checkpoint_dir: str, monkeypatch: pytest.MonkeyPatch, load_threads: str | None):
+    """Load the checkpoint; return every tensor's bytes and the thread counts the loads ran with."""
+    seen = {"model": set(), "modules": set()}
+    model_load = LlamaForCausalLM.load_weights
+    module_load = Linear.load_weights
+
+    def load_model(self, weights, weight_mapper=None, params_map=None, allow_partial_loading=False):
+        seen["model"].add(torch.get_num_threads())
+        model_load(
+            self,
+            weights,
+            weight_mapper=weight_mapper,
+            params_map=params_map,
+            allow_partial_loading=allow_partial_loading,
+        )
+
+    def load_module(self, weights, allow_partial_loading=False):
+        seen["modules"].add(torch.get_num_threads())
+        module_load(self, weights, allow_partial_loading=allow_partial_loading)
+
+    with monkeypatch.context() as patch:
+        if load_threads is None:
+            patch.delenv("TLLM_WEIGHT_LOAD_THREADS", raising=False)
+        else:
+            patch.setenv("TLLM_WEIGHT_LOAD_THREADS", load_threads)
+        patch.setattr(LlamaForCausalLM, "load_weights", load_model)
+        patch.setattr(Linear, "load_weights", load_module)
+        loader = ModelLoader(
+            llm_args=TorchLlmArgs(model=checkpoint_dir),
+            mapping=Mapping(),
+            spec_config=None,
+            sparse_attention_config=None,
+            max_num_tokens=64,
+            max_seq_len=64,
+        )
+        model, _ = loader.load(checkpoint_dir, construct_checkpoint_loader(None, "HF"))
+    tensors = dict(model.named_parameters())
+    tensors.update((f"buffer:{name}", buf) for name, buf in model.named_buffers())
+    data = {
+        name: t.detach().cpu().contiguous().view(-1).view(torch.uint8)
+        for name, t in tensors.items()
+    }
+    return data, seen
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
+def test_load_threads_same_bytes_and_restored(tmp_path, monkeypatch):
+    """Unset keeps 1 thread; N raises model.load_weights only, and 1 is back after the load."""
+    # Skip on the native backend only; if the loader stops seeing OpenMP, this fails.
+    if "ATen parallel backend: native" in torch.__config__.parallel_info():
+        pytest.skip("torch's native backend keeps one thread count per process")
+    config = LlamaConfig(
+        hidden_size=256,
+        intermediate_size=768,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        vocab_size=1024,
+        max_position_embeddings=64,
+    )
+    torch.manual_seed(0)
+    HFLlamaForCausalLM(config).to(torch.bfloat16).save_pretrained(tmp_path)
+    monkeypatch.delenv("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", raising=False)
+    raised = min(4, len(os.sched_getaffinity(0)))
+    saved = torch.get_num_threads()
+    # Start like a worker launched with OMP_NUM_THREADS=1.
+    torch.set_num_threads(1)
+    try:
+        off, off_threads = _load(str(tmp_path), monkeypatch, None)
+        assert torch.get_num_threads() == 1
+        on, on_threads = _load(str(tmp_path), monkeypatch, "4")
+        assert torch.get_num_threads() == 1
+    finally:
+        torch.set_num_threads(saved)
+
+    assert off_threads == {"model": {1}, "modules": {1}}
+    assert on_threads == {"model": {raised}, "modules": {1}}
+    assert off.keys() == on.keys()
+    for name, data in off.items():
+        assert torch.equal(data, on[name]), name