From b1f8b3119bf75b882a029017e68a04e345233146 Mon Sep 17 00:00:00 2001 From: functionstackx <47992694+functionstackx@users.noreply.github.com> Date: Sun, 20 Sep 2026 15:46:54 -0400 Subject: [PATCH 1/4] feat: add ModelScope model loading support Signed-off-by: functionstackx <47992694+functionstackx@users.noreply.github.com> --- docs/source/llm-api/index.md | 16 ++++++ tensorrt_llm/llmapi/llm.py | 16 +++--- tensorrt_llm/llmapi/llm_args.py | 24 ++++++++- tensorrt_llm/llmapi/utils.py | 76 +++++++++++++++++++------- tests/unittest/llmapi/test_utils.py | 82 +++++++++++++++++++++++++++++ 5 files changed, 188 insertions(+), 26 deletions(-) diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index 2d0a9f0cf64f..9a296e5826b7 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -31,6 +31,22 @@ llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0") You can also use [quantized checkpoints](https://huggingface.co/collections/nvidia/model-optimizer-66aa84f7966b3150262481a4) (FP4, FP8, etc) of popular models provided by NVIDIA in the same way. +### Using a Model from ModelScope + +To resolve remote model IDs through [ModelScope](https://modelscope.cn/) +instead of the Hugging Face Hub, install the optional client and enable the +ModelScope download path before starting TensorRT-LLM: + +```console +pip install modelscope +export TRTLLM_USE_MODELSCOPE=true +trtllm-serve Qwen/Qwen3-0.6B +``` + +The switch also applies to remote tokenizer and speculative-model IDs. Local +paths are used as-is. Unset `TRTLLM_USE_MODELSCOPE`, or set it to `false`, to +retain the default Hugging Face behavior. + ### 2. Using a Local Hugging Face Model To use a model from local storage, first download it manually: diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 21651b3b8e5c..da22aa1f08fc 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -1675,6 +1675,8 @@ def _try_load_tokenizer(self) -> Optional[TokenizerBase]: assert isinstance(self.args.tokenizer, TokenizerBase) return self.args.tokenizer + model_path = self._hf_model_dir or self.args.model + # TODO smor- need to refine what is the desired behavior if lora is enabled # in terms of the tokenizer initialization process if hasattr(self.args, "backend") and self.args.backend in [ @@ -1689,15 +1691,15 @@ def _try_load_tokenizer(self) -> Optional[TokenizerBase]: trust_remote_code=self.args.trust_remote_code, use_fast=self.args.tokenizer_mode != 'slow') if tokenizer is None: - tokenizer_path = self.args.model + tokenizer_path = model_path else: return tokenizer except Exception: - tokenizer_path = self.args.model + tokenizer_path = model_path else: - tokenizer_path = self.args.model + tokenizer_path = model_path else: - tokenizer_path = self.args.model + tokenizer_path = model_path return ModelLoader.load_hf_tokenizer( tokenizer_path, trust_remote_code=self.args.trust_remote_code, @@ -1716,7 +1718,8 @@ def tokenizer(self, tokenizer: TokenizerBase): def _try_load_generation_config( self) -> Optional[transformers.GenerationConfig]: - return ModelLoader.load_hf_generation_config(self.args.model) + model_dir = self._hf_model_dir or self.args.model + return ModelLoader.load_hf_generation_config(model_dir) def _try_load_generation_config_explicit_values(self) -> dict[str, Any]: if self.args.backend != "pytorch" or self.args.generation_config != "auto": @@ -1726,8 +1729,9 @@ def _try_load_generation_config_explicit_values(self) -> dict[str, Any]: def _try_load_hf_model_config( self) -> Optional[transformers.PretrainedConfig]: + model_dir = self._hf_model_dir or self.args.model return ModelLoader.load_hf_model_config( - self.args.model, trust_remote_code=self.args.trust_remote_code) + model_dir, trust_remote_code=self.args.trust_remote_code) @set_api_status("prototype") def start_profile(self, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e4ec65d8027e..1e4d8cd2abb4 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -75,8 +75,9 @@ from ..usage.config import UsageContext # noqa: F401 from ..usage.config import TelemetryConfig, TelemetryField from .tokenizer import TokenizerBase, tokenizer_factory -from .utils import (StrictBaseModel, generate_api_docs_as_docstring, - get_type_repr) +from .utils import (StrictBaseModel, download_hf_partial, + generate_api_docs_as_docstring, get_type_repr, + use_modelscope) TypeBaseModel = TypeVar("T", bound=BaseModel) @@ -5267,6 +5268,15 @@ def validate_and_init_tokenizer(self): # Use tokenizer path if specified, otherwise use model path. load_path = self.tokenizer if self.tokenizer else self.model + if (use_modelscope() and isinstance(load_path, (str, Path)) + and not Path(load_path).exists()): + load_path = download_hf_partial( + str(load_path), + [ + "*.json", "*.jinja", "*.j2", "*.model", "*.py", + "*.tiktoken", "*.txt" + ], + revision=self.tokenizer_revision) # The one loader for aliases and import paths; it raises # ValueError("Failed to load custom tokenizer ...") on failure. self.tokenizer = load_custom_tokenizer( @@ -5275,6 +5285,16 @@ def validate_and_init_tokenizer(self): trust_remote_code=self.trust_remote_code, use_fast=self.tokenizer_mode != 'slow') else: + if (use_modelscope() + and isinstance(self.tokenizer, (str, Path)) + and not Path(self.tokenizer).exists()): + self.tokenizer = download_hf_partial( + str(self.tokenizer), + [ + "*.json", "*.jinja", "*.j2", "*.model", "*.py", + "*.tiktoken", "*.txt" + ], + revision=self.tokenizer_revision) self.tokenizer = tokenizer_factory( self.tokenizer, trust_remote_code=self.trust_remote_code, diff --git a/tensorrt_llm/llmapi/utils.py b/tensorrt_llm/llmapi/utils.py index 386b296ef69d..94f14ea667f2 100644 --- a/tensorrt_llm/llmapi/utils.py +++ b/tensorrt_llm/llmapi/utils.py @@ -27,7 +27,7 @@ import huggingface_hub import psutil import torch -from huggingface_hub import snapshot_download +from huggingface_hub import snapshot_download as hf_snapshot_download from pydantic import BaseModel from tqdm.auto import tqdm @@ -231,22 +231,20 @@ def __init__(self, *args, **kwargs): def download_hf_model(model: str, revision: Optional[str] = None) -> Path: ignore_patterns = ["original/**/*"] - logger.info(f"Downloading model {model} from HuggingFace") + hub_name = "ModelScope" if use_modelscope() else "Hugging Face" + logger.info(f"Downloading model {model} from {hub_name}") with get_file_lock(model): - hf_folder = snapshot_download( - model, - local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, - ignore_patterns=ignore_patterns, - revision=revision, - tqdm_class=DisabledTqdm) - logger.info(f"Finished downloading model {model} from HuggingFace") - return Path(hf_folder) + model_folder = _snapshot_download(model, + ignore_patterns=ignore_patterns, + revision=revision) + logger.info(f"Finished downloading model {model} from {hub_name}") + return Path(model_folder) def download_hf_partial(model: str, allow_patterns: List[str], revision: Optional[str] = None) -> Path: - """Download a partial model from HuggingFace. + """Download selected model files from the configured model hub. Args: model: The model name or path. @@ -257,13 +255,55 @@ def download_hf_partial(model: str, The path to the downloaded model. """ with get_file_lock(model): - hf_folder = snapshot_download( - model, - local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, - revision=revision, - allow_patterns=allow_patterns, - tqdm_class=DisabledTqdm) - return Path(hf_folder) + model_folder = _snapshot_download(model, + revision=revision, + allow_patterns=allow_patterns) + return Path(model_folder) + + +def use_modelscope() -> bool: + """Return whether remote model IDs should resolve through ModelScope.""" + return os.environ.get("TRTLLM_USE_MODELSCOPE", "false").strip().lower( + ) in ("1", "true") + + +def _snapshot_download(model: str, + revision: Optional[str] = None, + ignore_patterns: Optional[List[str]] = None, + allow_patterns: Optional[List[str]] = None) -> str: + """Download a snapshot from ModelScope or Hugging Face. + + ModelScope uses different names for its file filters. Keep the optional + import in this boundary so standard TensorRT-LLM installations do not need + the ``modelscope`` package. + """ + local_files_only = huggingface_hub.constants.HF_HUB_OFFLINE + if use_modelscope(): + try: + from modelscope.hub.snapshot_download import snapshot_download + except ImportError as error: + raise ImportError( + "TRTLLM_USE_MODELSCOPE is enabled, but ModelScope is not " + "installed. Install it with `pip install modelscope`.") from error + + kwargs = { + "model_id": model, + "local_files_only": local_files_only, + "revision": revision, + } + if ignore_patterns: + kwargs["ignore_file_pattern"] = ignore_patterns + if allow_patterns: + kwargs["allow_file_pattern"] = allow_patterns + return snapshot_download(**kwargs) + + return hf_snapshot_download( + model, + local_files_only=local_files_only, + ignore_patterns=ignore_patterns, + allow_patterns=allow_patterns, + revision=revision, + tqdm_class=DisabledTqdm) def download_hf_pretrained_config(model: str, diff --git a/tests/unittest/llmapi/test_utils.py b/tests/unittest/llmapi/test_utils.py index b0681432a14b..e86b4eb93fe6 100644 --- a/tests/unittest/llmapi/test_utils.py +++ b/tests/unittest/llmapi/test_utils.py @@ -13,6 +13,8 @@ from tensorrt_llm.llmapi.utils import (ApiStatusRegistry, _set_affinity_all_threads, configure_cpu_affinity, + download_hf_model, + download_hf_partial, generate_api_docs_as_docstring) _TASK_DIR = "/proc/self/task" @@ -20,6 +22,86 @@ pytestmark = pytest.mark.cpu_only +def _stub_modelscope(monkeypatch, snapshot_download): + modelscope = types.ModuleType("modelscope") + hub = types.ModuleType("modelscope.hub") + snapshot_module = types.ModuleType("modelscope.hub.snapshot_download") + snapshot_module.snapshot_download = snapshot_download + monkeypatch.setitem(sys.modules, "modelscope", modelscope) + monkeypatch.setitem(sys.modules, "modelscope.hub", hub) + monkeypatch.setitem(sys.modules, "modelscope.hub.snapshot_download", + snapshot_module) + + +def test_modelscope_download_maps_snapshot_filters(monkeypatch, tmp_path): + calls = [] + + def snapshot_download(**kwargs): + calls.append(kwargs) + return str(tmp_path) + + _stub_modelscope(monkeypatch, snapshot_download) + monkeypatch.setenv("TRTLLM_USE_MODELSCOPE", "true") + monkeypatch.setattr(llmapi_utils.huggingface_hub.constants, + "HF_HUB_OFFLINE", True) + + downloaded = download_hf_partial("Qwen/Qwen3-0.6B", ["*.json"], + revision="v1") + + assert downloaded == tmp_path + assert calls == [{ + "model_id": "Qwen/Qwen3-0.6B", + "local_files_only": True, + "revision": "v1", + "allow_file_pattern": ["*.json"], + }] + + +def test_modelscope_download_maps_ignored_files(monkeypatch, tmp_path): + calls = [] + + def snapshot_download(**kwargs): + calls.append(kwargs) + return str(tmp_path) + + _stub_modelscope(monkeypatch, snapshot_download) + monkeypatch.setenv("TRTLLM_USE_MODELSCOPE", "1") + + downloaded = download_hf_model("Qwen/Qwen3-0.6B") + + assert downloaded == tmp_path + assert calls[0]["ignore_file_pattern"] == ["original/**/*"] + + +def test_hugging_face_download_remains_the_default(monkeypatch, tmp_path): + calls = [] + + def snapshot_download(model, **kwargs): + calls.append((model, kwargs)) + return str(tmp_path) + + monkeypatch.delenv("TRTLLM_USE_MODELSCOPE", raising=False) + monkeypatch.setattr(llmapi_utils, "hf_snapshot_download", + snapshot_download) + + downloaded = download_hf_partial("Qwen/Qwen3-0.6B", ["config.json"]) + + assert downloaded == tmp_path + assert calls[0][0] == "Qwen/Qwen3-0.6B" + assert calls[0][1]["allow_patterns"] == ["config.json"] + + +def test_modelscope_download_requires_optional_dependency(monkeypatch): + monkeypatch.setenv("TRTLLM_USE_MODELSCOPE", "true") + monkeypatch.setitem(sys.modules, "modelscope", None) + monkeypatch.delitem(sys.modules, + "modelscope.hub.snapshot_download", + raising=False) + + with pytest.raises(ImportError, match="pip install modelscope"): + download_hf_model("Qwen/Qwen3-0.6B") + + def test_api_status_registry(): @ApiStatusRegistry.set_api_status("beta") From 677d37c443d679f131564b6bd101b3646d4b64f4 Mon Sep 17 00:00:00 2001 From: functionstackx <47992694+functionstackx@users.noreply.github.com> Date: Sun, 20 Sep 2026 16:03:25 -0400 Subject: [PATCH 2/4] fix: preserve glob semantics for ModelScope filters Signed-off-by: functionstackx <47992694+functionstackx@users.noreply.github.com> --- docs/source/llm-api/index.md | 2 +- tensorrt_llm/llmapi/utils.py | 29 ++++++++++++++--------------- tests/unittest/llmapi/test_utils.py | 12 +++++------- 3 files changed, 20 insertions(+), 23 deletions(-) diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index 9a296e5826b7..6a043857ca9e 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -38,7 +38,7 @@ instead of the Hugging Face Hub, install the optional client and enable the ModelScope download path before starting TensorRT-LLM: ```console -pip install modelscope +pip install 'modelscope>=1.20' export TRTLLM_USE_MODELSCOPE=true trtllm-serve Qwen/Qwen3-0.6B ``` diff --git a/tensorrt_llm/llmapi/utils.py b/tensorrt_llm/llmapi/utils.py index 94f14ea667f2..ded0033a865f 100644 --- a/tensorrt_llm/llmapi/utils.py +++ b/tensorrt_llm/llmapi/utils.py @@ -263,8 +263,8 @@ def download_hf_partial(model: str, def use_modelscope() -> bool: """Return whether remote model IDs should resolve through ModelScope.""" - return os.environ.get("TRTLLM_USE_MODELSCOPE", "false").strip().lower( - ) in ("1", "true") + return os.environ.get("TRTLLM_USE_MODELSCOPE", + "false").strip().lower() in ("1", "true") def _snapshot_download(model: str, @@ -273,9 +273,8 @@ def _snapshot_download(model: str, allow_patterns: Optional[List[str]] = None) -> str: """Download a snapshot from ModelScope or Hugging Face. - ModelScope uses different names for its file filters. Keep the optional - import in this boundary so standard TensorRT-LLM installations do not need - the ``modelscope`` package. + Keep the optional import in this boundary so standard TensorRT-LLM + installations do not need the ``modelscope`` package. """ local_files_only = huggingface_hub.constants.HF_HUB_OFFLINE if use_modelscope(): @@ -284,7 +283,8 @@ def _snapshot_download(model: str, except ImportError as error: raise ImportError( "TRTLLM_USE_MODELSCOPE is enabled, but ModelScope is not " - "installed. Install it with `pip install modelscope`.") from error + "installed. Install it with `pip install 'modelscope>=1.20'`." + ) from error kwargs = { "model_id": model, @@ -292,18 +292,17 @@ def _snapshot_download(model: str, "revision": revision, } if ignore_patterns: - kwargs["ignore_file_pattern"] = ignore_patterns + kwargs["ignore_patterns"] = ignore_patterns if allow_patterns: - kwargs["allow_file_pattern"] = allow_patterns + kwargs["allow_patterns"] = allow_patterns return snapshot_download(**kwargs) - return hf_snapshot_download( - model, - local_files_only=local_files_only, - ignore_patterns=ignore_patterns, - allow_patterns=allow_patterns, - revision=revision, - tqdm_class=DisabledTqdm) + return hf_snapshot_download(model, + local_files_only=local_files_only, + ignore_patterns=ignore_patterns, + allow_patterns=allow_patterns, + revision=revision, + tqdm_class=DisabledTqdm) def download_hf_pretrained_config(model: str, diff --git a/tests/unittest/llmapi/test_utils.py b/tests/unittest/llmapi/test_utils.py index e86b4eb93fe6..e72685ef25a1 100644 --- a/tests/unittest/llmapi/test_utils.py +++ b/tests/unittest/llmapi/test_utils.py @@ -13,8 +13,7 @@ from tensorrt_llm.llmapi.utils import (ApiStatusRegistry, _set_affinity_all_threads, configure_cpu_affinity, - download_hf_model, - download_hf_partial, + download_hf_model, download_hf_partial, generate_api_docs_as_docstring) _TASK_DIR = "/proc/self/task" @@ -53,7 +52,7 @@ def snapshot_download(**kwargs): "model_id": "Qwen/Qwen3-0.6B", "local_files_only": True, "revision": "v1", - "allow_file_pattern": ["*.json"], + "allow_patterns": ["*.json"], }] @@ -70,7 +69,7 @@ def snapshot_download(**kwargs): downloaded = download_hf_model("Qwen/Qwen3-0.6B") assert downloaded == tmp_path - assert calls[0]["ignore_file_pattern"] == ["original/**/*"] + assert calls[0]["ignore_patterns"] == ["original/**/*"] def test_hugging_face_download_remains_the_default(monkeypatch, tmp_path): @@ -81,8 +80,7 @@ def snapshot_download(model, **kwargs): return str(tmp_path) monkeypatch.delenv("TRTLLM_USE_MODELSCOPE", raising=False) - monkeypatch.setattr(llmapi_utils, "hf_snapshot_download", - snapshot_download) + monkeypatch.setattr(llmapi_utils, "hf_snapshot_download", snapshot_download) downloaded = download_hf_partial("Qwen/Qwen3-0.6B", ["config.json"]) @@ -98,7 +96,7 @@ def test_modelscope_download_requires_optional_dependency(monkeypatch): "modelscope.hub.snapshot_download", raising=False) - with pytest.raises(ImportError, match="pip install modelscope"): + with pytest.raises(ImportError, match="modelscope>=1.20"): download_hf_model("Qwen/Qwen3-0.6B") From c5c002e64b898fdb9c524df04c07ca2c8e736cb5 Mon Sep 17 00:00:00 2001 From: functionstackx <47992694+functionstackx@users.noreply.github.com> Date: Sun, 20 Sep 2026 17:24:58 -0400 Subject: [PATCH 3/4] feat(modelscope): add optional install extra and document scoped routing Validate ModelScope 1.20 glob filters, explain the scoped design, and apply repository formatting and copyright conventions. Signed-off-by: functionstackx <47992694+functionstackx@users.noreply.github.com> --- docs/source/llm-api/index.md | 16 ++++++++++++++++ setup.py | 1 + tensorrt_llm/llmapi/llm_args.py | 9 +++------ tensorrt_llm/llmapi/utils.py | 2 ++ 4 files changed, 22 insertions(+), 6 deletions(-) diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index 6a043857ca9e..ca6980251c65 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -47,6 +47,22 @@ The switch also applies to remote tokenizer and speculative-model IDs. Local paths are used as-is. Unset `TRTLLM_USE_MODELSCOPE`, or set it to `false`, to retain the default Hugging Face behavior. +You can also install the client through the `tensorrt_llm[modelscope]` extra. +ModelScope remains optional and is imported only when this switch is enabled. + +The integration routes TensorRT-LLM snapshot downloads through ModelScope and +passes the resolved local directories to its tokenizer and configuration loaders. +Explicit remote tokenizers download only tokenizer/configuration files, honoring +`tokenizer_revision`. It deliberately does not call ModelScope's process-wide +`patch_hub()`: unrelated Hugging Face consumers in the same process retain their +own hub behavior. Applications that download additional assets outside these +TensorRT-LLM paths must configure those consumers separately. + +The minimum client version, [ModelScope 1.20.0](https://github.com/modelscope/modelscope/blob/v1.20.0/modelscope/hub/snapshot_download.py), +supports the `allow_patterns` and `ignore_patterns` glob arguments used here. +The legacy `ignore_file_pattern` argument also interprets valid patterns as +regular expressions; it is not interchangeable with the glob-only filter. + ### 2. Using a Local Hugging Face Model To use a model from local storage, first download it manually: diff --git a/setup.py b/setup.py index 5ad206aa4325..78014bcb8dea 100644 --- a/setup.py +++ b/setup.py @@ -747,6 +747,7 @@ def get_build_state_options(): scripts=['tensorrt_llm/llmapi/trtllm-llmapi-launch'], extras_require={ "devel": devel_deps + grpc_smg_deps, + "modelscope": ["modelscope>=1.20"], "openengine": openengine_deps, "mx": mx_deps, "grpc-smg": grpc_smg_deps, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1e4d8cd2abb4..c100a7b396ee 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5271,8 +5271,7 @@ def validate_and_init_tokenizer(self): if (use_modelscope() and isinstance(load_path, (str, Path)) and not Path(load_path).exists()): load_path = download_hf_partial( - str(load_path), - [ + str(load_path), [ "*.json", "*.jinja", "*.j2", "*.model", "*.py", "*.tiktoken", "*.txt" ], @@ -5285,12 +5284,10 @@ def validate_and_init_tokenizer(self): trust_remote_code=self.trust_remote_code, use_fast=self.tokenizer_mode != 'slow') else: - if (use_modelscope() - and isinstance(self.tokenizer, (str, Path)) + if (use_modelscope() and isinstance(self.tokenizer, (str, Path)) and not Path(self.tokenizer).exists()): self.tokenizer = download_hf_partial( - str(self.tokenizer), - [ + str(self.tokenizer), [ "*.json", "*.jinja", "*.j2", "*.model", "*.py", "*.tiktoken", "*.txt" ], diff --git a/tensorrt_llm/llmapi/utils.py b/tensorrt_llm/llmapi/utils.py index ded0033a865f..ab70b0d512aa 100644 --- a/tensorrt_llm/llmapi/utils.py +++ b/tensorrt_llm/llmapi/utils.py @@ -1,3 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 import asyncio import collections import ctypes From 2da0a1085153a883a7e1c127150bc36b35da1bac Mon Sep 17 00:00:00 2001 From: functionstackx <47992694+functionstackx@users.noreply.github.com> Date: Sun, 20 Sep 2026 20:25:43 -0400 Subject: [PATCH 4/4] docs: add docstrings for the ModelScope download paths The docstring coverage check flagged this diff at 25% against an 80% threshold. Document the download helpers, the model-directory loaders in LLM, and the ModelScope routing tests, which brings coverage of the functions touched by this PR to 93%. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: functionstackx <47992694+functionstackx@users.noreply.github.com> --- tensorrt_llm/llmapi/llm.py | 25 +++++++++++++++++++++++++ tensorrt_llm/llmapi/utils.py | 9 +++++++++ tests/unittest/llmapi/test_utils.py | 13 +++++++++++++ 3 files changed, 47 insertions(+) diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index da22aa1f08fc..30c6dce36e0e 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -1668,6 +1668,15 @@ def _build_model(self): self._engine_dir, self._hf_model_dir = model_loader() def _try_load_tokenizer(self) -> Optional[TokenizerBase]: + """Resolve the tokenizer for this LLM instance. + + Prefers an explicitly supplied tokenizer, then a single LoRA + directory on the PyTorch backends, and otherwise falls back to the + downloaded model directory or the configured model reference. + + Returns: + The resolved tokenizer, or None when tokenizer init is skipped. + """ if self.args.skip_tokenizer_init: return None @@ -1718,6 +1727,14 @@ def tokenizer(self, tokenizer: TokenizerBase): def _try_load_generation_config( self) -> Optional[transformers.GenerationConfig]: + """Load the Hugging Face generation config for this model. + + Reads from the downloaded model directory when one is available so + that remotely fetched snapshots are not re-resolved. + + Returns: + The generation config, or None when the model does not ship one. + """ model_dir = self._hf_model_dir or self.args.model return ModelLoader.load_hf_generation_config(model_dir) @@ -1729,6 +1746,14 @@ def _try_load_generation_config_explicit_values(self) -> dict[str, Any]: def _try_load_hf_model_config( self) -> Optional[transformers.PretrainedConfig]: + """Load the Hugging Face model config for this model. + + Reads from the downloaded model directory when one is available so + that remotely fetched snapshots are not re-resolved. + + Returns: + The model config, or None when the model does not ship one. + """ model_dir = self._hf_model_dir or self.args.model return ModelLoader.load_hf_model_config( model_dir, trust_remote_code=self.args.trust_remote_code) diff --git a/tensorrt_llm/llmapi/utils.py b/tensorrt_llm/llmapi/utils.py index ab70b0d512aa..7579234a55ec 100644 --- a/tensorrt_llm/llmapi/utils.py +++ b/tensorrt_llm/llmapi/utils.py @@ -232,6 +232,15 @@ def __init__(self, *args, **kwargs): def download_hf_model(model: str, revision: Optional[str] = None) -> Path: + """Download a full model snapshot from the configured model hub. + + Args: + model: The model name or path. + revision: The revision to use for the model. + + Returns: + The path to the downloaded model. + """ ignore_patterns = ["original/**/*"] hub_name = "ModelScope" if use_modelscope() else "Hugging Face" logger.info(f"Downloading model {model} from {hub_name}") diff --git a/tests/unittest/llmapi/test_utils.py b/tests/unittest/llmapi/test_utils.py index e72685ef25a1..768a9ae31041 100644 --- a/tests/unittest/llmapi/test_utils.py +++ b/tests/unittest/llmapi/test_utils.py @@ -22,6 +22,12 @@ def _stub_modelscope(monkeypatch, snapshot_download): + """Register a fake ``modelscope`` package exposing ``snapshot_download``. + + Args: + monkeypatch: The pytest monkeypatch fixture. + snapshot_download: The callable to install as the hub entry point. + """ modelscope = types.ModuleType("modelscope") hub = types.ModuleType("modelscope.hub") snapshot_module = types.ModuleType("modelscope.hub.snapshot_download") @@ -33,9 +39,11 @@ def _stub_modelscope(monkeypatch, snapshot_download): def test_modelscope_download_maps_snapshot_filters(monkeypatch, tmp_path): + """Partial downloads forward allow patterns and revision to ModelScope.""" calls = [] def snapshot_download(**kwargs): + """Record the hub call and return the temporary snapshot path.""" calls.append(kwargs) return str(tmp_path) @@ -57,9 +65,11 @@ def snapshot_download(**kwargs): def test_modelscope_download_maps_ignored_files(monkeypatch, tmp_path): + """Full downloads forward the default ignore patterns to ModelScope.""" calls = [] def snapshot_download(**kwargs): + """Record the hub call and return the temporary snapshot path.""" calls.append(kwargs) return str(tmp_path) @@ -73,9 +83,11 @@ def snapshot_download(**kwargs): def test_hugging_face_download_remains_the_default(monkeypatch, tmp_path): + """Downloads route to Hugging Face when ModelScope is not enabled.""" calls = [] def snapshot_download(model, **kwargs): + """Record the hub call and return the temporary snapshot path.""" calls.append((model, kwargs)) return str(tmp_path) @@ -90,6 +102,7 @@ def snapshot_download(model, **kwargs): def test_modelscope_download_requires_optional_dependency(monkeypatch): + """A missing ``modelscope`` install raises an actionable ImportError.""" monkeypatch.setenv("TRTLLM_USE_MODELSCOPE", "true") monkeypatch.setitem(sys.modules, "modelscope", None) monkeypatch.delitem(sys.modules,