From b23e97799cdeb439d053ce383df4631e1c19447b Mon Sep 17 00:00:00 2001 From: Rakib Hasan Date: Wed, 30 Jul 2025 23:46:18 +0000 Subject: [PATCH 1/3] refactoring input prep again to allow out-of-tree models to work with quickstart, trtllm-bench, etc. Signed-off-by: Rakib Hasan --- examples/llm-api/quickstart_multimodal.py | 44 +++- tensorrt_llm/__init__.py | 4 + .../_torch/models/modeling_gemma3vl.py | 12 +- .../_torch/models/modeling_hyperclovax.py | 22 +- tensorrt_llm/_torch/models/modeling_llama.py | 12 +- .../_torch/models/modeling_llava_next.py | 12 +- .../_torch/models/modeling_mistral.py | 20 +- tensorrt_llm/_torch/models/modeling_phi4mm.py | 16 +- .../_torch/models/modeling_qwen2vl.py | 24 +- tensorrt_llm/_torch/models/modeling_vila.py | 15 +- tensorrt_llm/bench/benchmark/throughput.py | 46 +++- tensorrt_llm/bench/utils/data.py | 6 +- tensorrt_llm/inputs/__init__.py | 16 +- tensorrt_llm/inputs/registry.py | 229 +++++++++++++++--- tensorrt_llm/inputs/utils.py | 115 +++------ tensorrt_llm/tools/importlib_utils.py | 134 ++++++++++ 16 files changed, 566 insertions(+), 161 deletions(-) create mode 100644 tensorrt_llm/tools/importlib_utils.py diff --git a/examples/llm-api/quickstart_multimodal.py b/examples/llm-api/quickstart_multimodal.py index fc18671ee280..c5f15a4ddce7 100644 --- a/examples/llm-api/quickstart_multimodal.py +++ b/examples/llm-api/quickstart_multimodal.py @@ -4,8 +4,9 @@ from quickstart_advanced import add_llm_args, setup_llm -from tensorrt_llm.inputs import (ALL_SUPPORTED_MULTIMODAL_MODELS, - default_multimodal_input_loader) +from tensorrt_llm.inputs import default_multimodal_input_loader +from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY +from tensorrt_llm.tools.importlib_utils import import_custom_module_from_dir example_medias_and_prompts = { "image": { @@ -79,10 +80,11 @@ def add_multimodal_args(parser): - parser.add_argument("--model_type", - type=str, - choices=ALL_SUPPORTED_MULTIMODAL_MODELS, - help="Model type.") + parser.add_argument( + "--model_type", + type=str, + choices=MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_model_types(), + help="Model type.") parser.add_argument("--modality", type=str, choices=[ @@ -108,6 +110,18 @@ def add_multimodal_args(parser): type=str, default="cpu", help="The device to have the input on.") + parser.add_argument( + "--custom_module_dirs", + type=str, + nargs="+", + default=None, + help= + ("Paths to an out-of-tree model directory which should be imported." + " This is useful to load a custom model. The directory should have a structure like:" + " " + " ├── __init__.py" + " ├── .py" + " └── ")) return parser @@ -140,6 +154,15 @@ def parse_arguments(): def main(): args = parse_arguments() + if args.custom_module_dirs is not None: + for custom_module_dir in args.custom_module_dirs: + try: + import_custom_module_from_dir(custom_module_dir) + except Exception as e: + print( + f"Failed to import custom module from {custom_module_dir}: {e}" + ) + raise e lora_config = None if args.load_lora: @@ -159,8 +182,11 @@ def main(): model_type = args.model_type else: model_type = json.load( - open(os.path.join(llm._hf_model_dir, 'config.json')))['model_type'] - assert model_type in ALL_SUPPORTED_MULTIMODAL_MODELS, f"Unsupported model_type: {model_type}" + open(os.path.join(str(llm._hf_model_dir), + 'config.json')))['model_type'] + assert model_type in MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_model_types(), \ + f"Unsupported model_type: {model_type} found!\n" \ + f"Supported types: {MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_model_types()}" # set prompts and media to example prompts and images if they are not provided if args.prompt is None: @@ -168,7 +194,7 @@ def main(): if args.media is None: args.media = example_medias_and_prompts[args.modality]["media"] inputs = default_multimodal_input_loader(tokenizer=llm.tokenizer, - model_dir=llm._hf_model_dir, + model_dir=str(llm._hf_model_dir), model_type=model_type, modality=args.modality, prompts=args.prompt, diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index f54026a8cbc0..911aab5e3d12 100644 --- a/tensorrt_llm/__init__.py +++ b/tensorrt_llm/__init__.py @@ -62,6 +62,9 @@ def _add_trt_llm_dll_directory(): from .sampling_params import SamplingParams from .version import __version__ +# Lazy import to avoid circular dependency on lora_manager +import tensorrt_llm._torch.models as torch_models # isort:skip + __all__ = [ 'AutoConfig', 'AutoModelForCausalLM', @@ -82,6 +85,7 @@ def _add_trt_llm_dll_directory(): 'default_trtnet', 'precision', 'net_guard', + 'torch_models', 'Network', 'Mapping', 'MnnvlMemory', diff --git a/tensorrt_llm/_torch/models/modeling_gemma3vl.py b/tensorrt_llm/_torch/models/modeling_gemma3vl.py index 699e464151d7..e305b82dbaa9 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3vl.py +++ b/tensorrt_llm/_torch/models/modeling_gemma3vl.py @@ -10,7 +10,9 @@ BaseWeightMapper from ..._utils import nvtx_range -from ...inputs import (ExtraProcessedInputs, InputProcessor, TextPrompt, +from ...inputs import (ExtraProcessedInputs, InputProcessor, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, TextPrompt, register_input_processor) from ...logger import logger from ...sampling_params import SamplingParams @@ -137,7 +139,13 @@ def forward(self, vision_outputs: torch.Tensor): @register_auto_model("Gemma3ForConditionalGeneration") -@register_input_processor(Gemma3InputProcessor, model_type="gemma3") +@register_input_processor( + Gemma3InputProcessor, + model_type="gemma3", + placeholder_metadata=MultimodalPlaceholderMetadata( + placeholder_map={"image": ""}, + placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, + )) class Gemma3VLM(PreTrainedModel): def __init__(self, model_config: ModelConfig[Gemma3Config]): diff --git a/tensorrt_llm/_torch/models/modeling_hyperclovax.py b/tensorrt_llm/_torch/models/modeling_hyperclovax.py index 56d56f244337..a05784b9d8d0 100644 --- a/tensorrt_llm/_torch/models/modeling_hyperclovax.py +++ b/tensorrt_llm/_torch/models/modeling_hyperclovax.py @@ -15,7 +15,9 @@ from tensorrt_llm.inputs.multimodal import MultimodalParams -from ...inputs import (ExtraProcessedInputs, InputProcessor, TextPrompt, +from ...inputs import (ExtraProcessedInputs, InputProcessor, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, TextPrompt, register_input_processor) from ...logger import logger from ...sampling_params import SamplingParams @@ -961,7 +963,23 @@ def forward(self, multimodal_params: List[MultimodalParams]): @register_auto_model("HCXVisionForCausalLM") -@register_input_processor(HCXVisionInputProcessor, model_type="hyperclovax_vlm") +@register_input_processor( + HCXVisionInputProcessor, + model_type="hyperclovax_vlm", + placeholder_metadata=MultimodalPlaceholderMetadata( + placeholder_map={ + "image": + ('\n<|im_start|>user (mime) \n' + '{"type": "image/jpeg", "filename": ""}<|im_end|>\n' + '<|im_start|>user (vector)\n<|dummy3|><|im_end|>\n' + '<|im_start|>image/aux\n' + '다음 중 ocr은 사진에서 검출된 글자이고, lens_keyword는 사진에서 추출된 ' + 'keyword와 bbox 위치입니다.bbox는 0~1 사이로 정규화된 [x1, y1, x2, y2]의 ' + '형태입니다. 참고하여 답변하세요. ' + '{"ocr": "", "lens_keywords": "", "lens_local_keywords": ""}') + }, + placeholder_placement=MultimodalPlaceholderPlacement.AFTER_TEXT, + )) class HCXVisionForCausalLM(PreTrainedModel): def __init__(self, model_config: ModelConfig): diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 6ec6557961e7..178cff1c2bf3 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/tensorrt_llm/_torch/models/modeling_llama.py @@ -20,7 +20,9 @@ from tensorrt_llm.lora_manager import HfLoraLoader from tensorrt_llm.models.convert_utils import split_matrix_tp -from ...inputs import (ExtraProcessedInputs, InputProcessor, TextPrompt, +from ...inputs import (ExtraProcessedInputs, InputProcessor, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, TextPrompt, register_input_processor) from ...sampling_params import SamplingParams from ..attention_backend import AttentionMetadata @@ -1168,7 +1170,13 @@ def __call__( @register_auto_model("Llama4ForConditionalGeneration") -@register_input_processor(Llama4InputProcessor, model_type="llama4") +@register_input_processor( + Llama4InputProcessor, + model_type="llama4", + placeholder_metadata=MultimodalPlaceholderMetadata( + placeholder_map={"image": "<|image|>"}, + placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, + )) class Llama4ForConditionalGeneration(SpecDecOneEngineForCausalLM[Llama4Model, Llama4Config]): diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py index ac26cb647366..8b52c7cf6392 100644 --- a/tensorrt_llm/_torch/models/modeling_llava_next.py +++ b/tensorrt_llm/_torch/models/modeling_llava_next.py @@ -14,7 +14,9 @@ from tensorrt_llm.inputs.multimodal import MultimodalParams -from ...inputs import (ExtraProcessedInputs, InputProcessor, TextPrompt, +from ...inputs import (ExtraProcessedInputs, InputProcessor, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, TextPrompt, register_input_processor) from ...llmapi.utils import download_hf_model from ...logger import logger @@ -263,7 +265,13 @@ def forward(self, multimodal_params: List[MultimodalParams]): @register_auto_model("LlavaNextForConditionalGeneration") -@register_input_processor(LlavaNextInputProcessor, model_type="llava_next") +@register_input_processor( + LlavaNextInputProcessor, + model_type="llava_next", + placeholder_metadata=MultimodalPlaceholderMetadata( + placeholder_map={"image": ""}, + placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, + )) class LlavaNextModel(PreTrainedModel): config_class = LlavaNextConfig diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index f10ea5368c96..cbf4d6722f0e 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -29,7 +29,9 @@ from tensorrt_llm._torch.speculative import SpecMetadata from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.inputs import (ExtraProcessedInputs, InputProcessor, - TextPrompt, register_input_processor) + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, TextPrompt, + register_input_processor) from tensorrt_llm.llmapi import SamplingParams from tensorrt_llm.logger import logger @@ -269,8 +271,20 @@ def __call__( @register_auto_model("Mistral3ForConditionalGeneration") -# The below informs the registry which input registry to create for this in `tensorrt_llm/llmapi/llm.py`. -@register_input_processor(Mistral3InputProcessor, model_type="mistral3") +@register_input_processor( + Mistral3InputProcessor, + model_type="mistral3", + placeholder_metadata=MultimodalPlaceholderMetadata( + placeholder_map={ + "image": "[IMG]", + }, + # NOTE: for mistral3 multimodal models, it does not strictly have to be before the text. + # Ref: https://github.com/mistralai/mistral-common/blob/039465db2bdc0486df36365c9bdb428188482a18/ + # src/mistral_common/tokens/tokenizers/base.py#L326 + # However, accuracy tests show that the model generates higher quality output when the image + # precedes the text (the relative difference can be as much as ~30% for both vLLM and TRT-LLM). + placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, + )) class Mistral3VLM(PreTrainedModel): """Mistral3VLM implementation for TRTLLM. diff --git a/tensorrt_llm/_torch/models/modeling_phi4mm.py b/tensorrt_llm/_torch/models/modeling_phi4mm.py index b5ad4f45203a..bf55965dd8b9 100644 --- a/tensorrt_llm/_torch/models/modeling_phi4mm.py +++ b/tensorrt_llm/_torch/models/modeling_phi4mm.py @@ -10,7 +10,9 @@ from PIL import Image from ...executor.request import LoRARequest -from ...inputs import (ExtraProcessedInputs, InputProcessor, TextPrompt, +from ...inputs import (ExtraProcessedInputs, InputProcessor, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, TextPrompt, register_input_processor) from ...logger import logger from ...lora_manager import LoraConfig @@ -138,7 +140,17 @@ def __call__( @register_auto_model("Phi4MMForCausalLM") -@register_input_processor(Phi4MMInputProcessor, model_type="phi4mm") +@register_input_processor( + Phi4MMInputProcessor, + model_type="phi4mm", + placeholder_metadata=MultimodalPlaceholderMetadata( + placeholder_map={ + "image": "<|image_{0}|>", + "audio": "<|audio_{0}|>", + }, + placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, + placeholders_separator="", + )) class Phi4MMForCausalLM(transformers.PreTrainedModel): _supports_flash_attn_2 = True diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index 03f15c37b42a..67e329e32696 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -12,7 +12,9 @@ from ..._utils import nvtx_range_debug from ...functional import RopeEmbeddingUtils, RotaryScalingType -from ...inputs import (ExtraProcessedInputs, InputProcessor, TextPrompt, +from ...inputs import (ExtraProcessedInputs, InputProcessor, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, TextPrompt, register_input_processor) from ...logger import logger from ...sampling_params import SamplingParams @@ -645,7 +647,16 @@ def forward( @register_auto_model("Qwen2VLForConditionalGeneration") -@register_input_processor(Qwen2VLInputProcessorBase, model_type="qwen2_vl") +@register_input_processor( + Qwen2VLInputProcessorBase, + model_type="qwen2_vl", + placeholder_metadata=MultimodalPlaceholderMetadata( + placeholder_map={ + "image": "<|vision_start|><|image_pad|><|vision_end|>", + "video": "<|vision_start|><|video_pad|><|vision_end|>" + }, + placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, + )) class Qwen2VLModel(Qwen2VLModelBase): def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, @@ -657,7 +668,14 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, @register_auto_model("Qwen2_5_VLForConditionalGeneration") -@register_input_processor(Qwen2VLInputProcessorBase, model_type="qwen2_5_vl") +@register_input_processor( + Qwen2VLInputProcessorBase, + model_type="qwen2_5_vl", + placeholder_metadata=MultimodalPlaceholderMetadata( + placeholder_map={ + "image": "<|vision_start|><|image_pad|><|vision_end|>", + "video": "<|vision_start|><|video_pad|><|vision_end|>" + })) class Qwen2_5_VLModel(Qwen2VLModelBase): def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, diff --git a/tensorrt_llm/_torch/models/modeling_vila.py b/tensorrt_llm/_torch/models/modeling_vila.py index 99820c1954c2..e69851cc2f5d 100644 --- a/tensorrt_llm/_torch/models/modeling_vila.py +++ b/tensorrt_llm/_torch/models/modeling_vila.py @@ -35,7 +35,9 @@ PreTrainedModel) from ..._utils import nvtx_range -from ...inputs import (ExtraProcessedInputs, InputProcessor, TextPrompt, +from ...inputs import (ExtraProcessedInputs, InputProcessor, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, TextPrompt, register_input_processor) from ...logger import logger from ...sampling_params import SamplingParams @@ -1118,7 +1120,16 @@ def __call__( @register_auto_model(VilaConfig.model_architecture) -@register_input_processor(VilaInputProcessor, model_type="llava_llama") +@register_input_processor( + VilaInputProcessor, + model_type="llava_llama", + placeholder_metadata=MultimodalPlaceholderMetadata( + placeholder_map={ + "image": "", + "video": "" + }, + placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, + )) class VilaModel(PreTrainedModel): config_class = VilaConfig diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 8b83c85d5178..49e5d167848a 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -13,6 +13,7 @@ from tensorrt_llm.bench.benchmark.utils.asynchronous import async_benchmark from tensorrt_llm.bench.benchmark.utils.processes import IterationWriter from tensorrt_llm.bench.build.build import get_model_config +from tensorrt_llm.tools.importlib_utils import import_custom_module_from_dir # isort: off from tensorrt_llm.bench.benchmark.utils.general import ( @@ -49,6 +50,16 @@ type=click.Choice(ALL_SUPPORTED_BACKENDS), default="pytorch", help="The backend to use when running benchmarking.") +@optgroup.option( + "--custom_module_dirs", + type=click.Path(exists=True, + readable=True, + path_type=Path, + resolve_path=True), + default=None, + multiple=True, + help="Paths to custom module directories to import.", +) @optgroup.option( "--extra_llm_api_options", type=str, @@ -98,6 +109,12 @@ required=False, help="Pass in a dataset file for parsing instead of stdin.", ) +@optgroup.option( + "--no_skip_tokenizer_init", + is_flag=True, + default=False, + help="Do not skip tokenizer initialization when loading the model.", +) @optgroup.option( "--eos_id", type=int, @@ -112,6 +129,18 @@ default=None, help="Modality of the multimodal requests.", ) +@optgroup.option( + "--image_data_format", + type=click.Choice(["pt", "pil"]), + default="pt", + help="Format of the image data for multimodal models.", +) +@optgroup.option( + "--data_device", + type=click.Choice(["cuda", "cpu"]), + default="cuda", + help="Device to load the multimodal data on.", +) @optgroup.option( "--max_input_len", type=int, @@ -256,7 +285,17 @@ def throughput_command( logger.info("Preparing to run throughput benchmark...") # Parameters from CLI # Model, experiment, and engine params + custom_module_dirs: list[Path] = params.pop("custom_module_dirs", []) + for custom_module_dir in custom_module_dirs: + try: + import_custom_module_from_dir(custom_module_dir) + except Exception as e: + logger.error( + f"Failed to import custom module from {custom_module_dir}: {e}") + raise e + dataset_path: Path = params.get("dataset") + no_skip_tokenizer_init: bool = params.get("no_skip_tokenizer_init", False) eos_id: int = params.get("eos_id") warmup: int = params.get("warmup") num_requests: int = params.get("num_requests") @@ -268,6 +307,8 @@ def throughput_command( backend: str = params.get("backend") modality: str = params.get("modality") max_input_len: int = params.get("max_input_len") + image_data_format: str = params.get("image_data_format", "pt") + data_device: str = params.get("data_device", "cpu") model_type = get_model_config(model, checkpoint_path).model_type # Reporting options @@ -280,7 +321,7 @@ def throughput_command( # Runtime kwargs and option tracking. kwargs = {} - # Initialize the HF tokenizer for the specified model. + # Initialize the HF tokenizer for the specified model. This is only used for data preparation. tokenizer = initialize_tokenizer(checkpoint_path) # Dataset Loading and Preparation @@ -292,6 +333,8 @@ def throughput_command( model_dir=checkpoint_path, model_type=model_type, modality=modality, + image_data_format=image_data_format, + data_device=data_device, max_input_seq_len_for_multimodal=max_input_len) metadata.dataset_path = dataset_path params["target_input_len"] = params.get( @@ -386,6 +429,7 @@ def ignore_trt_only_args(kwargs: dict): logger.info("Setting up throughput benchmark.") kwargs = kwargs | runtime_config.get_llm_args() kwargs['backend'] = backend + kwargs['skip_tokenizer_init'] = not no_skip_tokenizer_init if backend == "pytorch" and iteration_log is not None: kwargs["enable_iter_perf_stats"] = True diff --git a/tensorrt_llm/bench/utils/data.py b/tensorrt_llm/bench/utils/data.py index 6469d08739eb..080655c7271e 100644 --- a/tensorrt_llm/bench/utils/data.py +++ b/tensorrt_llm/bench/utils/data.py @@ -41,6 +41,8 @@ def create_dataset_from_stream( model_dir: str = None, model_type: str = None, modality: str = None, + image_data_format: str = "pt", + data_device: str = "cpu", max_input_seq_len_for_multimodal: int = 4096, ) -> Tuple[DatasetMetadata, List[InferenceRequest]]: """Generate metadata and a list of requests to drive benchmarking. @@ -130,7 +132,9 @@ def create_dataset_from_stream( model_type=model_type, modality=modality, prompts=prompts, - media=media_paths) # list of dicts + media=media_paths, # list of dicts + image_data_format=image_data_format, + device=data_device) all_isl = [] all_seq_len = [] diff --git a/tensorrt_llm/inputs/__init__.py b/tensorrt_llm/inputs/__init__.py index a20978cab430..f7e5ce97d7f2 100644 --- a/tensorrt_llm/inputs/__init__.py +++ b/tensorrt_llm/inputs/__init__.py @@ -1,16 +1,23 @@ from .data import PromptInputs, TextPrompt, TokensPrompt, prompt_inputs from .multimodal import MultimodalInput from .registry import (ExtraProcessedInputs, InputProcessor, - create_input_processor, create_input_processor_with_hash, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, create_input_processor, + create_input_processor_with_hash, register_input_processor) -from .utils import (ALL_SUPPORTED_MULTIMODAL_MODELS, ConversationMessage, - MultimodalData, MultimodalDataTracker, +from .utils import (ALL_SUPPORTED_AUDIO_MODELS, ALL_SUPPORTED_IMAGE_MODELS, + ALL_SUPPORTED_MULTIMODAL_MODELS, ALL_SUPPORTED_VIDEO_MODELS, + ConversationMessage, MultimodalData, MultimodalDataTracker, add_multimodal_placeholders, async_load_audio, async_load_image, async_load_video, default_multimodal_input_loader, encode_base64_content_from_url, load_image, load_video) __all__ = [ + "ALL_SUPPORTED_MULTIMODAL_MODELS", + "ALL_SUPPORTED_IMAGE_MODELS", + "ALL_SUPPORTED_VIDEO_MODELS", + "ALL_SUPPORTED_AUDIO_MODELS", "PromptInputs", "prompt_inputs", "TextPrompt", @@ -20,7 +27,8 @@ "create_input_processor_with_hash", "register_input_processor", "ExtraProcessedInputs", - "ALL_SUPPORTED_MULTIMODAL_MODELS", + "MultimodalPlaceholderMetadata", + "MultimodalPlaceholderPlacement", "ConversationMessage", "MultimodalDataTracker", "MultimodalData", diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 2b2211f41fcd..16fcf6fcd540 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -1,3 +1,5 @@ +import enum +from dataclasses import dataclass, field from typing import (Any, Callable, Dict, List, Optional, Protocol, Tuple, Type, TypeVar) @@ -10,7 +12,6 @@ from .multimodal import (MultimodalInput, apply_mm_hashes, default_hasher, find_mm_token_lengths, find_mm_token_positions, hexdigest_to_int32, validate_mm_inputs) -from .utils import ALL_SUPPORTED_MULTIMODAL_MODELS N = TypeVar("N", bound=Type[nn.Module]) @@ -108,6 +109,106 @@ def __call__( return token_ids, None +class MultimodalPlaceholderPlacement(enum.Enum): + INVALID = -1 + BEFORE_TEXT = 0 + AFTER_TEXT = 1 + + +@dataclass +class MultimodalPlaceholderMetadata: + placeholder_map: Dict[str, str] = field(default_factory=dict) + placeholder_placement: MultimodalPlaceholderPlacement = MultimodalPlaceholderPlacement.AFTER_TEXT + placeholders_separator: str = "\n" + + +class MultimodalPlaceholderRegistry: + + def __init__(self) -> None: + self._multimodal_placeholder_by_model_type: Dict[ + str, MultimodalPlaceholderMetadata] = {} + + def __str__(self) -> str: + s = "" + for model_type, placeholder_metadata in self._multimodal_placeholder_by_model_type.items( + ): + s += "-" * 100 + "\n" + s += f"Model type: {model_type}\n" + s += f"Placeholder map: {placeholder_metadata.placeholder_map}\n" + s += f"Placeholder placement: {placeholder_metadata.placeholder_placement}\n" + s += f"Placeholders separator: \"{placeholder_metadata.placeholders_separator}\"\n" + s += "-" * 80 + "\n" + return s + + def register(self, model_type: str, + placeholder_metadata: MultimodalPlaceholderMetadata): + self._multimodal_placeholder_by_model_type[ + model_type] = placeholder_metadata + + def is_valid(self, model_type: str, modality: str) -> bool: + return model_type in self._multimodal_placeholder_by_model_type and \ + modality in self._multimodal_placeholder_by_model_type[model_type].placeholder_map + + def get_placeholder_metadata( + self, model_type: str) -> MultimodalPlaceholderMetadata: + if model_type not in self._multimodal_placeholder_by_model_type: + raise ValueError( + f"Model type {model_type} is not registered in MultimodalPlaceholderRegistry" + ) + return self._multimodal_placeholder_by_model_type[model_type] + + def get_placeholder(self, model_type: str, modality: str) -> str: + if not self.is_valid(model_type, modality): + raise ValueError( + f"Model type '{model_type}' with modality '{modality}' is not registered." + ) + return self._multimodal_placeholder_by_model_type[ + model_type].placeholder_map[modality] + + def get_placeholder_placement( + self, model_type: str) -> MultimodalPlaceholderPlacement: + if model_type not in self._multimodal_placeholder_by_model_type: + raise ValueError(f"Model type '{model_type}' is not registered") + return self._multimodal_placeholder_by_model_type[ + model_type].placeholder_placement + + def get_placeholders_separator(self, model_type: str) -> str: + if model_type not in self._multimodal_placeholder_by_model_type: + raise ValueError(f"Model type '{model_type}' is not registered") + return self._multimodal_placeholder_by_model_type[ + model_type].placeholders_separator + + def get_registered_image_model_types(self) -> List[str]: + return [ + model_type + for model_type in self._multimodal_placeholder_by_model_type + if "image" in self. + _multimodal_placeholder_by_model_type[model_type].placeholder_map + ] + + def get_registered_video_model_types(self) -> List[str]: + return [ + model_type + for model_type in self._multimodal_placeholder_by_model_type + if "video" in self. + _multimodal_placeholder_by_model_type[model_type].placeholder_map + ] + + def get_registered_audio_model_types(self) -> List[str]: + return [ + model_type + for model_type in self._multimodal_placeholder_by_model_type + if "audio" in self. + _multimodal_placeholder_by_model_type[model_type].placeholder_map + ] + + def get_registered_model_types(self) -> List[str]: + return list(self._multimodal_placeholder_by_model_type.keys()) + + +MULTIMODAL_PLACEHOLDER_REGISTRY = MultimodalPlaceholderRegistry() + + class InputProcessorRegistry: def __init__(self) -> None: @@ -118,9 +219,10 @@ def __init__(self) -> None: INPUT_PROCESSOR_REGISTRY = InputProcessorRegistry() -def register_input_processor(processor_cls: Type[InputProcessor], - model_type: str, - out_of_tree: bool = False): +def register_input_processor( + processor_cls: Type[InputProcessor], + model_type: str, + placeholder_metadata: MultimodalPlaceholderMetadata = None): """ Register an input processor to a model class. NOTE: @@ -128,17 +230,18 @@ def register_input_processor(processor_cls: Type[InputProcessor], the model type only for that. 2. If this is used for other models in the future, this logic needs to be updated e.g. adding another version of this API without the model_type. - 3. If the model is not in the tree, user needs to set out_of_tree to True - to bypass the model type check and provide their own input preparation. """ def wrapper(model_cls: N) -> N: INPUT_PROCESSOR_REGISTRY._input_processors_cls_by_model_type[ model_cls] = processor_cls - if not out_of_tree: - assert model_type in ALL_SUPPORTED_MULTIMODAL_MODELS, \ - f"Model type {model_type} not in {ALL_SUPPORTED_MULTIMODAL_MODELS}.\n" \ - "Please see the tensorrt_llm/inputs/utils.py file for more information." + if placeholder_metadata is None: + raise ValueError( + f"A valid placeholder_metadata must be provided but got {placeholder_metadata}" + ) + + MULTIMODAL_PLACEHOLDER_REGISTRY.register(model_type, + placeholder_metadata) return model_cls @@ -192,41 +295,87 @@ def create_input_processor_with_hash( A wrapped processor that modifies prompts before processing. """ + def multimodal_hashing_process( + inputs: TextPrompt, sampling_params: SamplingParams + ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: + assert 'multi_modal_data' in inputs, "multi_modal_data must be provided for hashing support." + mm_data = inputs['multi_modal_data'] + num_mm_tokens = find_mm_token_lengths(mm_data, input_processor) + if len(num_mm_tokens) > 0: + mm_hashes = apply_mm_hashes(mm_data, hash_lib) + prompt_token_ids, extra_processed_inputs = input_processor( + inputs, sampling_params) + start_positions = find_mm_token_positions( + input_ids=prompt_token_ids, # token sequence + num_mm_tokens= + num_mm_tokens, # list of lengths of each chunk of visual tokens + vocab_size=input_processor.model_config.vocab_size, + ) + # flatten the hashes from dict to a single list + mm_hashes = [h for hashes in mm_hashes.values() for h in hashes] + validate_mm_inputs(prompt_token_ids, mm_hashes, start_positions, + num_mm_tokens) + mm_hashes_int32 = [hexdigest_to_int32(h) for h in mm_hashes + ] # nested list w/ multiple int32 per hash + + extra_processed_inputs[ + "multimodal_input"] = MultimodalInput.from_components( + mm_hashes_int32, start_positions, num_mm_tokens) + return prompt_token_ids, extra_processed_inputs + def input_processor_wrapper( inputs: TextPrompt, sampling_params: SamplingParams ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: - try: - assert 'multi_modal_data' in inputs, "multi_modal_data must be provided for hashing support." - mm_data = inputs['multi_modal_data'] - num_mm_tokens = find_mm_token_lengths(mm_data, input_processor) - if len(num_mm_tokens) > 0: - mm_hashes = apply_mm_hashes(mm_data, hash_lib) - prompt_token_ids, extra_processed_inputs = input_processor( + try_multimodal_hashing = False # only used for first time + use_multimodal_hashing = False # used for subsequent calls + modalities = list(set(inputs['multi_modal_data'].keys()) + ) if 'multi_modal_data' in inputs else [] + if len(modalities) > 0: + # NOTE: tensorrt_llm/inputs/multimodal.py:find_mm_token_lengths only supports image data for now + if len(modalities) == 1 and modalities[0] == "image": + # only try multimodal hashing if the inputs only contain image data + if hasattr(input_processor, "multimodal_hashing_supported"): + if input_processor.__getattribute__( + "multimodal_hashing_supported"): + use_multimodal_hashing = True + else: + try_multimodal_hashing = True + + if try_multimodal_hashing or use_multimodal_hashing: + try: + prompt_token_ids, extra_processed_inputs = multimodal_hashing_process( inputs, sampling_params) - start_positions = find_mm_token_positions( - input_ids=prompt_token_ids, # token sequence - num_mm_tokens= - num_mm_tokens, # list of lengths of each chunk of visual tokens - vocab_size=input_processor.model_config.vocab_size, - ) - # flatten the hashes from dict to a single list - mm_hashes = [h for hashes in mm_hashes.values() for h in hashes] - validate_mm_inputs(prompt_token_ids, mm_hashes, start_positions, - num_mm_tokens) - mm_hashes_int32 = [hexdigest_to_int32(h) for h in mm_hashes - ] # nested list w/ multiple int32 per hash - - extra_processed_inputs[ - "multimodal_input"] = MultimodalInput.from_components( - mm_hashes_int32, start_positions, num_mm_tokens) + if try_multimodal_hashing: + # if trying for first time, set the flag to True + input_processor.__setattr__("multimodal_hashing_supported", + True) return prompt_token_ids, extra_processed_inputs - else: + except Exception as e: + import traceback + traceback.print_exc() + logger.warning(f"Multimodal hashing failed: {e}.") + if try_multimodal_hashing: + # if trying for first time, fall back to basic input processor + # and set the flag to False so that we don't try again + input_processor.__setattr__("multimodal_hashing_supported", + False) + logger.warning("Falling back to basic input processor.") + try: + return input_processor(inputs, sampling_params) + except Exception as e2: + import traceback + traceback.print_exc() + logger.warning(f"Basic input processor failed: {e}.") + raise e2 + else: + raise e + else: + try: return input_processor(inputs, sampling_params) - except Exception as e: - # Fall back to basic input processor if multimodal processing fails - logger.warning( - f"Multimodal hashing failed: {e}. Falling back to basic input processor." - ) - return input_processor(inputs, sampling_params) + except Exception as e: + import traceback + traceback.print_exc() + logger.warning(f"Basic input processor failed: {e}.") + raise e return input_processor_wrapper diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index 7764436e2521..9f0e9f16706c 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -1,6 +1,5 @@ import asyncio import base64 -import enum import tempfile from collections import defaultdict from io import BytesIO @@ -18,6 +17,8 @@ from transformers import AutoProcessor, ProcessorMixin from transformers.utils import logging +from tensorrt_llm.inputs.registry import (MULTIMODAL_PLACEHOLDER_REGISTRY, + MultimodalPlaceholderPlacement) from tensorrt_llm.llmapi.llm_utils import ModelLoader from tensorrt_llm.llmapi.tokenizer import TokenizerBase, TransformersTokenizer @@ -209,60 +210,25 @@ def encode_base64_content_from_url(content_url: str) -> str: placeholder for the model needs to be added in retrieve_multimodal_placeholder(). """ -SUPPORTED_QWEN_MODEL_GROUP = ["qwen2_vl", "qwen2_5_vl"] -SUPPORTED_GEMMA_MODEL_GROUP = ["gemma3"] -SUPPORTED_LLAMA_MODEL_GROUP = ["mllama", "llama4"] -SUPPORTED_LLAVA_IMAGE_MODEL_GROUP = ["llava_llama", "llava_next"] -SUPPORTED_LLAVA_VIDEO_MODEL_GROUP = ["llava_llama"] -SUPPORTED_MISTRAL_IMAGE_MODEL_GROUP = ["mistral3"] -SUPPORTED_HYPERCLOVAX_MODEL_GROUP = ["hyperclovax_vlm"] -SUPPORTED_PHI_MODEL_GROUP = ["phi4mm"] +HF_CHAT_TEMPLATE_EXCEPTIONS = ["llava_llama"] +PLACEHOLDER_EXCEPTIONS = ["llava_next"] -ALL_SUPPORTED_IMAGE_MODELS = SUPPORTED_QWEN_MODEL_GROUP \ - + SUPPORTED_LLAMA_MODEL_GROUP \ - + SUPPORTED_LLAVA_IMAGE_MODEL_GROUP \ - + SUPPORTED_HYPERCLOVAX_MODEL_GROUP \ - + SUPPORTED_GEMMA_MODEL_GROUP \ - + SUPPORTED_MISTRAL_IMAGE_MODEL_GROUP \ - + SUPPORTED_PHI_MODEL_GROUP -ALL_SUPPORTED_VIDEO_MODELS = SUPPORTED_QWEN_MODEL_GROUP \ - + SUPPORTED_LLAVA_VIDEO_MODEL_GROUP +# Helpers to always get the latest supported multimodal model types from the registry +def ALL_SUPPORTED_MULTIMODAL_MODELS(): + return MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_model_types() -ALL_SUPPORTED_AUDIO_MODELS = SUPPORTED_PHI_MODEL_GROUP -ALL_SUPPORTED_MULTIMODAL_MODELS = list(set(ALL_SUPPORTED_IMAGE_MODELS) \ - | set(ALL_SUPPORTED_VIDEO_MODELS) \ - | set(ALL_SUPPORTED_AUDIO_MODELS)) +def ALL_SUPPORTED_IMAGE_MODELS(): + return MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_image_model_types() -HF_CHAT_TEMPLATE_EXCEPTIONS = ["llava_llama"] -PLACEHOLDER_EXCEPTIONS = ["llava_next"] + +def ALL_SUPPORTED_VIDEO_MODELS(): + return MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_video_model_types() -class MultimodalPlaceholderPlacement(enum.Enum): - INVALID = -1 - BEFORE_TEXT = 0 - AFTER_TEXT = 1 - - -PLACEHOLDER_PLACEMENT_MAP = { - "qwen2_vl": MultimodalPlaceholderPlacement.BEFORE_TEXT, - "qwen2_5_vl": MultimodalPlaceholderPlacement.BEFORE_TEXT, - "llava_llama": MultimodalPlaceholderPlacement.BEFORE_TEXT, - "llava_next": MultimodalPlaceholderPlacement.BEFORE_TEXT, - "llama4": MultimodalPlaceholderPlacement.BEFORE_TEXT, - "mllama": MultimodalPlaceholderPlacement.BEFORE_TEXT, - "hyperclovax_vlm": MultimodalPlaceholderPlacement.AFTER_TEXT, - "gemma3": MultimodalPlaceholderPlacement.BEFORE_TEXT, - # NOTE: for mistral3 multimodal models, it does not strictly have to be before the text. - # Ref: https://github.com/mistralai/mistral-common/blob/039465db2bdc0486df36365c9bdb428188482a18/ - # src/mistral_common/tokens/tokenizers/base.py#L326 - # However, accuracy tests show that the model generates higher quality output when the image - # precedes the text (the relative difference can be as much as ~30% for both vLLM and TRT-LLM). - "mistral3": MultimodalPlaceholderPlacement.BEFORE_TEXT, - "phi4mm": MultimodalPlaceholderPlacement.BEFORE_TEXT, -} -assert len(PLACEHOLDER_PLACEMENT_MAP) == len(ALL_SUPPORTED_MULTIMODAL_MODELS) +def ALL_SUPPORTED_AUDIO_MODELS(): + return MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_audio_model_types() def retrieve_multimodal_placeholder(model_type: str, modality: str, @@ -276,41 +242,16 @@ def retrieve_multimodal_placeholder(model_type: str, modality: str, current_count: The number of multimodal data already added. """ - - if modality == "image": - if model_type in SUPPORTED_QWEN_MODEL_GROUP: - return "<|vision_start|><|image_pad|><|vision_end|>" - elif model_type in SUPPORTED_LLAMA_MODEL_GROUP: - return "<|image|>" - elif model_type in SUPPORTED_LLAVA_IMAGE_MODEL_GROUP: - return "" - elif model_type in SUPPORTED_GEMMA_MODEL_GROUP: - return "" - elif model_type in SUPPORTED_HYPERCLOVAX_MODEL_GROUP: - return '\n<|im_start|>user (mime) \n{"type": "image/jpeg", "filename": ""}<|im_end|>\n' + \ - '<|im_start|>user (vector)\n<|dummy3|><|im_end|>\n' + \ - '<|im_start|>image/aux\n다음 중 ocr은 사진에서 검출된 글자이고, lens_keyword는 사진에서 추출된 keyword와 bbox 위치입니다.' + \ - 'bbox는 0~1 사이로 정규화된 [x1, y1, x2, y2]의 형태입니다. 참고하여 답변하세요. {"ocr": "", "lens_keywords": "", "lens_local_keywords": ""}' - elif model_type in SUPPORTED_MISTRAL_IMAGE_MODEL_GROUP: - # Ref: https://github.com/mistralai/mistral-common/blob/26a6bb3a07ee0b78a3808f2797f23e1d28514b93/ - # src/mistral_common/tokens/tokenizers/base.py#L60 - return "[IMG]" - elif model_type in SUPPORTED_PHI_MODEL_GROUP: - return f"<|image_{current_count}|>" - raise TypeError( - f"For image modality, only {ALL_SUPPORTED_IMAGE_MODELS} are supported but got {model_type}" - ) - elif modality == "video": - if model_type in SUPPORTED_QWEN_MODEL_GROUP: - return "<|vision_start|><|video_pad|><|vision_end|>" - elif model_type in SUPPORTED_LLAVA_VIDEO_MODEL_GROUP: - return "" - raise TypeError( - f"For video modality, only {ALL_SUPPORTED_VIDEO_MODELS} are supported but got {model_type}" - ) - elif modality == "audio": - if model_type in SUPPORTED_PHI_MODEL_GROUP: - return f"<|audio_{current_count}|>" + if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality): + """ + The placeholder is a string with a single placeholder for the current count. + - For example, if the placeholder is "<|image_{0}|>", and the current count is 1, + the placeholder will be "<|image_1|>". + - However, if the placeholder is "<|image|>", the current count would be ignored. + In this case, the placeholder would be "<|image|>". + """ + return MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder( + model_type, modality).format(current_count) raise TypeError(f"Unknown modality: {modality}") @@ -379,17 +320,15 @@ def add_multimodal_placeholders(model_type: str, text_prompt: str, for placeholder in mm_placeholder_counts: placeholders.extend([placeholder] * mm_placeholder_counts[placeholder]) parts = [] - match PLACEHOLDER_PLACEMENT_MAP[model_type]: + match MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder_placement(model_type): case MultimodalPlaceholderPlacement.BEFORE_TEXT: parts.extend(placeholders) parts.append(text_prompt) case MultimodalPlaceholderPlacement.AFTER_TEXT: parts.append(text_prompt) parts.extend(placeholders) - if model_type == "phi4mm": - return "".join(parts) - else: - return "\n".join(parts) + return MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholders_separator( + model_type).join(parts) def resolve_hf_chat_template( diff --git a/tensorrt_llm/tools/importlib_utils.py b/tensorrt_llm/tools/importlib_utils.py new file mode 100644 index 000000000000..866740ef3e0e --- /dev/null +++ b/tensorrt_llm/tools/importlib_utils.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +from pathlib import Path +from typing import Union + + +def import_custom_module_from_file(custom_module_path: Union[str, Path]): + """Import a custom module from a single file. + + Args: + custom_module_path (Union[str, Path]): The path to the custom module file. + + Returns: + The imported module object. + + Raises: + ImportError: If the module cannot be imported. + """ + if isinstance(custom_module_path, str): + custom_module_path = Path(custom_module_path) + print(f"Importing custom module from file: {custom_module_path}") + + # Import single Python file + module = None + spec = importlib.util.spec_from_file_location(custom_module_path.stem, + str(custom_module_path)) + if spec is not None: + module = importlib.util.module_from_spec(spec) + if spec.loader is not None: + spec.loader.exec_module(module) + print( + f"Successfully imported custom module from file: {custom_module_path}" + ) + else: + raise ImportError( + f"Failed to import custom module from {custom_module_path}") + else: + raise ImportError( + f"Failed to import custom module from {custom_module_path}") + return module + + +def import_custom_module_from_dir(custom_module_path: Union[str, Path]): + """Import a custom module from a directory. + + Args: + custom_module_path (Union[str, Path]): The path to the custom module directory. + + Returns: + The imported module object. + + Raises: + ImportError: If the module cannot be imported. + + Note: + This function will add the parent directory of the custom module directory to sys.path. + This is useful for importing modules that are not in the current working directory. + """ + if isinstance(custom_module_path, str): + custom_module_path = Path(custom_module_path) + print(f"Importing custom module from directory: {custom_module_path}") + + # Import directory as a package + # Add the parent directory to sys.path so we can import the package + import sys + parent_dir = str(custom_module_path.parent) + path_added = False + if parent_dir not in sys.path: + sys.path.insert(0, parent_dir) + path_added = True + + # Import the package + module = None + package_name = custom_module_path.name + try: + module = importlib.import_module(package_name) + print( + f"Successfully imported custom module from directory: {custom_module_path}" + ) + except ImportError as e: + raise ImportError( + f"Failed to import package {package_name} from {custom_module_path}: {e}" + ) + finally: + # Clean up sys.path if we added to it + if path_added and parent_dir in sys.path: + sys.path.remove(parent_dir) + return module + + +def import_custom_module(custom_module_path: Union[str, Path]): + """Import a custom module from a file or directory. + + Args: + custom_module_path (Union[str, Path]): The path to the custom module file or directory. + + Returns: + The imported module object. + + Raises: + ImportError: If the module cannot be imported. + FileNotFoundError: If the custom module path does not exist. + """ + if isinstance(custom_module_path, str): + custom_module_path = Path(custom_module_path) + print(f"Importing custom module from: {custom_module_path}") + + if custom_module_path.exists(): + if custom_module_path.is_file(): + return import_custom_module_from_file(custom_module_path) + elif custom_module_path.is_dir(): + return import_custom_module_from_dir(custom_module_path) + else: + raise FileNotFoundError( + f"Custom module path {custom_module_path} is neither a file nor a directory." + ) + else: + raise FileNotFoundError( + f"Custom module path {custom_module_path} does not exist.") + return None From 12c8156278843093c8dfc0d1e3f4765189fc3a99 Mon Sep 17 00:00:00 2001 From: Rakib Hasan Date: Fri, 8 Aug 2025 00:18:17 +0000 Subject: [PATCH 2/3] polish: adding docs, tests and addressing review comments Signed-off-by: Rakib Hasan --- .../llm-api/out_of_tree_example/readme.md | 52 +++++++++ examples/llm-api/quickstart_multimodal.py | 4 +- tensorrt_llm/bench/benchmark/throughput.py | 4 + tensorrt_llm/inputs/registry.py | 78 ++++++++----- tensorrt_llm/tools/importlib_utils.py | 18 ++- .../others/test_multimodal_registry.py | 106 ++++++++++++++++++ 6 files changed, 223 insertions(+), 39 deletions(-) create mode 100644 examples/llm-api/out_of_tree_example/readme.md create mode 100644 tests/unittest/others/test_multimodal_registry.py diff --git a/examples/llm-api/out_of_tree_example/readme.md b/examples/llm-api/out_of_tree_example/readme.md new file mode 100644 index 000000000000..1b26ea3cd67c --- /dev/null +++ b/examples/llm-api/out_of_tree_example/readme.md @@ -0,0 +1,52 @@ +# Out-of-tree Model Development +The file `modeling_opt.py` shows an example of how a custom model can be defined using TRT-LLM APIs without modifying the source code of TRT-LLM. + +The file `main.py` shows how to run inference for such custom models using the LLM API. + + +## Out-of-tree Multimodal Models + +For multimodal models, TRT-LLM provides `quickstart_multimodal.py` to quickly run a multimodal model that is defined within TRT-LLM. `trtllm-bench` can be used for benchmarking such models. +However, the following sections describe how to use those tools for out-of-tree models. + +### Pre-requisite +To use an out-of-tree model with the quickstart example and trtllm-bench, you need to prepare the model definition files similar to a python module. +Consider the following file structure as an example: +``` +modeling_custom_phi +|-- __init__.py +|-- configuration.py +|-- modeling_custom_phi.py +|-- encoder + |-- __init__.py + |-- configuration.py + |-- modeling_encoder.py +```` +The files `__init__.py` should be populated with the right imports for the custom model. For example, the `modeling_custom_phi/__init__.py` can contain something like: +``` +from .modeling_custom_phi import MyVLMForConditionalGeneration +from . import encoder +``` + +### Quickstart Example + +Once the model definition files are prepared as a python module (as described above), you can use the `--custom_module_dirs` flag in `quickstart_multimodal.py` to load your model and run inference. + +``` +python3 quickstart_multimodal.py --model_dir ./model_ckpt --modality image --max_tokens 10 --prompt "Describe the image." --media ./demo_lower.png --image_format pil --custom_module_dirs ../modeling_custom_phi +``` + +### Benchmarking + +Similar to the quickstart example, you can use the same CLI argument with `trtllm-bench` to benchmark a custom model. + +Prepare the dataset: +``` +python ./benchmarks/cpp/prepare_dataset.py --tokenizer ./model_ckpt --stdout dataset --dataset-name lmms-lab/MMMU --dataset-split test --dataset-image-key image --dataset-prompt-key "question" --num-requests 100 --output-len-dist 128,5 > mm_data.jsonl +``` + + +Run the benchmark: +``` +trtllm-bench --model ./model_ckpt --model_path ./model_ckpt throughput --dataset mm_data.jsonl --backend pytorch --num_requests 100 --max_batch_size 4 --modality image --streaming --custom_module_dirs ../modeling_custom_phi +``` diff --git a/examples/llm-api/quickstart_multimodal.py b/examples/llm-api/quickstart_multimodal.py index c5f15a4ddce7..25401e1c9558 100644 --- a/examples/llm-api/quickstart_multimodal.py +++ b/examples/llm-api/quickstart_multimodal.py @@ -84,7 +84,7 @@ def add_multimodal_args(parser): "--model_type", type=str, choices=MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_model_types(), - help="Model type.") + help="Model type as specified in the HuggingFace model config.") parser.add_argument("--modality", type=str, choices=[ @@ -92,7 +92,7 @@ def add_multimodal_args(parser): "multiple_image", "mixture_text_image" ], default="image", - help="Media type.") + help="Media type being used for inference.") parser.add_argument("--media", type=str, nargs="+", diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 49e5d167848a..f3072ecdecfa 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -109,6 +109,10 @@ required=False, help="Pass in a dataset file for parsing instead of stdin.", ) +# For text models, tokenizer initialization is not needed when loading the model since the dataset is already tokenized. +# For this reason, we skip tokenizer initialization by default. +# However, for VLM models, tokenizer initialization is needed inside the model since the dataset contains texts and +# raw media data. We cannot skip tokenizer initialization in this case. @optgroup.option( "--no_skip_tokenizer_init", is_flag=True, diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 16fcf6fcd540..e75ba7390878 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -33,6 +33,7 @@ class InputProcessor(Protocol): model_path: any model_config: any tokenizer: any + multimodal_hashing_supported: Optional[bool] = None def __call__( self, inputs: TextPrompt, sampling_params: SamplingParams @@ -51,6 +52,7 @@ def __init__(self, self.tokenizer = tokenizer self.model_config = model_config self.model_path = model_path + self.multimodal_hashing_supported = None def __call__( self, inputs: TextPrompt, sampling_params: SamplingParams @@ -110,19 +112,37 @@ def __call__( class MultimodalPlaceholderPlacement(enum.Enum): + """ + The placement of the multimodal placeholder in the prompt. Valid values are: + - BEFORE_TEXT: the placeholders are placed before the text prompt. + - AFTER_TEXT: the placeholders are placed after the text prompt. + """ INVALID = -1 BEFORE_TEXT = 0 AFTER_TEXT = 1 -@dataclass +@dataclass(frozen=True) class MultimodalPlaceholderMetadata: + """ + Metadata for the multimodal placeholder. It has 3 components: + - placeholder_map: + A mapping from modality to placeholder string. + Modality can be "image", "video", "audio", etc. + - placeholder_placement: + The placement of the placeholders, e.g. before or after the text prompt. + - placeholders_separator: + The separator between the placeholders, e.g. some models use "\n" to separate the placeholders. + """ placeholder_map: Dict[str, str] = field(default_factory=dict) placeholder_placement: MultimodalPlaceholderPlacement = MultimodalPlaceholderPlacement.AFTER_TEXT placeholders_separator: str = "\n" class MultimodalPlaceholderRegistry: + """ + Registry for the multimodal models to keep track of the placeholder information. + """ def __init__(self) -> None: self._multimodal_placeholder_by_model_type: Dict[ @@ -140,11 +160,17 @@ def __str__(self) -> str: s += "-" * 80 + "\n" return s - def register(self, model_type: str, - placeholder_metadata: MultimodalPlaceholderMetadata): + def set_placeholder_metadata( + self, model_type: str, + placeholder_metadata: MultimodalPlaceholderMetadata): self._multimodal_placeholder_by_model_type[ model_type] = placeholder_metadata + def remove_placeholder_metadata(self, model_type: str): + if model_type not in self._multimodal_placeholder_by_model_type: + raise ValueError(f"Model type '{model_type}' is not registered") + del self._multimodal_placeholder_by_model_type[model_type] + def is_valid(self, model_type: str, modality: str) -> bool: return model_type in self._multimodal_placeholder_by_model_type and \ modality in self._multimodal_placeholder_by_model_type[model_type].placeholder_map @@ -178,32 +204,29 @@ def get_placeholders_separator(self, model_type: str) -> str: return self._multimodal_placeholder_by_model_type[ model_type].placeholders_separator - def get_registered_image_model_types(self) -> List[str]: - return [ + def get_registered_image_model_types(self) -> Tuple[str, ...]: + return ( model_type for model_type in self._multimodal_placeholder_by_model_type if "image" in self. - _multimodal_placeholder_by_model_type[model_type].placeholder_map - ] + _multimodal_placeholder_by_model_type[model_type].placeholder_map) - def get_registered_video_model_types(self) -> List[str]: - return [ + def get_registered_video_model_types(self) -> Tuple[str, ...]: + return ( model_type for model_type in self._multimodal_placeholder_by_model_type if "video" in self. - _multimodal_placeholder_by_model_type[model_type].placeholder_map - ] + _multimodal_placeholder_by_model_type[model_type].placeholder_map) - def get_registered_audio_model_types(self) -> List[str]: - return [ + def get_registered_audio_model_types(self) -> Tuple[str, ...]: + return ( model_type for model_type in self._multimodal_placeholder_by_model_type if "audio" in self. - _multimodal_placeholder_by_model_type[model_type].placeholder_map - ] + _multimodal_placeholder_by_model_type[model_type].placeholder_map) - def get_registered_model_types(self) -> List[str]: - return list(self._multimodal_placeholder_by_model_type.keys()) + def get_registered_model_types(self) -> Tuple[str, ...]: + return tuple(self._multimodal_placeholder_by_model_type.keys()) MULTIMODAL_PLACEHOLDER_REGISTRY = MultimodalPlaceholderRegistry() @@ -240,8 +263,8 @@ def wrapper(model_cls: N) -> N: f"A valid placeholder_metadata must be provided but got {placeholder_metadata}" ) - MULTIMODAL_PLACEHOLDER_REGISTRY.register(model_type, - placeholder_metadata) + MULTIMODAL_PLACEHOLDER_REGISTRY.set_placeholder_metadata( + model_type, placeholder_metadata) return model_cls @@ -298,6 +321,9 @@ def create_input_processor_with_hash( def multimodal_hashing_process( inputs: TextPrompt, sampling_params: SamplingParams ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: + """ + Process the multinmodal hashing for media tokens if possible. + """ assert 'multi_modal_data' in inputs, "multi_modal_data must be provided for hashing support." mm_data = inputs['multi_modal_data'] num_mm_tokens = find_mm_token_lengths(mm_data, input_processor) @@ -322,6 +348,7 @@ def multimodal_hashing_process( "multimodal_input"] = MultimodalInput.from_components( mm_hashes_int32, start_positions, num_mm_tokens) return prompt_token_ids, extra_processed_inputs + return [], None def input_processor_wrapper( inputs: TextPrompt, sampling_params: SamplingParams @@ -334,11 +361,10 @@ def input_processor_wrapper( # NOTE: tensorrt_llm/inputs/multimodal.py:find_mm_token_lengths only supports image data for now if len(modalities) == 1 and modalities[0] == "image": # only try multimodal hashing if the inputs only contain image data - if hasattr(input_processor, "multimodal_hashing_supported"): - if input_processor.__getattribute__( - "multimodal_hashing_supported"): - use_multimodal_hashing = True + if input_processor.multimodal_hashing_supported is not None: + use_multimodal_hashing = input_processor.multimodal_hashing_supported else: + # we need to try the multimodal hashing for the first time to determine if it is supported try_multimodal_hashing = True if try_multimodal_hashing or use_multimodal_hashing: @@ -347,8 +373,7 @@ def input_processor_wrapper( inputs, sampling_params) if try_multimodal_hashing: # if trying for first time, set the flag to True - input_processor.__setattr__("multimodal_hashing_supported", - True) + input_processor.multimodal_hashing_supported = True return prompt_token_ids, extra_processed_inputs except Exception as e: import traceback @@ -357,8 +382,7 @@ def input_processor_wrapper( if try_multimodal_hashing: # if trying for first time, fall back to basic input processor # and set the flag to False so that we don't try again - input_processor.__setattr__("multimodal_hashing_supported", - False) + input_processor.multimodal_hashing_supported = False logger.warning("Falling back to basic input processor.") try: return input_processor(inputs, sampling_params) diff --git a/tensorrt_llm/tools/importlib_utils.py b/tensorrt_llm/tools/importlib_utils.py index 866740ef3e0e..281258612e84 100644 --- a/tensorrt_llm/tools/importlib_utils.py +++ b/tensorrt_llm/tools/importlib_utils.py @@ -15,10 +15,12 @@ import importlib.util from pathlib import Path -from typing import Union +from types import ModuleType +from typing import Optional, Union -def import_custom_module_from_file(custom_module_path: Union[str, Path]): +def import_custom_module_from_file( + custom_module_path: Union[str, Path]) -> Optional[ModuleType]: """Import a custom module from a single file. Args: @@ -54,7 +56,8 @@ def import_custom_module_from_file(custom_module_path: Union[str, Path]): return module -def import_custom_module_from_dir(custom_module_path: Union[str, Path]): +def import_custom_module_from_dir( + custom_module_path: Union[str, Path]) -> Optional[ModuleType]: """Import a custom module from a directory. Args: @@ -78,10 +81,8 @@ def import_custom_module_from_dir(custom_module_path: Union[str, Path]): # Add the parent directory to sys.path so we can import the package import sys parent_dir = str(custom_module_path.parent) - path_added = False if parent_dir not in sys.path: sys.path.insert(0, parent_dir) - path_added = True # Import the package module = None @@ -95,14 +96,11 @@ def import_custom_module_from_dir(custom_module_path: Union[str, Path]): raise ImportError( f"Failed to import package {package_name} from {custom_module_path}: {e}" ) - finally: - # Clean up sys.path if we added to it - if path_added and parent_dir in sys.path: - sys.path.remove(parent_dir) return module -def import_custom_module(custom_module_path: Union[str, Path]): +def import_custom_module( + custom_module_path: Union[str, Path]) -> Optional[ModuleType]: """Import a custom module from a file or directory. Args: diff --git a/tests/unittest/others/test_multimodal_registry.py b/tests/unittest/others/test_multimodal_registry.py new file mode 100644 index 000000000000..a6eb5345e0c9 --- /dev/null +++ b/tests/unittest/others/test_multimodal_registry.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import unittest + +from tensorrt_llm.inputs.registry import (MULTIMODAL_PLACEHOLDER_REGISTRY, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement) + + +class TestMultimodalPlaceholderRegistry(unittest.TestCase): + + def setUp(self): + self.model_type = "test_model_type" + self.placeholder_metadata = MultimodalPlaceholderMetadata( + placeholder_map={ + "image": "IMAGE_PLACEHOLDER", + "video": "VIDEO_PLACEHOLDER" + }, + placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, + placeholders_separator="\n") + + def test_new_registration(self): + MULTIMODAL_PLACEHOLDER_REGISTRY.set_placeholder_metadata( + self.model_type, self.placeholder_metadata) + self.assertEqual( + MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder_metadata( + self.model_type), self.placeholder_metadata) + MULTIMODAL_PLACEHOLDER_REGISTRY.remove_placeholder_metadata( + self.model_type) + + def test_registered_model_types(self): + pre_reg_model_types = list( + MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_model_types()) + + # register the model type + MULTIMODAL_PLACEHOLDER_REGISTRY.set_placeholder_metadata( + self.model_type, self.placeholder_metadata) + + post_reg_model_types = list( + MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_model_types()) + self.assertEqual( + len(pre_reg_model_types) + 1, len(post_reg_model_types)) + self.assertIn(self.model_type, post_reg_model_types) + + MULTIMODAL_PLACEHOLDER_REGISTRY.remove_placeholder_metadata( + self.model_type) + + def test_validity(self): + MULTIMODAL_PLACEHOLDER_REGISTRY.set_placeholder_metadata( + self.model_type, self.placeholder_metadata) + self.assertTrue( + MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(self.model_type, "image")) + self.assertTrue( + MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(self.model_type, "video")) + self.assertFalse( + MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(self.model_type, "audio")) + + MULTIMODAL_PLACEHOLDER_REGISTRY.remove_placeholder_metadata( + self.model_type) + + def test_model_types_per_modality(self): + pre_reg_image_model_types = list( + MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_image_model_types()) + pre_reg_video_model_types = list( + MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_video_model_types()) + pre_reg_audio_model_types = list( + MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_audio_model_types()) + + # register the model type for image and video + MULTIMODAL_PLACEHOLDER_REGISTRY.set_placeholder_metadata( + self.model_type, self.placeholder_metadata) + + post_reg_image_model_types = list( + MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_image_model_types()) + post_reg_video_model_types = list( + MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_video_model_types()) + post_reg_audio_model_types = list( + MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_audio_model_types()) + self.assertEqual( + len(pre_reg_image_model_types) + 1, len(post_reg_image_model_types)) + self.assertEqual( + len(pre_reg_video_model_types) + 1, len(post_reg_video_model_types)) + self.assertEqual(len(pre_reg_audio_model_types), + len(post_reg_audio_model_types)) + self.assertIn(self.model_type, post_reg_image_model_types) + self.assertIn(self.model_type, post_reg_video_model_types) + self.assertNotIn(self.model_type, post_reg_audio_model_types) + + MULTIMODAL_PLACEHOLDER_REGISTRY.remove_placeholder_metadata( + self.model_type) + + +if __name__ == "__main__": + unittest.main() From 8295b85b5ff62a114354d72b51d8886b1d5f66ff Mon Sep 17 00:00:00 2001 From: Rakib Hasan Date: Tue, 12 Aug 2025 00:28:34 +0000 Subject: [PATCH 3/3] rebasing and removing the now redundant import of torch models Signed-off-by: Rakib Hasan --- tensorrt_llm/__init__.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index 35c4ed485e5b..5cffa985460a 100644 --- a/tensorrt_llm/__init__.py +++ b/tensorrt_llm/__init__.py @@ -63,9 +63,6 @@ def _add_trt_llm_dll_directory(): from .sampling_params import SamplingParams from .version import __version__ -# Lazy import to avoid circular dependency on lora_manager -import tensorrt_llm._torch.models as torch_models # isort:skip - __all__ = [ 'AutoConfig', 'AutoModelForCausalLM',