diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index 5cffa985460a..d545b785b954 100644 --- a/tensorrt_llm/__init__.py +++ b/tensorrt_llm/__init__.py @@ -51,7 +51,7 @@ def _add_trt_llm_dll_directory(): from .builder import BuildConfig, Builder, BuilderConfig, build from .disaggregated_params import DisaggregatedParams from .functional import Tensor, constant -from .llmapi import LLM, LlmArgs +from .llmapi import LLM, MultimodalEncoder from .llmapi.llm_args import LlmArgs, TorchLlmArgs, TrtLlmArgs from .logger import logger from .mapping import Mapping @@ -105,6 +105,7 @@ def _add_trt_llm_dll_directory(): 'quantization', 'tools', 'LLM', + 'MultimodalEncoder', 'LlmArgs', 'TorchLlmArgs', 'TrtLlmArgs', diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 7e310f934acc..27b47acea390 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -103,6 +103,9 @@ class ModelConfig(Generic[TConfig]): _frozen: bool = field(default=False, init=False, repr=False) + # If true, ONLY the vision encoder part of the full model is loaded/executed. + mm_encoder_only: bool = False + def __setattr__(self, key, value): """ Prevent modification of frozen instance attributes. @@ -122,7 +125,8 @@ def __post_init__(self): if self.pretrained_config and hasattr(self.pretrained_config, "architectures"): self.is_generation = self.is_generation_model( - self.pretrained_config.architectures) + self.pretrained_config.architectures, + mm_encoder_only=self.mm_encoder_only) def get_all_reduce_strategy(strategy: str = "AUTO"): maps = { @@ -181,12 +185,15 @@ def get_quant_config(self, name: Optional[str] = None) -> QuantConfig: raise ValueError(f'quant config of {name} is not found') @staticmethod - def is_generation_model(model_architectures: Optional[List[str]]) -> bool: + def is_generation_model(model_architectures: Optional[List[str]], + mm_encoder_only: bool = False) -> bool: if model_architectures is None: logger.warning( "Model architectures is None, default to is_generation_model=True" ) return True + if mm_encoder_only: + return False return model_architectures[0] not in [ "BertForSequenceClassification", "Qwen2ForProcessRewardModel", "Qwen2ForRewardModel", "LlamaForTextEmbedding" diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py index 17bb6bab55cf..cd73919ca2cf 100644 --- a/tensorrt_llm/_torch/models/modeling_auto.py +++ b/tensorrt_llm/_torch/models/modeling_auto.py @@ -2,8 +2,9 @@ from ..model_config import ModelConfig from ..utils import model_extra_attrs -from .modeling_utils import (MODEL_CLASS_MAPPING, DecoderModelForCausalLM, - TConfig, TModel) +from .modeling_utils import (MODEL_CLASS_MAPPING, + MODEL_CLASS_VISION_ENCODER_MAPPING, + DecoderModelForCausalLM, TConfig, TModel) class AutoModelForCausalLM(Generic[TModel, TConfig]): @@ -13,6 +14,16 @@ def from_config( config: ModelConfig[TConfig], ) -> DecoderModelForCausalLM[TModel, TConfig]: model_arch = config.pretrained_config.architectures[0] + if config.mm_encoder_only: + vision_encoder_info = MODEL_CLASS_VISION_ENCODER_MAPPING.get( + model_arch) + if vision_encoder_info is None: + raise ValueError( + f"Unknown architecture for AutoModelForMultimodalEncoder: {model_arch}" + ) + vision_encoder_cls, vlm_base_model = vision_encoder_info + return vision_encoder_cls(config, vlm_base_model) + # Hack to detect eagle3 checkpoints. TODO: should we provide # our own checkpoints with the correct arch? It would let us # avoid nasty stuff like this. diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py index 8b52c7cf6392..282ec0c5cc7c 100644 --- a/tensorrt_llm/_torch/models/modeling_llava_next.py +++ b/tensorrt_llm/_torch/models/modeling_llava_next.py @@ -1,6 +1,6 @@ import copy import os -from typing import List, Optional, Tuple +from typing import Dict, List, Optional, Tuple, Union import numpy as np import torch @@ -26,7 +26,8 @@ from .modeling_auto import AutoModelForCausalLM from .modeling_clip import CLIPVisionModel from .modeling_multimodal_utils import fuse_input_embeds -from .modeling_utils import ModelConfig, filter_weights, register_auto_model +from .modeling_utils import (filter_weights, register_auto_model, + register_vision_encoder) DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1' @@ -53,6 +54,140 @@ def __init__(self, self.image_token_index = model_config.image_token_index self.vocab_size = model_config.vocab_size + self.config = model_config.vision_config + + def get_num_tokens_per_image( + self, + *, + image_width: int, + image_height: int, + ) -> int: + image_size = (image_height, image_width) + num_image_tokens = self.processor._get_num_multimodal_tokens( + [image_size])["num_image_tokens"][0] + return num_image_tokens + + def _postprocess( + self, input_ids: torch.Tensor, mm_features: Union[torch.Tensor, + List[torch.Tensor]] + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Define model specific variables here before shared logic + mm_tokens = torch.tensor([self.model_config.image_token_index + ]).to(input_ids.device) + model_hidden_size = self.model_config.text_config.hidden_size + start_len = end_len = 0 # for llava, need not append start/end token around each image token + # End model specific variables + + ## find mm token positions in input_ids + mm_token_positions = torch.where(torch.isin(input_ids, mm_tokens))[0] + num_medias = num_mm_tokens = len(mm_token_positions) + if num_medias > 1 and isinstance(mm_features, torch.Tensor): + mm_features = list( + mm_features.split(mm_features.shape[0] // num_medias)) + + if isinstance(mm_features, torch.Tensor): + # 1 prompt + 1 media + # "split" means what a single mm_token in the input_ids should represent + # image: one split --> one frame + # video: one split --> N frames + num_frames, mm_feature_length, mm_hidden_dim = mm_features.shape + mm_lengths_per_split = [mm_feature_length * num_frames] + mm_lengths_per_frame = [mm_feature_length] + elif isinstance(mm_features, list): + # 1 prompt + N media + num_frames = len(mm_features) if mm_features[0].dim() == 2 else sum( + [f.shape[0] for f in mm_features]) + mm_lengths_per_split = [ + f.shape[0] if f.dim() == 2 else f.shape[0] * f.shape[1] + for f in mm_features + ] + mm_lengths_per_frame = [ + f.shape[0] if f.dim() == 2 else f.shape[1] for f in mm_features + ] + mm_hidden_dim = mm_features[0].shape[-1] + mm_features = torch.cat(mm_features, dim=0) + else: + raise ValueError( + f"Invalid multimodal features type: {type(mm_features)}") + mm_total_length = sum(mm_lengths_per_split) + assert mm_hidden_dim == model_hidden_size, "Multimodal embedding_dim must match model hidden_size" + + ## split input_ids into segments by isolating mm tokens + mm_split_positions = torch.cat( + [mm_token_positions, mm_token_positions + 1]).unique() + input_ids_splits = list(input_ids.tensor_split(mm_split_positions.cpu( + ))) # len(input_ids_splits) = num_segments after mm tokens are isolated + mm_ids_splits = list( + torch.arange(self.vocab_size, + self.vocab_size + mm_total_length, + device=input_ids.device).split(mm_lengths_per_split) + ) # len(mm_ids_splits) = num_mm_segments + + for i, mm_ids in enumerate(mm_ids_splits): + mm_ids = mm_ids.reshape(-1, mm_lengths_per_frame[i]) + mm_ids_splits[i] = mm_ids.flatten() + + ## replace mm token ids with the expanded out-of-vocab ids + mm_split_idx = 0 + for i, split in enumerate(input_ids_splits): + if torch.isin(split, mm_tokens).any().item(): + input_ids_splits[i] = mm_ids_splits[mm_split_idx] + mm_split_idx += 1 + assert mm_split_idx == len( + mm_ids_splits), "All mm_ids_splits should be consumed" + + ## concat text & mm input_ids, wrap mm feature in prompt tuning config + fused_input_ids = torch.cat(input_ids_splits).to( + device=input_ids.device) + fused_length = len(input_ids) + mm_total_length + num_frames * ( + start_len + end_len) - num_medias + assert len( + fused_input_ids + ) == fused_length, f"Fused input_ids length {len(fused_input_ids)} should match the sum of text and multimodal embedding lengths {fused_length}" + + # [num_frames, feature_length, hidden_dim] -> [num_frames * feature_length, hidden_dim] + mm_features = mm_features.view(-1, mm_features.shape[-1]) + return fused_input_ids, mm_features + + def attach_multimodal_embeddings( + self, inputs: TextPrompt, + multimodal_embedding: Dict[str, List[torch.Tensor]], + sampling_params: SamplingParams + ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: + """ + Attach pre-processed multimodal embeddings into text token stream for LlavaNext model. + This method skips vision processing and works with externally provided embeddings. + It replaces/expands image placeholders in the text with appropriate tokens and prepares + the embeddings for model forward pass. + Args: + inputs: Text prompt containing image placeholders + multimodal_embedding: Dictionary containing pre-processed image embedding data + Returns: + Tuple of (token_ids, extra_processed_inputs) where: + - token_ids: List of processed token IDs with image placeholders + - extra_processed_inputs: Optional dictionary containing multimodal embeddings + """ + text_prompt = inputs.get("prompt") + if not text_prompt: + raise ValueError("Text prompt is required but not provided") + + if not isinstance(multimodal_embedding, dict): + raise ValueError("multimodal_embedding must be a dictionary") + + if 'image' not in multimodal_embedding: + raise ValueError( + "Only image modality is supported for external multimodal embedding" + ) + + input_ids = self.tokenizer(text_prompt, + return_tensors="pt").input_ids[0] + mm_features = multimodal_embedding['image'] + fused_input_ids, mm_features = self._postprocess(input_ids, mm_features) + multimodal_data = {} + multimodal_data["multimodal_embedding"] = mm_features + return fused_input_ids.to(torch.int32).tolist(), { + "multimodal_data": multimodal_data + } @torch.inference_mode() def __call__( @@ -92,6 +227,7 @@ class LlavaNextVisionModel(nn.Module): def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs) -> None: super().__init__() + self.model_config = model_config self.pretrained_config = model_config.pretrained_config self.device = f"cuda:{model_config.mapping.rank}" model_path = self.pretrained_config._name_or_path @@ -134,7 +270,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, self.vision_tower = hf_vision_tower.to(self.device) else: vision_model_config = ModelConfig( - pretrained_config=model_config.pretrained_config.vision_config, + pretrained_config=self.pretrained_config.vision_config, attn_backend="TRTLLM") self.vision_tower = CLIPVisionModel(vision_model_config).to( self.device).to(self.dtype) @@ -144,8 +280,12 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, self.mm_projector = hf_mm_projector self.image_newline = hf_image_newline self.vision_feature_select_strategy = getattr( - model_config.pretrained_config, "vision_feature_select_strategy", - "default") + self.pretrained_config, "vision_feature_select_strategy", "default") + + self.post_config() + + def post_config(self): + self.config = self.pretrained_config.vision_config # Copied from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_next/modeling_llava_next.py#L284 def pack_image_features(self, @@ -159,12 +299,12 @@ def pack_image_features(self, if image_feature.shape[0] > 1: base_image_feature = image_feature[0] image_feature = image_feature[1:] - height = width = self.pretrained_config.vision_config.image_size // self.pretrained_config.vision_config.patch_size + height = width = self.config.image_size // self.config.patch_size num_patch_height, num_patch_width = get_anyres_image_grid_shape( image_sizes[image_idx], self.pretrained_config.image_grid_pinpoints, - self.pretrained_config.vision_config.image_size, + self.config.image_size, ) if (np.prod(image_feature.shape) % @@ -226,7 +366,7 @@ def forward(self, multimodal_params: List[MultimodalParams]): image_size_to_num_patches( image_size=imsize, grid_pinpoints=self.pretrained_config.image_grid_pinpoints, - patch_size=self.pretrained_config.vision_config.image_size, + patch_size=self.config.image_size, ) for imsize in image_sizes ] @@ -264,6 +404,7 @@ def forward(self, multimodal_params: List[MultimodalParams]): return [image_features] +@register_vision_encoder(LlavaNextVisionModel) @register_auto_model("LlavaNextForConditionalGeneration") @register_input_processor( LlavaNextInputProcessor, @@ -295,7 +436,6 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, self.llm = AutoModelForCausalLM.from_config(llm_model_config) self.model_config = model_config - self.vocab_size = config.vocab_size self.model_dtype = getattr(config.text_config, "torch_dtype", torch.float16) logger.info(f"{self.dtype=} {self.model_dtype=}") @@ -331,7 +471,14 @@ def forward( mm_embeds = [] if len(multimodal_params) > 0: if not DISAGG: - mm_embeds = self.mm_encoder.forward(multimodal_params) + if multimodal_params[0].multimodal_data.get( + "multimodal_embedding", None) is not None: + mm_embeds = [ + multimodal_param.multimodal_data["multimodal_embedding"] + for multimodal_param in multimodal_params + ] + else: + mm_embeds = self.mm_encoder.forward(multimodal_params) else: mm_embeds = [ multimodal_param.multimodal_data["multimodal_embedding"] diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index 67e329e32696..74b41e8c93a0 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -3,6 +3,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union import torch +import torch.nn as nn from transformers import (AutoProcessor, AutoTokenizer, PretrainedConfig, PreTrainedModel, Qwen2_5_VLForConditionalGeneration, Qwen2VLForConditionalGeneration) @@ -23,7 +24,7 @@ from .modeling_auto import AutoModelForCausalLM from .modeling_multimodal_utils import (find_uncached_mm_embeds, fuse_input_embeds) -from .modeling_utils import register_auto_model +from .modeling_utils import register_auto_model, register_vision_encoder DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1' @@ -309,7 +310,8 @@ def __call__( mm_processor_kwargs) if not mm_data: fused_input_ids = processed_inputs['input_ids'] - return fused_input_ids.to(torch.int32).tolist(), {} + # Flatten the tensor to get a simple list of integers + return fused_input_ids.flatten().to(torch.int32).tolist(), {} pixel_values = processed_inputs.get('pixel_values', None) pixel_values_videos = processed_inputs.get('pixel_values_videos', None) @@ -343,22 +345,29 @@ def __call__( } -class Qwen2VisionModelBase: +class Qwen2VisionModelBase(nn.Module): def __init__(self, model_config: ModelConfig[PretrainedConfig], model_class: type[PreTrainedModel]): - self.pretrained_config = model_config.pretrained_config + super().__init__() + pretrained_config = model_config.pretrained_config + self.model_config = model_config self.device = f"cuda:{model_config.mapping.rank}" - model_path = self.pretrained_config._name_or_path + model_path = pretrained_config._name_or_path # TODO: Change the model class to TRT-LLM's Qwen2VisionModel # Currently, copying vision encoder on all devices. # NOTE: Using attn_implementation='flash_attention_2' to avoid the issue of vision model's GPU OOM. model = model_class.from_pretrained( model_path, - torch_dtype=self.pretrained_config.torch_dtype, + torch_dtype=pretrained_config.torch_dtype, attn_implementation='flash_attention_2').eval() + # TODO: Make vision model compatible with meta init mode and load_weights at the same place self.visual = model.visual.to(self.device) + self.post_config() + + def post_config(self): + self.config = self.visual.config def _to_device( self, input_tensor: Union[torch.Tensor, List, None] @@ -646,6 +655,8 @@ def forward( return output_prob +@register_vision_encoder(Qwen2VisionModelBase, + vlm_base_model=Qwen2VLForConditionalGeneration) @register_auto_model("Qwen2VLForConditionalGeneration") @register_input_processor( Qwen2VLInputProcessorBase, @@ -661,12 +672,14 @@ class Qwen2VLModel(Qwen2VLModelBase): def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs): + super().__init__(model_config, *args, **kwargs) if not DISAGG: self.mm_encoder = Qwen2VisionModelBase( model_config, Qwen2VLForConditionalGeneration) - super().__init__(model_config, *args, **kwargs) +@register_vision_encoder(Qwen2VisionModelBase, + vlm_base_model=Qwen2_5_VLForConditionalGeneration) @register_auto_model("Qwen2_5_VLForConditionalGeneration") @register_input_processor( Qwen2VLInputProcessorBase, @@ -680,7 +693,7 @@ class Qwen2_5_VLModel(Qwen2VLModelBase): def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs): + super().__init__(model_config, *args, **kwargs) if not DISAGG: self.mm_encoder = Qwen2VisionModelBase( model_config, Qwen2_5_VLForConditionalGeneration) - super().__init__(model_config, *args, **kwargs) diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 238ac97ffe02..ad200f1a3e66 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -563,8 +563,14 @@ def infer_max_seq_len(self) -> int: # Step 1: Find the upper bound of max_seq_len inferred_max_seq_len = 2048 - if getattr(self.config, 'max_position_embeddings', None) is not None: - inferred_max_seq_len = self.config.max_position_embeddings + max_position_embeddings = getattr(self.config, + 'max_position_embeddings', None) + if max_position_embeddings is None and hasattr(self.config, + 'text_config'): + max_position_embeddings = getattr(self.config.text_config, + 'max_position_embeddings', None) + if max_position_embeddings is not None: + inferred_max_seq_len = max_position_embeddings # Step 2: Scale max_seq_len with rotary scaling if rope_factor != 1: @@ -579,6 +585,7 @@ def infer_max_seq_len(self) -> int: MODEL_CLASS_MAPPING = {} +MODEL_CLASS_VISION_ENCODER_MAPPING = {} MODEL_CLASS_MAPPER_MAPPING = {} MODEL_CLASS_CHECKPOINT_WEIGHT_LOADER_DEFAULT_MAPPING = {} MODEL_CLASS_CONFIG_LOADER_DEFAULT_MAPPING = {} @@ -594,6 +601,38 @@ def decorator(cls): return decorator +def register_vision_encoder( + vision_encoder_cls: Type[nn.Module], + vlm_base_model: Optional[Type[nn.Module]] = None, +): + """Decorator to register a vision encoder implementation for a pre-registered model architecture. + + Usage: + @register_vision_encoder(MyVisionEncoder, MyVLMBaseModel) + @register_auto_model("SomeVLModel") + class SomeVLModel(...): + ... + The register_auto_model decorator must be applied (executed) before this one (i.e., placed lower) + so that the architecture name is present in MODEL_CLASS_MAPPING. + """ + + def wrapper(model_cls: Type[nn.Module]) -> Type[nn.Module]: + for arch_name, registered_cls in MODEL_CLASS_MAPPING.items(): + if registered_cls.__name__ == model_cls.__name__: + MODEL_CLASS_VISION_ENCODER_MAPPING[arch_name] = ( + vision_encoder_cls, vlm_base_model) + break + else: + raise ValueError( + f"register_vision_encoder: model class {model_cls.__name__} is not registered " + f"via register_auto_model; decorator order must ensure registration occurs first." + ) + + return model_cls + + return wrapper + + def register_mapper(format: str, name: Optional[str] = None): def decorator(cls): diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index fed6a715374a..2f0753ed31a3 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -29,7 +29,8 @@ from .py_executor import PyExecutor from .resource_manager import (KVCacheManager, PeftCacheManager, ResourceManager, ResourceManagerType) -from .sampler import EarlyStopSampler, TorchSampler, TRTLLMSampler +from .sampler import (EarlyStopSampler, EarlyStopWithMMResult, TorchSampler, + TRTLLMSampler) from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler, SimpleScheduler) from .seq_slot_manager import SeqSlotManager @@ -601,6 +602,10 @@ def instantiate_sampler(engine: PyTorchModelEngine, if engine.spec_config is not None and engine.spec_config.spec_dec_mode.has_spec_decoder( ): return get_spec_decoder(sampler_args, engine.spec_config) + + if executor_config.mm_encoder_only: + # NOTE: handle model outputs specially for mm encoder executor/engine + return EarlyStopWithMMResult() if pytorch_backend_config.sampler_type == SamplerType.TRTLLMSampler or ( pytorch_backend_config.sampler_type == SamplerType.auto and decoding_mode.isBeamSearch()): diff --git a/tensorrt_llm/_torch/pyexecutor/config.py b/tensorrt_llm/_torch/pyexecutor/config.py index 3958f5d1f7a1..3be1e5558fc5 100644 --- a/tensorrt_llm/_torch/pyexecutor/config.py +++ b/tensorrt_llm/_torch/pyexecutor/config.py @@ -106,6 +106,9 @@ class PyTorchConfig: force_dynamic_quantization: bool = False + # If true, ONLY the vision encoder part of the full model is loaded/executed. + mm_encoder_only: bool = False + # If true, adjust PyTorch CUDA memory fraction to correspond to the # total GPU memory minus the statically allocated engine memory. # If false, set the PyTorch CUDA memory fraction to 1.0. @@ -119,6 +122,7 @@ class PyTorchConfig: 'tokens_per_block', 'mapping', 'hf_model_dir', + 'mm_encoder_only', ] @@ -133,7 +137,8 @@ def update_executor_config( max_input_len: Optional[int] = None, max_seq_len: Optional[int] = None, checkpoint_format: Optional[str] = None, - checkpoint_loader: Optional[BaseCheckpointLoader] = None): + checkpoint_loader: Optional[BaseCheckpointLoader] = None, + mm_encoder_only: bool = False): if backend is None: return @@ -147,6 +152,7 @@ def update_executor_config( executor_config.pytorch_backend_config = pytorch_backend_config executor_config.mapping = mapping executor_config.speculative_config = speculative_config + executor_config.mm_encoder_only = mm_encoder_only logger.info(f"{executor_config.pytorch_backend_config}") diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 80f1153e5048..96af64fcebb9 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1,10 +1,11 @@ from copy import deepcopy from dataclasses import dataclass -from typing import List, Optional, Union +from typing import Any, Dict, List, Optional, Union import torch import tensorrt_llm.bindings +from tensorrt_llm._torch.shared_tensor import SharedTensorContainer from tensorrt_llm.bindings import executor as tllm_executor from tensorrt_llm.executor.result import TokenLogprobs @@ -172,6 +173,7 @@ def __init__(self, max_new_tokens, use_device_memory, exclude_last_generation_logits ) if return_generation_logits else None self._log_probs = LogProbStorage() if return_log_probs else None + self._mm_embeddings = None def append_context_logits(self, context_logits: torch.Tensor): if self._context_logits: @@ -187,6 +189,10 @@ def append_log_probs(self, if self._log_probs: self._log_probs.append(log_probs, cum_log_probs) + def append_mm_embeddings(self, mm_embeddings: torch.Tensor): + self._mm_embeddings = SharedTensorContainer.from_tensor( + mm_embeddings).dump_to_dict() + def set_log_probs(self, log_probs: list[TokenLogprobs], cum_log_probs: list[float]): """ @@ -224,11 +230,16 @@ def log_probs(self) -> list[TokenLogprobs] | None: def cum_log_probs(self) -> list[float] | None: return self._log_probs and self._log_probs.cum_log_probs + @property + def mm_embedding_handle(self) -> Dict[str, Any] | None: + return self._mm_embeddings + class LlmResult: """LlmResult wraps `bindings.executor.Result` but detour some features to Python implementation""" py_result_properties = frozenset( - ('context_logits', 'generation_logits', 'log_probs', 'cum_log_probs')) + ('context_logits', 'generation_logits', 'log_probs', 'cum_log_probs', + 'mm_embedding_handle')) def __init__(self, result: Union[bytes, tensorrt_llm.bindings.executor.Result], diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b698285c22c3..d9f180c0fc32 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1,5 +1,6 @@ import bisect import contextlib +import copy import functools import gc import inspect @@ -1061,6 +1062,7 @@ def _load_model(self, moe_load_balancer=moe_load_balancer, lora_config=lora_config, allreduce_strategy=self.pytorch_backend_config.allreduce_strategy, + mm_encoder_only=self.pytorch_backend_config.mm_encoder_only, **kwargs) validate_and_set_kv_cache_quant( @@ -1078,9 +1080,12 @@ def _load_model(self, with timing("Model init total"), maybe_create_moe_load_balancer( config, self.mapping) as moe_load_balancer: + try: + # config will be modified in-place for some models, like Qwen2 + config_copy = copy.deepcopy(config) with MetaInitMode(): - model = AutoModelForCausalLM.from_config(config) + model = AutoModelForCausalLM.from_config(config_copy) memo = dict() @@ -1092,6 +1097,7 @@ def init_meta_tensor(t: torch.Tensor): return memo[t] model._apply(init_meta_tensor) + config = config_copy except Exception: logger.info( @@ -1126,6 +1132,12 @@ def init_meta_tensor(t: torch.Tensor): elif load_format == LoadFormat.DUMMY: initialize_dummy_weights(model) + elif load_format == LoadFormat.VISION_ONLY: + # Vision weights are already loaded within the model. + logger.info( + "LoadFormat.VISION_ONLY: skipping weight loading; using preloaded vision weights." + ) + else: raise NotImplementedError( f"No load support for load format: {load_format}") @@ -1151,7 +1163,12 @@ def _call_load_weights(self, load_method, weights, weight_mapper): load_method(weights) def _init_max_seq_len(self): - inferred_max_seq_len = self.model.infer_max_seq_len() + # For mm_encoder_only mode, infer_max_seq_len() is for LLM decoder models + if hasattr(self.model, 'infer_max_seq_len'): + inferred_max_seq_len = self.model.infer_max_seq_len() + else: + inferred_max_seq_len = self._infer_max_seq_len_from_config() + if self.max_seq_len is None: logger.info( f"max_seq_len is not specified, using inferred value {inferred_max_seq_len}" @@ -1167,6 +1184,46 @@ def _init_max_seq_len(self): ) self.max_seq_len = inferred_max_seq_len + def _infer_max_seq_len_from_config(self) -> int: + + if hasattr(self.model, 'model_config') and self.model.model_config: + model_config = self.model.model_config.pretrained_config + rope_scaling = getattr(model_config, 'rope_scaling', None) + rope_factor = 1 + if rope_scaling is not None: + rope_type = rope_scaling.get('type', + rope_scaling.get('rope_type')) + if rope_type not in ("su", "longrope", "llama3", "yarn"): + rope_factor = rope_scaling.get('factor', 1.0) + + # Step 1: Find the upper bound of max_seq_len + inferred_max_seq_len = 2048 + max_position_embeddings = getattr(model_config, + 'max_position_embeddings', None) + if max_position_embeddings is None and hasattr( + model_config, 'text_config'): + max_position_embeddings = getattr(model_config.text_config, + 'max_position_embeddings', + None) + if max_position_embeddings is not None: + inferred_max_seq_len = max_position_embeddings + + # Step 2: Scale max_seq_len with rotary scaling + if rope_factor != 1: + inferred_max_seq_len = int( + math.ceil(inferred_max_seq_len * rope_factor)) + logger.warning( + f'max_seq_len is scaled to {inferred_max_seq_len} by rope scaling {rope_factor}' + ) + + return inferred_max_seq_len + + default_max_seq_len = 8192 + logger.warning( + f"Could not infer max_seq_len from model config, using default value: {default_max_seq_len}" + ) + return default_max_seq_len + def _init_max_num_tokens(self): # Modified from tensorrt_llm/_common.py check_max_num_tokens if self.max_num_tokens is None: @@ -1677,6 +1734,7 @@ def _prepare_tp_inputs_no_cache( multi_modal_data = [] draft_lens = [] request_ids = [] + multimodal_params_list = [] for request in scheduled_requests.context_requests: prompt_tokens = request.get_tokens(0) @@ -1693,6 +1751,17 @@ def _prepare_tp_inputs_no_cache( if multimodal_embedding is not None: multi_modal_data.append(multimodal_embedding) + # Multimodal + if request.py_multimodal_data is not None: + multimodal_params = MultimodalParams( + multimodal_data=request.py_multimodal_data, ) + multimodal_params.to_device("multimodal_data", + "cuda", + pin_memory=True) + multimodal_params_list.append(multimodal_params) + + request.py_batch_idx = request.py_seq_slot + num_tokens = len(input_ids) assert num_tokens <= self.max_num_tokens, ( "num_tokens should be less than or equal to max_num_tokens") @@ -1744,7 +1813,7 @@ def _prepare_tp_inputs_no_cache( 'input_ids': self.input_ids_cuda[:num_tokens], 'position_ids': self.position_ids_cuda[:num_tokens].unsqueeze(0), 'inputs_embeds': None, - 'multi_modal_data': multi_modal_data + "multimodal_params": multimodal_params_list } if bool(lora_params): @@ -2184,8 +2253,13 @@ def forward( scheduled_requests, attn_metadata, spec_metadata) with MoeLoadBalancerIterContext(moe_load_balancer): - return self._forward_step(inputs, gather_ids, - gather_context_logits) + # Special handling for multimodal encoder only mode + if self.pytorch_backend_config.mm_encoder_only: + return self._forward_step_mm_encoder_only( + inputs, scheduled_requests) + else: + return self._forward_step(inputs, gather_ids, + gather_context_logits) with self._maybe_pad_batch(scheduled_requests, kv_cache_manager, spec_resource_manager) as scheduled_requests: maybe_graph = self._maybe_get_cuda_graph(scheduled_requests) @@ -2279,6 +2353,44 @@ def _forward_step(self, else: return {'logits': logits} + @nvtx_range("_forward_step_mm_encoder_only") + def _forward_step_mm_encoder_only( + self, inputs: Dict[str, Any], + scheduled_requests: ScheduledRequests) -> Dict[str, Any]: + """Forward step for multimodal encoder only mode - returns mm_embeddings instead of logits.""" + # Get multimodal parameters from inputs + multimodal_params = inputs.get("multimodal_params", []) + if not multimodal_params or len(multimodal_params) == 0: + # Return empty embeddings if no multimodal data + return {'mm_embeddings': []} + if getattr(scheduled_requests.context_requests[0], 'multimodal_lengths', + None) is None: + multimodal_chunks = None + else: + multimodal_chunks = [ + sum(request.multimodal_lengths) + for request in scheduled_requests.context_requests + if request.multimodal_lengths is not None + ] + # For mm_encoder_only mode, we only run the vision encoder part + # The model should be a vision encoder (e.g., Qwen2VisionModelBase) + mm_embeddings = self.model.forward(multimodal_params) + assert len( + mm_embeddings + ) == 1, "mm_embeddings should be a 1-element list, mix modality (video+image) is not supported" + + if multimodal_chunks is None or len(multimodal_chunks) != len( + multimodal_params): + mm_embeddings = list( + torch.chunk(mm_embeddings[0], + len(scheduled_requests.context_requests), + dim=0)) + else: + mm_embeddings = list( + torch.split(mm_embeddings[0], multimodal_chunks, dim=0)) + + return {'mm_embeddings': mm_embeddings, 'logits': None} + def _init_userbuffers(self, hidden_size): if self.mapping.tp_size <= 1: return False diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index d72173745e05..12686728cdac 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -178,6 +178,17 @@ def _mangle_executor_config(executor_config: ExecutorConfig): ) executor_config.pytorch_backend_config.disable_overlap_scheduler = True + if executor_config.mm_encoder_only: + from tensorrt_llm.llmapi.llm_args import LoadFormat + pytorch_backend_config.mm_encoder_only = True + pytorch_backend_config.load_format = LoadFormat.VISION_ONLY + # Disable overlap scheduler for multimodal encoder-only mode + logger.warning( + "Disabling overlap scheduler for multimodal encoder-only mode. " + "The overlap scheduler is designed for generation models and is not needed " + "when only processing vision encoder inputs.") + pytorch_backend_config.disable_overlap_scheduler = True + def _get_mapping(executor_config: ExecutorConfig) -> Mapping: if executor_config.mapping is None: diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 2fa0e4e3310e..919b99be2d8d 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from collections.abc import Iterable from dataclasses import dataclass -from typing import Literal, Optional +from typing import List, Literal, Optional import torch @@ -97,6 +97,51 @@ def update_requests(self, state: SampleState) -> None: request.py_result.append_context_logits(logits) +@dataclass(kw_only=True) +class MultimodalResult: + mm_embeddings: List[torch.Tensor] = None + + def values(self): + return vars(self).values() + + +@dataclass(kw_only=True) +class SampleStateWithMMResult: + scheduled_requests: ScheduledRequests + + data: MultimodalResult = None + + +class EarlyStopWithMMResult(Sampler): + """ + Use for skipping decoding step for non generation model, and return the batch_output (such as mm_embeddings) + """ + + def sample_async(self, scheduled_requests: ScheduledRequests, + model_outputs) -> SampleStateWithMMResult: + # from model_outputs to MultimodalResult + data = MultimodalResult(mm_embeddings=model_outputs['mm_embeddings']) + return SampleStateWithMMResult(scheduled_requests=scheduled_requests, + data=data) + + def update_requests(self, state: SampleStateWithMMResult) -> None: + assert isinstance(state, SampleStateWithMMResult) + scheduled_requests = state.scheduled_requests + assert (not scheduled_requests.generation_requests) + mm_embeddings = state.data.mm_embeddings + for request, mm_embedding in zip(scheduled_requests.context_requests, + mm_embeddings): + request.state = LlmRequestState.GENERATION_COMPLETE + # NOTE: This is a hack: set finish reason manually and set the beam 0 + request.set_finished_reason(FinishReason.LENGTH, 0) + if len(mm_embedding) != sum(request.multimodal_lengths): + raise ValueError( + f"mm_embedding shape mismatch: {len(mm_embedding)} != {sum(request.multimodal_lengths)}" + ) + + request.py_result.append_mm_embeddings(mm_embedding) + + def top_k_sampling_batch(logits, top_k=50, generator: Optional[torch.Generator] = None): diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index f949fda0d9f3..07eb13d7968c 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -12,6 +12,7 @@ from torch.cuda import device_count from tensorrt_llm import LLM as PyTorchLLM +from tensorrt_llm import MultimodalEncoder from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm._utils import mpi_rank from tensorrt_llm.executor.utils import LlmLauncherEnvs @@ -175,6 +176,22 @@ def launch_server(host: str, asyncio.run(server(host, port)) +def launch_mm_encoder_server( + host: str, + port: int, + encoder_args: dict, + metadata_server_cfg: Optional[MetadataServerConfig] = None, +): + model = encoder_args["model"] + mm_encoder = MultimodalEncoder(**encoder_args) + + server = OpenAIServer(llm=mm_encoder, + model=model, + server_role=ServerRole.MM_ENCODER, + metadata_server_cfg=metadata_server_cfg) + asyncio.run(server(host, port)) + + @click.command("serve") @click.argument("model", type=str) @click.option("--tokenizer", @@ -342,6 +359,79 @@ def serve(model: str, tokenizer: Optional[str], host: str, port: int, launch_server(host, port, llm_args, metadata_server_cfg, server_role) +@click.command("mm_embedding_serve") +@click.argument("model", type=str) +@click.option("--host", + type=str, + default="localhost", + help="Hostname of the server.") +@click.option("--port", type=int, default=8000, help="Port of the server.") +@click.option('--log_level', + type=click.Choice(severity_map.keys()), + default='info', + help="The logging level.") +@click.option("--max_batch_size", + type=int, + default=BuildConfig.max_batch_size, + help="Maximum number of requests that the engine can schedule.") +@click.option( + "--max_num_tokens", + type=int, + default=16384, # set higher default max_num_tokens for multimodal encoder + help= + "Maximum number of batched input tokens after padding is removed in each batch." +) +@click.option("--gpus_per_node", + type=int, + default=None, + help="Number of GPUs per node. Default to None, and it will be " + "detected automatically.") +@click.option("--trust_remote_code", + is_flag=True, + default=False, + help="Flag for HF transformers.") +@click.option( + "--extra_encoder_options", + type=str, + default=None, + help= + "Path to a YAML file that overwrites the parameters specified by trtllm-serve." +) +@click.option("--metadata_server_config_file", + type=str, + default=None, + help="Path to metadata server config file") +def serve_encoder(model: str, host: str, port: int, log_level: str, + max_batch_size: int, max_num_tokens: int, + gpus_per_node: Optional[int], trust_remote_code: bool, + extra_encoder_options: Optional[str], + metadata_server_config_file: Optional[str]): + """Running an OpenAI API compatible server + + MODEL: model name | HF checkpoint path | TensorRT engine path + """ + logger.set_level(log_level) + + # TODO: expose more argument progressivly + llm_args, _ = get_llm_args(model=model, + max_batch_size=max_batch_size, + max_num_tokens=max_num_tokens, + gpus_per_node=gpus_per_node, + trust_remote_code=trust_remote_code) + + encoder_args_extra_dict = {} + if extra_encoder_options is not None: + with open(extra_encoder_options, 'r') as f: + encoder_args_extra_dict = yaml.safe_load(f) + encoder_args = update_llm_args_with_extra_dict(llm_args, + encoder_args_extra_dict) + + metadata_server_cfg = parse_metadata_server_config_file( + metadata_server_config_file) + + launch_mm_encoder_server(host, port, encoder_args, metadata_server_cfg) + + def get_ctx_gen_server_urls( server_configs: List[CtxGenServerConfig]) -> List[str]: ctx_server_urls = [] @@ -647,7 +737,8 @@ def resolve_command(self, ctx, args): commands={ "serve": serve, "disaggregated": disaggregated, - "disaggregated_mpi_worker": disaggregated_mpi_worker + "disaggregated_mpi_worker": disaggregated_mpi_worker, + "mm_embedding_serve": serve_encoder }) if __name__ == "__main__": diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 3d571dfc9d0b..f3ba0bced03a 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -3,8 +3,8 @@ import weakref from dataclasses import dataclass, field from queue import Empty, Queue -from typing import (TYPE_CHECKING, Any, Callable, List, Literal, NamedTuple, - Optional, TypeAlias, Union) +from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, + NamedTuple, Optional, TypeAlias, Union) from weakref import WeakMethod import torch @@ -178,6 +178,7 @@ def __init__(self, CompletionOutput(i) for i in range(self.sampling_params.best_of) ] self._context_logits: Optional[torch.Tensor] = None + self._mm_embedding_handle: Optional[Dict[str, Any]] = None self._background_error_handler = None if background_error_handler is not None: @@ -214,6 +215,10 @@ def outputs(self) -> List[CompletionOutput]: def context_logits(self) -> Optional[torch.Tensor]: return self._context_logits + @property + def mm_embedding_handle(self) -> Optional[Dict[str, Any]]: + return self._mm_embedding_handle + def _handle_sequence(self, finish_reasons, response_tensors, @@ -358,6 +363,10 @@ def _handle_response(self, if response_result.context_logits is not None: self._context_logits = response_result.context_logits + if hasattr(response_result, 'mm_embedding_handle' + ) and response_result.mm_embedding_handle is not None: + self._mm_embedding_handle = response_result.mm_embedding_handle + # Processing background errors here ASAF during generation. if self._background_error_handler and ( handler := self._background_error_handler()): @@ -606,7 +615,7 @@ def _exception(self, timeout: Optional[float] = None): def _repr_fields(self): return [ 'request_id', 'prompt_token_ids', 'outputs', 'finished', - "context_logits" + "context_logits", "mm_embedding_handle" ] def __repr__(self) -> str: diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index a5f49917886c..4981b1639170 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -16,10 +16,12 @@ UserProvidedDecodingConfig) from .llm_utils import (BuildConfig, KvCacheRetentionConfig, QuantAlgo, QuantConfig) +from .mm_encoder import MultimodalEncoder from .mpi_session import MpiCommSession __all__ = [ 'LLM', + 'MultimodalEncoder', 'CompletionOutput', 'RequestOutput', 'GuidedDecodingParams', diff --git a/tensorrt_llm/llmapi/disagg_utils.py b/tensorrt_llm/llmapi/disagg_utils.py index f929c701fe4c..faa9bd08b34a 100644 --- a/tensorrt_llm/llmapi/disagg_utils.py +++ b/tensorrt_llm/llmapi/disagg_utils.py @@ -19,6 +19,7 @@ class ServerRole(Enum): CONTEXT = 0 GENERATION = 1 + MM_ENCODER = 2 @dataclass diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 02298b17431d..9022f7070c70 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -53,6 +53,7 @@ class RequestOutput(DetokenizedGenerationResultBase, GenerationResult): prompt_token_ids (List[int]): The token ids of the prompt. outputs (List[CompletionOutput]): The output sequences of the request. context_logits (torch.Tensor, optional): The logits on the prompt token ids. + mm_embedding_handle (Dict[str, Any], optional): The multimodal embedding handle of the request. finished (bool): Whether the whole request is finished. """ @@ -81,7 +82,8 @@ def prompt(self) -> Optional[str]: def _repr_fields(self): return [ - "request_id", "prompt", "prompt_token_ids", "outputs", "finished" + "request_id", "prompt", "prompt_token_ids", "outputs", "finished", + "mm_embedding_handle" ] @@ -789,7 +791,7 @@ def _build_model(self): # 2. May need to modify model weights for MM (e.g., resize vocab embedding). We must do such operation via input processor's __init__ self.input_processor = create_input_processor(self._hf_model_dir, self.tokenizer) - self.tokenizer = self.input_processor.tokenizer + self._tokenizer = self.input_processor.tokenizer max_batch_size = self.args.max_batch_size max_num_tokens = self.args.max_num_tokens @@ -954,7 +956,7 @@ def _build_model(self): # 2. May need to modify model weights for MM (e.g., resize vocab embedding). We must do such operation via input processor's __init__ self.input_processor = create_input_processor(self._hf_model_dir, self.tokenizer) - self.tokenizer = self.input_processor.tokenizer + self._tokenizer = self.input_processor.tokenizer max_batch_size = self.args.max_batch_size max_num_tokens = self.args.max_num_tokens diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 477c66b2b2c3..da5071e3b0a2 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1966,6 +1966,8 @@ class LoadFormat(Enum): AUTO = 0 # Initialize all weights randomly. DUMMY = 1 + # Only load the multimodal(vision) encoder weights + VISION_ONLY = 2 class SamplerType(StrEnum): diff --git a/tensorrt_llm/llmapi/mm_encoder.py b/tensorrt_llm/llmapi/mm_encoder.py new file mode 100644 index 000000000000..541068a9a6df --- /dev/null +++ b/tensorrt_llm/llmapi/mm_encoder.py @@ -0,0 +1,169 @@ +from pathlib import Path +from typing import Any, List, Literal, Optional, Sequence, Union + +from tqdm import tqdm + +from tensorrt_llm._utils import nvtx_range_debug +from tensorrt_llm.bindings import executor as tllm +from tensorrt_llm.inputs import create_input_processor, prompt_inputs +from tensorrt_llm.inputs.data import PromptInputs +from tensorrt_llm.sampling_params import SamplingParams + +from .llm import BaseLLM, RequestOutput, _TorchLLM +from .llm_args import PybindMirror +from .mpi_session import external_mpi_comm_available + + +class MultimodalEncoder(_TorchLLM): + """MultimodalEncoder class is the main class for running a multimodal encoder model using PyTorch backend. +""" + + def __init__(self, + model: Union[str, Path], + trust_remote_code: bool = False, + tensor_parallel_size: int = 1, + dtype: Literal["auto", "float16", "float32", + "bfloat16"] = "auto", + **kwargs: Any) -> None: + + # Validate that users don't pass LLM-specific or TRT-specific arguments + self._validate_mm_args_for_torch_backend(kwargs) + + # Set higher default max_num_tokens for multimodal encoder (16384 vs 8192 default) + # Vision encoders can handle more tokens than text-only models + # TODO: Make this adaptive based on model-specific max_mm_token_length (see _test_llm_multimodal_general) + if 'max_num_tokens' not in kwargs or kwargs['max_num_tokens'] is None: + kwargs['max_num_tokens'] = 16384 + + super().__init__(model, + trust_remote_code=trust_remote_code, + tensor_parallel_size=tensor_parallel_size, + dtype=dtype, + **kwargs) + + def _build_model(self): + BaseLLM._build_model(self) + assert self._engine_dir is None + + # Tokenizer loading should be after calling model_loader(), since model_loader() may download the model from HF hub. + # It should also be before bindings ExecutorConfig, which may depend on tokenizer info. + self._tokenizer = self._try_load_tokenizer() + + # Multimodal special handling: + # 1. Default load_tokenizer may fail because MM has different tokenizer configuration. Hence we initialize it inside input processor + # 2. May need to modify model weights for MM (e.g., resize vocab embedding). We must do such operation via input processor's __init__ + self.input_processor = create_input_processor(self._hf_model_dir, + self.tokenizer) + self._tokenizer = self.input_processor.tokenizer + + max_batch_size = self.args.max_batch_size + max_num_tokens = self.args.max_num_tokens + max_seq_len = self.args.max_seq_len + + kwargs = {} + if self._on_trt_backend: + kwargs[ + "batching_type"] = self.args.batching_type or tllm.BatchingType.INFLIGHT + + self._executor_config = tllm.ExecutorConfig( + scheduler_config=PybindMirror.maybe_to_pybind( + self.args.scheduler_config), + max_batch_size=max_batch_size, + max_num_tokens=max_num_tokens, + **kwargs) + from tensorrt_llm._torch.pyexecutor.config import update_executor_config + max_batch_size = self._executor_config.max_batch_size + update_executor_config( + self._executor_config, + backend=self.args.backend, + pytorch_backend_config=self.args.get_pytorch_backend_config() + if self.args.backend in ["pytorch", "_autodeploy"] else None, + mapping=self.args.parallel_config.to_mapping(), + hf_model_dir=self._hf_model_dir, + max_input_len=self.args.max_input_len, + max_seq_len=max_seq_len, + checkpoint_format=None if self.args.backend == "_autodeploy" else + self.args.checkpoint_format, + checkpoint_loader=None if self.args.backend == "_autodeploy" else + self.args.checkpoint_loader, + mm_encoder_only=True) + + self._executor = self._executor_cls.create( + self._engine_dir, + executor_config=self._executor_config, + model_world_size=self.args.parallel_config.world_size, + mpi_session=self.mpi_session, + reuse_mpi_comm=external_mpi_comm_available( + self.args.parallel_config.world_size), + is_llm_executor=True, # TODO: check if this is correct or needed + garbage_collection_gen0_threshold=self.args. + garbage_collection_gen0_threshold) + + def _validate_mm_args_for_torch_backend(self, kwargs: dict) -> None: + """Validate that users don't pass LLM-specific arguments when using MultimodalEncoder (PyTorch). + Placeholder for now. + """ + + def generate( + self, + inputs: Union[PromptInputs, Sequence[PromptInputs]], + use_tqdm: bool = True, + ) -> Union[RequestOutput, List[RequestOutput]]: + """Generate output for the given prompts in the synchronous mode. + Synchronous generation accepts either single prompt or batched prompts. + + Args: + inputs (tensorrt_llm.inputs.data.PromptInputs, Sequence[tensorrt_llm.inputs.data.PromptInputs]): The prompt text or token ids. + It can be single prompt or batched prompts. + Returns: + Union[tensorrt_llm.llmapi.RequestOutput, List[tensorrt_llm.llmapi.RequestOutput]]: The output data of the completion request to the LLM. + """ + unbatched = not isinstance(inputs, list) + if not unbatched: + if isinstance(inputs[0], int): + unbatched = True + + if unbatched: + inputs = [inputs] + + inputs = [prompt_inputs(i) for i in inputs] + + def _item_at(maybe_batched: Union[Any, Sequence[Any]], pos: int) -> Any: + if isinstance(maybe_batched, list): + return maybe_batched[pos] + else: + return maybe_batched + + futures = [] + for i, request_inputs in enumerate(inputs): + future = self.generate_async(request_inputs) + futures.append(future) + + for future in tqdm(futures, + desc="Processed requests", + dynamic_ncols=True, + disable=not use_tqdm): + future.result() + + if unbatched: + futures = futures[0] + + return futures + + @nvtx_range_debug("MM_encoder.generate_async", + color="green", + category="VisionEncoder") + def generate_async( + self, + inputs: PromptInputs, + sampling_params: Optional[SamplingParams] = None, + ) -> RequestOutput: + """Generate output for the given multimodal request in the asynchronous mode. + Asynchronous generation accepts single multimodal request only. + + Returns: + Future that resolves to tensorrt_llm.llmapi.RequestOutput containing mm_embeddings + """ + result = super().generate_async(inputs, sampling_params) + # TODO: possible postprocess the result for disaggregated serving + return result diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 24f316028484..7778f8230282 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -392,6 +392,9 @@ class ChatCompletionResponseChoice(OpenAIBaseModel): logprobs: Optional[ChatCompletionLogProbs] = None finish_reason: Optional[str] = None stop_reason: Optional[Union[int, str]] = None + # TODO: progressivly add more info like input_ids, specific_token_ids, mrope, mm_hashes, etc + # TODO: and use a JSON-safe handle to refer to the server-side output + mm_embedding_handle: Optional[Dict[str, Any]] = None disaggregated_params: Optional[DisaggregatedParams] = Field(default=None) avg_decoded_tokens_per_iter: Optional[float] = Field(default=None) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 1b1e15ec6259..cecbdb5d1d4c 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -8,7 +8,7 @@ from datetime import datetime from http import HTTPStatus from pathlib import Path -from typing import Any, AsyncGenerator, AsyncIterator, List, Optional +from typing import Any, AsyncGenerator, AsyncIterator, List, Optional, Union import uvicorn from fastapi import FastAPI, Request @@ -24,6 +24,7 @@ from tensorrt_llm.inputs import prompt_inputs from tensorrt_llm.inputs.utils import ConversationMessage, apply_chat_template from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams +from tensorrt_llm.llmapi import MultimodalEncoder from tensorrt_llm.llmapi.disagg_utils import MetadataServerConfig, ServerRole from tensorrt_llm.llmapi.llm import RequestOutput from tensorrt_llm.logger import logger @@ -33,7 +34,8 @@ from tensorrt_llm.serve.metadata_server import create_metadata_server from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest, ChatCompletionResponse, - CompletionRequest, + ChatCompletionResponseChoice, + ChatMessage, CompletionRequest, CompletionResponse, CompletionResponseChoice, ErrorResponse, ModelCard, @@ -54,7 +56,7 @@ class OpenAIServer: def __init__(self, - llm: LLM, + llm: Union[LLM, MultimodalEncoder], model: str, server_role: Optional[ServerRole], metadata_server_cfg: MetadataServerConfig): @@ -118,7 +120,11 @@ async def lifespan(app: FastAPI): async def validation_exception_handler(_, exc): return self.create_error_response(message=str(exc)) - self.register_routes() + if self.server_role is not ServerRole.MM_ENCODER: + self.register_routes() + else: + assert isinstance(self.llm, MultimodalEncoder), "llm must be a MultimodalEncoder for multimodal encoder" + self.register_mm_encoder_routes() async def await_disconnected(self, raw_request: Request, promise): if raw_request is None: @@ -188,6 +194,16 @@ def mount_metrics(self): metrics_route.path_regex = re.compile("^/prometheus/metrics(?P.*)$") self.app.routes.append(metrics_route) + def register_mm_encoder_routes(self): + self.app.add_api_route("/health", self.health, methods=["GET"]) + self.app.add_api_route("/version", self.version, methods=["GET"]) + self.app.add_api_route("/v1/models", self.get_model, methods=["GET"]) + # TODO: the metrics endpoint only reports iteration stats, not the runtime stats for now + self.app.add_api_route("/metrics", self.get_iteration_stats, methods=["GET"]) + self.app.add_api_route("/v1/chat/completions", + self.openai_mm_encoder, + methods=["POST"]) + async def health(self) -> Response: return Response(status_code=200) @@ -364,6 +380,82 @@ async def create_chat_response( logger.error(traceback.format_exc()) return self.create_error_response(str(e)) + async def openai_mm_encoder(self, request: ChatCompletionRequest, raw_request: Request) -> Response: + + async def create_mm_embedding_response(promise: RequestOutput): + await promise.aresult() + # TODO: Replace mm_embedding_handle with a dedicated OpenAIBaseModel(JSON-safe), when enable multimodal disagg E2E + mm_embedding_handle = getattr(promise, "mm_embedding_handle", None) + if not mm_embedding_handle or "tensor_size" not in mm_embedding_handle: + return self.create_error_response( + message="Multimodal embedding handle missing in response", + err_type="InternalServerError", + status_code=HTTPStatus.INTERNAL_SERVER_ERROR) + num_tokens = int(mm_embedding_handle["tensor_size"][0]) + return ChatCompletionResponse( + id=str(promise.request_id), + model=self.model, + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="dummy"), + mm_embedding_handle=mm_embedding_handle, + finish_reason="length", + ) + ], + usage=UsageInfo( + prompt_tokens=num_tokens, + completion_tokens=1, + total_tokens=num_tokens + 1, + ), + ) + + try: + check_multiple_response(request.n, self.llm.args.backend) + conversation: List[ConversationMessage] = [] + tool_dicts = None if request.tools is None else [ + tool.model_dump() for tool in request.tools + ] + + conversation, mm_coroutines, mm_placeholder_counts = parse_chat_messages_coroutines(request.messages, self.model_config) + + if request.prompt_token_ids is not None: + prompt = request.prompt_token_ids + else: + prompt: str = apply_chat_template( + model_type=self.model_config.model_type, + tokenizer=self.tokenizer, + processor=self.processor, + conversation=conversation, + add_generation_prompt=request.add_generation_prompt, + mm_placeholder_counts=mm_placeholder_counts, + tools=tool_dicts, + documents=request.documents, + chat_template=request.chat_template, + chat_template_kwargs=request.chat_template_kwargs or {}, + ) + prompt = prompt_inputs(prompt) + + mm_data = await mm_coroutines + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + + promise = self.llm.generate_async( + inputs=prompt, + ) + asyncio.create_task(self.await_disconnected(raw_request, promise)) + + response = await create_mm_embedding_response(promise) + return JSONResponse(content=response.model_dump()) + + except CppExecutorError: + logger.error(traceback.format_exc()) + # If internal executor error is raised, shutdown the server + signal.raise_signal(signal.SIGINT) + except Exception as e: + logger.error(traceback.format_exc()) + return self.create_error_response(str(e)) + async def openai_completion(self, request: CompletionRequest, raw_request: Request) -> Response: async def completion_response(promise: RequestOutput, diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py index a46bfd4c912f..da4faf578b5b 100644 --- a/tests/integration/defs/test_e2e.py +++ b/tests/integration/defs/test_e2e.py @@ -1516,6 +1516,13 @@ def test_openai_chat_multimodal_example(llm_root, llm_venv): str(test_root / "_test_openai_chat_multimodal.py")]) +def test_openai_mmencoder_example(llm_root, llm_venv): + test_root = unittest_path() / "llmapi" / "apps" + llm_venv.run_cmd( + ["-m", "pytest", + str(test_root / "_test_openai_mmencoder.py")]) + + def test_openai_chat_structural_tag_example(llm_venv): test_root = unittest_path() / "llmapi" / "apps" llm_venv.run_cmd([ diff --git a/tests/unittest/_torch/multimodal/test_external_embedding.py b/tests/unittest/_torch/multimodal/test_external_embedding.py new file mode 100644 index 000000000000..c2734b67765f --- /dev/null +++ b/tests/unittest/_torch/multimodal/test_external_embedding.py @@ -0,0 +1,196 @@ +from unittest.mock import Mock, patch + +import pytest + +from tensorrt_llm._torch.models.modeling_llava_next import \ + LlavaNextInputProcessor +from tensorrt_llm.inputs.data import TextPrompt +from tensorrt_llm.sampling_params import SamplingParams + +# Model configurations for different processors +MODEL_CONFIGS = { + "LlavaNextInputProcessor": { + "processor_class": + LlavaNextInputProcessor, + "module_path": + "tensorrt_llm._torch.models.modeling_llava_next", + "config_setup": + lambda mock_config: setattr_multiple( + mock_config, { + "image_token_index": 32001, + "vocab_size": 32001, + "text_config.hidden_size": 4096, + "text_config.vocab_size": 32000, + "vision_config": Mock(), + "vision_feature_select_strategy": "default" + }) + }, + # TODO: Add test for more VLM models + # Future model configurations can be added here + # "LlamaInputProcessor": { + # "processor_class": LlamaInputProcessor, + # "module_path": "tensorrt_llm._torch.models.modeling_llama", + # }, + # "QwenInputProcessor": { + # "processor_class": QwenInputProcessor, + # "module_path": "tensorrt_llm._torch.models.modeling_qwen2vl", + # }, +} + + +def setattr_multiple(obj, attr_dict): + """Helper function to set multiple nested attributes.""" + for attr_path, value in attr_dict.items(): + if '.' in attr_path: + # Handle nested attributes like 'text_config.hidden_size' + attrs = attr_path.split('.') + current_obj = obj + for attr in attrs[:-1]: + if not hasattr(current_obj, attr): + setattr(current_obj, attr, Mock()) + current_obj = getattr(current_obj, attr) + setattr(current_obj, attrs[-1], value) + else: + setattr(obj, attr_path, value) + + +@pytest.fixture(params=["LlavaNextInputProcessor"]) +def processor_setup(request): + """Fixture to set up different input processors based on the parameter.""" + processor_name = request.param + config = MODEL_CONFIGS[processor_name] + + # Mock model configuration + mock_config = Mock() + config["config_setup"](mock_config) + + # Mock tokenizer and processor + mock_tokenizer = Mock() + mock_processor = Mock() + + # Create processor instance with mocks + with patch(f'{config["module_path"]}.AutoTokenizer') as mock_auto_tokenizer, \ + patch(f'{config["module_path"]}.AutoProcessor') as mock_auto_processor: + + mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer + mock_auto_processor.from_pretrained.return_value = mock_processor + + processor = config["processor_class"](model_path="dummy_path", + model_config=mock_config, + tokenizer=mock_tokenizer, + trust_remote_code=True) + + # Return processor along with mocks for test usage + return { + "processor": processor, + "mock_tokenizer": mock_tokenizer, + "mock_processor": mock_processor, + "mock_config": mock_config, + "processor_name": processor_name + } + + +class TestExternalEmbedding: + # TODO: Add test for more modalities (audio, video, etc.) + def test_attach_multimodal_embeddings_image_basic(self, processor_setup): + """Test basic functionality of attach_multimodal_embeddings.""" + import torch + + # Prepare test inputs + text_prompt: TextPrompt = {"prompt": "What is in this image? "} + + # Create mock embedding tensors (batch_size=1, seq_len=144, hidden_dim=768) + image_embeddings = [ + torch.randn(144, + 4096), # Single image embedding so num_frames=1 always + ] + + multimodal_embedding = {"image": image_embeddings} + + sampling_params = SamplingParams() + + # Mock tokenizer output - simulate tokenizing the text with image token + mock_token_ids = torch.tensor([1, 2, 3, 32001, 4, + 5]) # image token at index 3 + processor_setup["mock_tokenizer"].return_value.input_ids = [ + mock_token_ids + ] + + result_token_ids, extra_processed_inputs = processor_setup[ + "processor"].attach_multimodal_embeddings(text_prompt, + multimodal_embedding, + sampling_params) + + # Verify outputs + assert isinstance(result_token_ids, list) + assert len(result_token_ids) == len(mock_token_ids) + len( + image_embeddings[0]) - 1 + assert isinstance(extra_processed_inputs, dict) + assert "multimodal_data" in extra_processed_inputs + assert "multimodal_embedding" in extra_processed_inputs[ + "multimodal_data"] + + # Check that multimodal embedding is properly formatted + mm_embedding = extra_processed_inputs["multimodal_data"][ + "multimodal_embedding"] + assert mm_embedding.shape[-1] == 4096 # Hidden dimension should match + assert mm_embedding.shape[0] == 144 + + # Verify tokenizer was called with correct prompt + processor_setup["mock_tokenizer"].assert_called_once_with( + text_prompt["prompt"], return_tensors="pt") + + def test_attach_multimodal_embeddings_multiple_images( + self, processor_setup): + """Test with multiple image embeddings.""" + import torch + + text_prompt: TextPrompt = { + "prompt": "Compare these images: and " + } + + # Create multiple image embeddings + image_embeddings = [ + torch.randn(144, 4096), # First image + torch.randn(288, 4096), # Second image + ] + + # multiple images are concatenated along the first dimension + multimodal_embedding = {"image": image_embeddings} + + sampling_params = SamplingParams() + + # Mock tokenizer output with two image tokens + mock_token_ids = torch.tensor([1, 2, 32001, 3, 4, 32001, 5]) + processor_setup["mock_tokenizer"].return_value.input_ids = [ + mock_token_ids + ] + + # Call the function + result_token_ids, extra_processed_inputs = processor_setup[ + "processor"].attach_multimodal_embeddings(text_prompt, + multimodal_embedding, + sampling_params) + + # Verify outputs + assert isinstance(result_token_ids, list) + assert len(result_token_ids) == len(mock_token_ids) + torch.cat( + image_embeddings, dim=0).shape[0] - 2 + assert isinstance(extra_processed_inputs, dict) + assert "multimodal_data" in extra_processed_inputs + assert "multimodal_embedding" in extra_processed_inputs[ + "multimodal_data"] + + # Check that multimodal embedding is properly formatted + mm_embedding = extra_processed_inputs["multimodal_data"][ + "multimodal_embedding"] + assert mm_embedding.shape[-1] == 4096 # Hidden dimension should match + assert mm_embedding.shape[0] == 144 + 288 + + # Verify tokenizer was called with correct prompt + processor_setup["mock_tokenizer"].assert_called_once_with( + text_prompt["prompt"], return_tensors="pt") + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py b/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py new file mode 100644 index 000000000000..915970d1cea0 --- /dev/null +++ b/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py @@ -0,0 +1,153 @@ +import io + +import pytest +import requests +from PIL import Image +from transformers import AutoConfig, AutoTokenizer + +from tensorrt_llm import MultimodalEncoder +from tensorrt_llm._torch.models.modeling_llava_next import \ + LlavaNextInputProcessor +from tensorrt_llm._torch.models.modeling_qwen2vl import \ + Qwen2VLInputProcessorBase +from tensorrt_llm._torch.shared_tensor import SharedTensorContainer +from tensorrt_llm.inputs import default_multimodal_input_loader + +example_images = [ + "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png", + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png", + "https://huggingface.co/datasets/Sayali9141/traffic_signal_images/resolve/main/61.jpg", +] + + +def download_image(url: str) -> Image.Image: + """Download image from URL and return as PIL Image.""" + response = requests.get(url, timeout=30) + response.raise_for_status() + img = Image.open(io.BytesIO(response.content)) + return img.convert("RGB") + + +@pytest.fixture(scope="function") +def multimodal_model_configs(): + """Get multimodal model configurations for testing.""" + model_configs = { + 'llava-v1.6-mistral-7b-hf': { + 'hf_model_dir': 'llava-hf/llava-v1.6-mistral-7b-hf', + 'model_type': 'llava_next', + }, + 'qwen2.5-vl': { + 'hf_model_dir': 'Qwen/Qwen2.5-VL-3B-Instruct', + 'model_type': 'qwen2_5_vl', + } + } + return model_configs + + +@pytest.mark.parametrize("model_key", [ + "llava-v1.6-mistral-7b-hf", + "qwen2.5-vl", +]) +def test_get_num_tokens_per_image(model_key, multimodal_model_configs): + """Test that get_num_tokens_per_image predicts the correct number of tokens. + + This test verifies that the get_num_tokens_per_image method correctly predicts + the number of image tokens by comparing against actual encoder output shapes. + Tests all example images in a single test run. + """ + # Get model configuration + if model_key not in multimodal_model_configs: + pytest.skip(f"Skipping test for {model_key} - model not available") + + model_config = multimodal_model_configs[model_key] + encoder_model_dir = model_config['hf_model_dir'] + model_type = model_config['model_type'] + + # Test configuration + max_batch_size = len(example_images) # Batch size to handle all test images + + encoder = None + + try: + encoder = MultimodalEncoder(model=encoder_model_dir, + max_batch_size=max_batch_size) + + # Load model configuration and create input processor once + model_config_dict = AutoConfig.from_pretrained(encoder_model_dir, + trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(encoder_model_dir, + trust_remote_code=True) + + # Create input processor once + if model_type == 'llava_next': + input_processor = LlavaNextInputProcessor( + model_path=encoder_model_dir, + model_config=model_config_dict, + tokenizer=tokenizer, + trust_remote_code=True) + elif model_type == 'qwen2_5_vl': + input_processor = Qwen2VLInputProcessorBase( + model_path=encoder_model_dir, + model_config=model_config_dict, + tokenizer=tokenizer, + trust_remote_code=True) + else: + pytest.fail(f"Unsupported model type: {model_type}") + + # Prepare batch inputs with all 3 images + prompts = ["dummy"] * len(example_images) # One prompt per image + media = example_images # All image URLs + + # Prepare inputs and get actual embeddings for all images + inputs = default_multimodal_input_loader(tokenizer=tokenizer, + model_dir=encoder_model_dir, + model_type=model_type, + modality="image", + prompts=prompts, + media=media, + image_data_format="pt") + + # Get actual embeddings from encoder (batch processing) + encoder_outputs = encoder.generate(inputs) + assert len(encoder_outputs) == len( + example_images + ), f"Expected {len(example_images)} encoder outputs, got {len(encoder_outputs)}" + + for image_idx, test_image_url in enumerate(example_images): + + # Get test image dimensions + test_image = download_image(test_image_url) + image_width, image_height = test_image.size + + # Get actual embedding tensor for this image + actual_embedding = SharedTensorContainer.from_dict( + encoder_outputs[image_idx].mm_embedding_handle).get_local_view( + ) + + # The first dimension should be the number of image tokens + actual_num_tokens = actual_embedding.shape[0] + + # Get predicted number of tokens using get_num_tokens_per_image + if model_type == 'llava_next': + predicted_num_tokens = input_processor.get_num_tokens_per_image( + image_width=image_width, image_height=image_height) + elif model_type == 'qwen2_5_vl': + predicted_num_tokens = input_processor.get_num_tokens_per_image( + image_width=image_width, + image_height=image_height, + num_frames=1, + do_resize=True) + + # The key assertion: predicted should match actual + assert predicted_num_tokens == actual_num_tokens, \ + f"Token count mismatch for {model_key} with image {image_idx} ({image_width}x{image_height}): " \ + f"predicted={predicted_num_tokens}, actual={actual_num_tokens}" + + # Additional validation: ensure we got reasonable token counts + assert actual_num_tokens > 0, f"Got zero image tokens for {model_key} image {image_idx}" + assert actual_embedding.ndim >= 2, f"Expected at least 2D embedding, got {actual_embedding.ndim}D" + + finally: + # Cleanup resources + if encoder is not None: + del encoder diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py new file mode 100644 index 000000000000..e42e07654fbf --- /dev/null +++ b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py @@ -0,0 +1,163 @@ +import json +import os + +import pytest + +from tensorrt_llm import MultimodalEncoder +from tensorrt_llm._torch.shared_tensor import SharedTensorContainer +from tensorrt_llm.inputs import default_multimodal_input_loader +from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.llmapi.llm import LLM, SamplingParams + +example_images = [ + "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png", + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png", + "https://huggingface.co/datasets/Sayali9141/traffic_signal_images/resolve/main/61.jpg", +] + + +@pytest.fixture(scope="function") +def multimodal_model_config(): + """Get multimodal model configuration similar to integration tests""" + # You can extend this to support multiple models or get from environment + model_configs = { + 'llava-v1.6-mistral-7b-hf': { + 'model_name': 'llava-v1.6-mistral-7b-hf', + 'hf_model_dir': 'llava-hf/llava-v1.6-mistral-7b-hf', + } + } + + return model_configs['llava-v1.6-mistral-7b-hf'] + + +@pytest.mark.parametrize("model_key", [ + "llava-v1.6-mistral-7b-hf", +]) +def test_single_image_chat(model_key, multimodal_model_config): + """Test processing single image using encoder (pass mm_embeddings) + LLM API. + + This test verifies that encoder (pass mm_embeddings) + LLM API produces identical + results to standard llm generation (pass raw image) by comparing outputs. + """ + # Get model configuration + if model_key != "llava-v1.6-mistral-7b-hf": + #TODO: add more model tests progressively here + pytest.skip( + f"Skipping test for {model_key} - only testing llava-v1.6-mistral-7b-hf for now" + ) + + # Extract model information from config + encoder_model_dir = multimodal_model_config['hf_model_dir'] + + # Test configuration + max_tokens = 64 + free_gpu_memory_fraction = 0.6 + max_batch_size = 1 + + # Test data - OpenAI chat completion format + prompts = ["Describe the natural environment in the image."] + media = [example_images[0]] + + # Sampling configuration + sampling_params = SamplingParams(max_tokens=max_tokens) + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + free_gpu_memory_fraction=free_gpu_memory_fraction, + ) + + # Process multimodal data using encoder (pass mm_embeddings) + encoder = MultimodalEncoder(model=encoder_model_dir, + max_batch_size=max_batch_size) + llm = LLM(model=encoder_model_dir, + backend='pytorch', + kv_cache_config=kv_cache_config, + trust_remote_code=True) + + # Load model configuration + config_path = os.path.join(llm._hf_model_dir, 'config.json') + assert os.path.exists( + config_path), f"Model config not found at {config_path}" + + with open(config_path, 'r') as f: + model_config = json.load(f) + model_type = model_config['model_type'] + + # Prepare multimodal inputs + inputs = default_multimodal_input_loader(tokenizer=llm.tokenizer, + model_dir=llm._hf_model_dir, + model_type=model_type, + modality="image", + prompts=prompts, + media=media, + image_data_format="pt") + + # Validate inputs structure + assert len(inputs) == len( + prompts), f"Expected {len(prompts)} inputs, got {len(inputs)}" + # Generate reference output with raw multimodal inputs + outputs_ref = llm.generate(inputs, sampling_params=sampling_params) + + # Validate reference outputs + assert outputs_ref is not None, "Reference generation returned None" + assert len(outputs_ref) == len( + prompts + ), f"Expected {len(prompts)} reference outputs, got {len(outputs_ref)}" + for i, output in enumerate(outputs_ref): + assert len( + output.outputs + ) > 0, f"Reference generation has no output text for input {i}" + + # Prepare inputs for llm (pass mm_embeddings) + encoder_outputs = encoder.generate(inputs) + inputs = default_multimodal_input_loader( + tokenizer=llm.tokenizer, + model_dir=llm._hf_model_dir, + model_type=model_type, + modality="image", + prompts=prompts, + mm_embeddings=[ + SharedTensorContainer.from_dict( + output.mm_embedding_handle).get_local_view() + for output in encoder_outputs + ], + image_data_format="pt") + + # Generate output using llm (pass mm_embeddings) + outputs = llm.generate(inputs, sampling_params=sampling_params) + + # Validate outputs + assert len(outputs) == len( + prompts), f"Expected {len(prompts)} outputs, got {len(outputs)}" + for i, output in enumerate(outputs): + assert len( + output.outputs) > 0, f"generation has no output text for input {i}" + + # Compare outputs - they should match exactly + assert len(outputs_ref) == len( + outputs + ), f"Number of outputs don't match: {len(outputs_ref)} vs {len(outputs)}" + + for i, (ref_output, test_output) in enumerate(zip(outputs_ref, outputs)): + # Compare prompts + assert ref_output.prompt == test_output.prompt, \ + f"Prompts don't match for output {i}:\nReference: {ref_output.prompt!r}\nTest: {test_output.prompt!r}" + + # Compare number of generated outputs + assert len(ref_output.outputs) == len(test_output.outputs), \ + f"Number of generated outputs don't match for output {i}: {len(ref_output.outputs)} vs {len(test_output.outputs)}" + + # Compare generated text and other attributes + for j, (ref_gen, test_gen) in enumerate( + zip(ref_output.outputs, test_output.outputs)): + assert ref_gen.text == test_gen.text, \ + f"Generated text doesn't match for output {i}, generation {j}:\nReference: {ref_gen.text!r}\nTest: {test_gen.text!r}" + + # Compare token IDs if available + if hasattr(ref_gen, 'token_ids') and hasattr(test_gen, 'token_ids'): + assert ref_gen.token_ids == test_gen.token_ids, \ + f"Token IDs don't match for output {i}, generation {j}" + + # Compare log probabilities if available + if hasattr(ref_gen, 'logprobs') and hasattr(test_gen, 'logprobs'): + assert ref_gen.logprobs == test_gen.logprobs, \ + f"Log probabilities don't match for output {i}, generation {j}" diff --git a/tests/unittest/api_stability/references_committed/request_output.yaml b/tests/unittest/api_stability/references_committed/request_output.yaml index c8f745129cb4..62e8ec234710 100644 --- a/tests/unittest/api_stability/references_committed/request_output.yaml +++ b/tests/unittest/api_stability/references_committed/request_output.yaml @@ -27,3 +27,6 @@ properties: finished: annotation: bool default: inspect._empty + mm_embedding_handle: + annotation: Optional[Dict[str, Any]] + default: inspect._empty diff --git a/tests/unittest/llmapi/apps/_test_openai_mmencoder.py b/tests/unittest/llmapi/apps/_test_openai_mmencoder.py new file mode 100644 index 000000000000..15a1f66cd501 --- /dev/null +++ b/tests/unittest/llmapi/apps/_test_openai_mmencoder.py @@ -0,0 +1,152 @@ +import os +import tempfile +from typing import List + +import openai +import pytest +import requests +import yaml + +from ..test_llm import get_model_path +from .openai_server import RemoteOpenAIServer + +pytestmark = pytest.mark.threadleak(enabled=False) + + +class RemoteMMEncoderServer(RemoteOpenAIServer): + """Remote server for testing multimodal encoder endpoints.""" + + def __init__(self, + model: str, + cli_args: List[str] = None, + port: int = None) -> None: + # Reuse parent initialization but change the command + import subprocess + import sys + + from tensorrt_llm.llmapi.mpi_session import find_free_port + + self.host = "localhost" + self.port = port if port is not None else find_free_port() + self.rank = os.environ.get("SLURM_PROCID", 0) + + args = ["--host", f"{self.host}", "--port", f"{self.port}"] + if cli_args: + args += cli_args + + # Use mm_embedding_serve command instead of regular serve + launch_cmd = ["trtllm-serve", "mm_embedding_serve"] + [model] + args + + self.proc = subprocess.Popen(launch_cmd, + stdout=sys.stdout, + stderr=sys.stderr) + self._wait_for_server(url=self.url_for("health"), + timeout=self.MAX_SERVER_START_WAIT_S) + + +@pytest.fixture(scope="module", ids=["Qwen2.5-VL-3B-Instruct"]) +def model_name(): + return "Qwen2.5-VL-3B-Instruct" + + +@pytest.fixture(scope="module", + params=[True, False], + ids=["extra_options", "no_extra_options"]) +def extra_encoder_options(request): + return request.param + + +@pytest.fixture(scope="module") +def temp_extra_encoder_options_file(request): + temp_dir = tempfile.gettempdir() + temp_file_path = os.path.join(temp_dir, "extra_encoder_options.yaml") + try: + extra_encoder_options_dict = { + "max_batch_size": 8, + "max_num_tokens": 16384 + } + + with open(temp_file_path, 'w') as f: + yaml.dump(extra_encoder_options_dict, f) + + yield temp_file_path + finally: + if os.path.exists(temp_file_path): + os.remove(temp_file_path) + + +@pytest.fixture(scope="module") +def server(model_name: str, extra_encoder_options: bool, + temp_extra_encoder_options_file: str): + model_path = get_model_path(model_name) + args = ["--max_batch_size", "8"] + if extra_encoder_options: + args.extend( + ["--extra_encoder_options", temp_extra_encoder_options_file]) + + with RemoteMMEncoderServer(model_path, args) as remote_server: + yield remote_server + + +@pytest.fixture(scope="module") +def client(server: RemoteMMEncoderServer): + return server.get_client() + + +@pytest.fixture(scope="module") +def async_client(server: RemoteMMEncoderServer): + return server.get_async_client() + + +def test_multimodal_content_mm_encoder(client: openai.OpenAI, model_name: str): + + content_text = "Describe the natural environment in the image." + image_url = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png" + messages = [{ + "role": + "user", + "content": [{ + "type": "text", + "text": content_text + }, { + "type": "image_url", + "image_url": { + "url": image_url + } + }], + }] + + chat_completion = client.chat.completions.create( + model=model_name, + messages=messages, + temperature=0.0, + ) + assert chat_completion.id is not None + assert len(chat_completion.choices) == 1 + choice = chat_completion.choices[0] + # Verify mm_embedding_handle is present + assert hasattr(choice, 'mm_embedding_handle') + assert choice.mm_embedding_handle is not None + # Verify the handle contains tensor information + mm_handle = choice.mm_embedding_handle + assert "tensor_size" in mm_handle + assert mm_handle["tensor_size"][ + 0] == 324 # qwen2.5-vl: 324 tokens for the same image + assert mm_handle["tensor_size"][ + 1] == 2048 # qwen2.5-vl: hidden_size of the vision encoder + + +def test_health(server: RemoteMMEncoderServer): + health_url = server.url_for("health") + response = requests.get(health_url) + assert response.status_code == 200 + + +def test_models_endpoint(client: openai.OpenAI, model_name: str): + models = client.models.list() + assert len(models.data) >= 1 + + model_names = [model.id for model in models.data] + # The model name might be transformed, so check if any model contains our base name + expected_name = model_name.split('/')[-1] + assert any(expected_name in name for name in model_names)