diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index ac3efd14bd61..496acf193912 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -29,6 +29,7 @@ from tensorrt_llm.bench.utils.data import (create_dataset_from_stream, initialize_tokenizer, update_metadata_for_multimodal) +from tensorrt_llm.commands.common_llm_options import common_llm_options from tensorrt_llm.logger import logger from tensorrt_llm.sampling_params import SamplingParams @@ -45,30 +46,6 @@ default=None, help="Path to a serialized TRT-LLM engine.", ) -@optgroup.option( - "--extra_llm_api_options", - type=str, - default=None, - help= - "Path to a YAML file that overwrites the parameters specified by trtllm-bench." -) -@optgroup.option( - "--backend", - type=click.Choice(ALL_SUPPORTED_BACKENDS), - default="pytorch", - help="The backend to use for benchmark. Default is pytorch backend.") -@optgroup.option( - "--kv_cache_free_gpu_mem_fraction", - type=float, - default=.90, - help="The percentage of memory to use for KV Cache after model load.", -) -@optgroup.option( - "--max_seq_len", - type=int, - default=None, - help="Maximum sequence length.", -) @optgroup.group( "Engine Input Configuration", help="Input configuration for driving the engine.", @@ -109,24 +86,6 @@ default=2, help="Number of requests warm up benchmark.", ) -@optgroup.option( - "--tp", - type=int, - default=1, - help="tensor parallelism size", -) -@optgroup.option( - "--pp", - type=int, - default=1, - help="pipeline parallelism size", -) -@optgroup.option( - "--ep", - type=int, - default=None, - help="expert parallelism size", -) @optgroup.group("Request Load Control Options", cls=MutuallyExclusiveOptionGroup, help="Limits how requests are loaded.") @@ -185,6 +144,7 @@ required=False, help="Path where iteration logging is written to.", ) +@common_llm_options @click.pass_obj def latency_command( bench_env: BenchmarkEnvironment, @@ -285,6 +245,7 @@ def latency_command( # Construct the runtime configuration dataclass. runtime_config = RuntimeConfig(**exec_settings) + # TODO: unify LlmArgs parsing via Pydantic llm = None kwargs = kwargs | runtime_config.get_llm_args() kwargs['backend'] = options.backend diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 6406b755c766..d9277245ee49 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -28,6 +28,7 @@ from tensorrt_llm.bench.utils.data import (create_dataset_from_stream, initialize_tokenizer, update_metadata_for_multimodal) +from tensorrt_llm.commands.common_llm_options import common_llm_options from tensorrt_llm.llmapi import CapacitySchedulerPolicy from tensorrt_llm.logger import logger from tensorrt_llm.sampling_params import SamplingParams @@ -45,11 +46,6 @@ default=None, help="Path to a serialized TRT-LLM engine.", ) -@optgroup.option( - "--backend", - type=click.Choice(ALL_SUPPORTED_BACKENDS), - default="pytorch", - help="The backend to use for benchmark. Default is pytorch backend.") @optgroup.option( "--custom_module_dirs", type=click.Path(exists=True, @@ -60,13 +56,6 @@ multiple=True, help="Paths to custom module directories to import.", ) -@optgroup.option( - "--extra_llm_api_options", - type=str, - default=None, - help= - "Path to a YAML file that overwrites the parameters specified by trtllm-bench." -) @optgroup.option("--sampler_options", type=click.Path(exists=True, readable=True, @@ -74,34 +63,12 @@ resolve_path=True), default=None, help="Path to a YAML file that sets sampler options.") -@optgroup.option( - "--max_batch_size", - type=int, - help="Maximum runtime batch size to run the engine with.", -) -@optgroup.option( - "--max_num_tokens", - type=int, - help="Maximum runtime tokens that an engine can accept.", -) -@optgroup.option( - "--max_seq_len", - type=int, - default=None, - help="Maximum sequence length.", -) @optgroup.option( "--beam_width", type=int, default=1, help="Number of search beams.", ) -@optgroup.option( - "--kv_cache_free_gpu_mem_fraction", - type=float, - default=.90, - help="The percentage of memory to use for KV Cache after model load.", -) @optgroup.group( "Engine Input Configuration", help="Input configuration for driving the engine.", @@ -116,16 +83,6 @@ 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, - default=False, - help="Do not skip tokenizer initialization when loading the model.", -) @optgroup.option( "--eos_id", type=int, @@ -152,14 +109,6 @@ default="cuda", help="Device to load the multimodal data on.", ) -@optgroup.option( - "--max_input_len", - type=int, - default=4096, - help= - "Maximum input sequence length to use for multimodal models. This is used only when --modality " - "is specified since the actual number of vision tokens is unknown before the model is run.", -) @optgroup.option( "--num_requests", type=int, @@ -186,34 +135,6 @@ type=click.IntRange(min=1), help="Target (average) sequence length for tuning heuristics.", ) -@optgroup.group( - "World Configuration", - help="Options for configuring the backend multi-GPU world.", -) -@optgroup.option( - "--tp", - type=int, - default=1, - help="tensor parallelism size", -) -@optgroup.option( - "--pp", - type=int, - default=1, - help="pipeline parallelism size", -) -@optgroup.option( - "--ep", - type=int, - default=None, - help="expert parallelism size", -) -@optgroup.option( - "--cluster_size", - type=int, - default=None, - help="expert cluster parallelism size", -) @optgroup.group("Request Load Control Options", cls=MutuallyExclusiveOptionGroup, help="Limits how requests are loaded.") @@ -286,6 +207,7 @@ help= "KV cache scheduler policy: guaranteed_no_evict prevents request eviction, max_utilization optimizes for throughput.", ) +@common_llm_options @click.pass_obj def throughput_command( bench_env: BenchmarkEnvironment, @@ -296,7 +218,11 @@ def throughput_command( # Parameters from CLI image_data_format: str = params.get("image_data_format", "pt") data_device: str = params.get("data_device", "cpu") - no_skip_tokenizer_init: bool = params.get("no_skip_tokenizer_init", False) + # 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, hence we expose the option to the user. + skip_tokenizer_init: bool = params.get("skip_tokenizer_init", True) # Get general CLI options using the centralized function options: GeneralExecSettings = get_general_cli_options(params, bench_env) @@ -396,6 +322,7 @@ def throughput_command( # Dynamic runtime features. exec_settings["settings_config"]["dynamic_max_batch_size"] = True + # TODO: unify LlmArgs parsing via Pydantic # LlmArgs exec_settings["extra_llm_api_options"] = params.pop("extra_llm_api_options") exec_settings["iteration_log"] = options.iteration_log @@ -407,7 +334,7 @@ def throughput_command( try: logger.info("Setting up throughput benchmark.") kwargs = kwargs | runtime_config.get_llm_args() - kwargs['skip_tokenizer_init'] = not no_skip_tokenizer_init + kwargs['skip_tokenizer_init'] = skip_tokenizer_init kwargs['backend'] = options.backend llm = get_llm(runtime_config, kwargs) diff --git a/tensorrt_llm/commands/common_llm_options.py b/tensorrt_llm/commands/common_llm_options.py new file mode 100644 index 000000000000..12e7eaae4079 --- /dev/null +++ b/tensorrt_llm/commands/common_llm_options.py @@ -0,0 +1,153 @@ +import typing +from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Type + +import click +from click_option_group import OptionGroup, optgroup +from pydantic import BaseModel + +from tensorrt_llm.llmapi.llm_args import BaseLlmArgs, KvCacheConfig +from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory + + +def option_from_field( + field_name: str, + option_names: Optional[List[str]] = None, + model_class: Type[BaseModel] = BaseLlmArgs, +) -> Callable: + """Create a CLI option decorator from a field in a Pydantic model. + + Args: + field_name: Name of the field in the pydantic model + option_name: CLI option name (default: --{field_name}) + model_class: Pydantic model class containing the field + + Returns: + A decorator function that adds the option + """ + if option_names is None: + option_names = [f"--{field_name}"] + + field_info = model_class.model_fields[field_name] + + field_type = field_info.annotation + # Handle Optional types + if typing.get_origin(field_type) is typing.Union: + args = typing.get_args(field_type) + # Get the non-None type + field_type = next((arg for arg in args if arg is not type(None)), + args[0]) + + elif field_type not in (int, float, str, bool): + raise ValueError(f"Unsupported field type: {field_type}") + + return optgroup.option(*option_names, + type=field_type, + is_flag=field_type is bool, + default=field_info.default, + help=field_info.description) + + +class ChoiceWithAlias(click.Choice): + + def __init__(self, + choices: Sequence[str], + aliases: Mapping[str, str], + case_sensitive: bool = True) -> None: + super().__init__(choices, case_sensitive) + self.aliases = aliases + + def to_info_dict(self) -> Dict[str, Any]: + info_dict = super().to_info_dict() + info_dict["aliases"] = self.aliases + return info_dict + + def convert(self, value: Any, param: Optional["click.Parameter"], + ctx: Optional["click.Context"]) -> Any: + if value in self.aliases: + value = self.aliases[value] + return super().convert(value, param, ctx) + + +def common_llm_options(f: Callable) -> Callable: + """Apply all common LLM API options to a command. + + This decorator adds all LLM API-related configuration options and organizes them into + logical groups. + """ + # Model and backend configuration + f = optgroup.group("Model Configuration", + help="Model, tokenizer, and backend settings.")(f) + f = option_from_field("tokenizer")(f) + f = optgroup.option( + "--backend", + type=ChoiceWithAlias(["pytorch", "tensorrt", "_autodeploy"], + {"trt": "tensorrt"}), + default="pytorch", + help=BaseLlmArgs.model_fields["backend"].description)(f) + f = option_from_field("trust_remote_code")(f) + f = optgroup.option( + "--extra_llm_api_options", + type=str, + default=None, + help= + "Path to a YAML file with LLM API configuration options for serving the model." + )(f) + + # Parallelism configuration + f = optgroup.group("Parallelism Configuration", + help="Multi-GPU and distributed execution settings.", + cls=OptionGroup)(f) + f = option_from_field("tensor_parallel_size", + option_names=["--tp_size", "--tp"])(f) + f = option_from_field("pipeline_parallel_size", + option_names=["--pp_size", "--pp"])(f) + f = option_from_field("moe_expert_parallel_size", + option_names=["--ep_size", "--ep"])(f) + f = option_from_field("moe_cluster_parallel_size", + option_names=["--cluster_size"])(f) + + # Build and runtime limits + f = optgroup.group( + "Build and Runtime Limits", + help="Maximum batch size, sequence length, and token limits.", + cls=OptionGroup)(f) + f = option_from_field("max_batch_size")(f) + f = option_from_field("max_num_tokens")(f) + f = option_from_field("max_seq_len")(f) + f = option_from_field("max_beam_width")(f) + f = option_from_field("max_input_len")(f) + + # KV cache configuration + f = optgroup.group("KV Cache Configuration", + help="KV cache memory and reuse settings.", + cls=OptionGroup)(f) + f = option_from_field("free_gpu_memory_fraction", + option_names=[ + "--kv_cache_free_gpu_memory_fraction", + "--kv_cache_free_gpu_mem_fraction" + ], + model_class=KvCacheConfig)(f) + f = optgroup.option("--disable_kv_cache_reuse", + is_flag=True, + default=False, + help="Flag for disabling KV cache reuse.")(f) + + # Postprocessing options + f = optgroup.group("Postprocessing Options", + help="Output processing and formatting settings.", + cls=OptionGroup)(f) + f = option_from_field("num_postprocess_workers")(f) + f = optgroup.option( + "--reasoning_parser", + type=click.Choice(ReasoningParserFactory.parsers.keys()), + default=None, + help=BaseLlmArgs.model_fields["reasoning_parser"].description)(f) + + # Advanced options + f = optgroup.group("Advanced Options", + help="Advanced configuration and debugging settings.", + cls=OptionGroup)(f) + f = option_from_field("fail_fast_on_attention_window_too_large")(f) + f = option_from_field("skip_tokenizer_init")(f) + + return f diff --git a/tensorrt_llm/commands/eval.py b/tensorrt_llm/commands/eval.py index 0b5d7ba4774a..045503663d13 100644 --- a/tensorrt_llm/commands/eval.py +++ b/tensorrt_llm/commands/eval.py @@ -12,14 +12,13 @@ # 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. -from typing import Optional - import click import tensorrt_llm.profiler as profiler from .. import LLM as PyTorchLLM from .._tensorrt_engine import LLM +from ..commands.common_llm_options import common_llm_options from ..evaluate import (GSM8K, MMLU, MMMU, CnnDailymail, GPQADiamond, GPQAExtended, GPQAMain, JsonModeEval) from ..llmapi import BuildConfig, KvCacheConfig @@ -34,106 +33,44 @@ type=str, help="model name | HF checkpoint path | TensorRT engine path", ) -@click.option("--tokenizer", - type=str, - default=None, - help="Path | Name of the tokenizer." - "Specify this value only if using TensorRT engine as model.") -@click.option( - "--backend", - type=click.Choice(["pytorch", "tensorrt"]), - default="pytorch", - help="The backend to use for evaluation. Default is pytorch backend.") @click.option('--log_level', type=click.Choice(severity_map.keys()), default='info', help="The logging level.") -@click.option("--max_beam_width", - type=int, - default=BuildConfig.max_beam_width, - help="Maximum number of beams for beam search decoding.") -@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=BuildConfig.max_num_tokens, - help= - "Maximum number of batched input tokens after padding is removed in each batch." -) -@click.option( - "--max_seq_len", - type=int, - default=BuildConfig.max_seq_len, - help="Maximum total length of one request, including prompt and outputs. " - "If unspecified, the value is deduced from the model config.") -@click.option("--tp_size", type=int, default=1, help='Tensor parallelism size.') -@click.option("--pp_size", - type=int, - default=1, - help='Pipeline parallelism size.') -@click.option("--ep_size", - type=int, - default=None, - help="expert parallelism size") -@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("--kv_cache_free_gpu_memory_fraction", - type=float, - default=0.9, - help="Free GPU memory fraction reserved for KV Cache, " - "after allocating model weights and buffers.") -@click.option("--trust_remote_code", - is_flag=True, - default=False, - help="Flag for HF transformers.") -@click.option("--extra_llm_api_options", - type=str, - default=None, - help="Path to a YAML file that overwrites the parameters") -@click.option("--disable_kv_cache_reuse", - is_flag=True, - default=False, - help="Flag for disabling KV cache reuse.") +@common_llm_options @click.pass_context -def main(ctx, model: str, tokenizer: Optional[str], log_level: str, - backend: str, max_beam_width: int, max_batch_size: int, - max_num_tokens: int, max_seq_len: int, tp_size: int, pp_size: int, - ep_size: Optional[int], gpus_per_node: Optional[int], - kv_cache_free_gpu_memory_fraction: float, trust_remote_code: bool, - extra_llm_api_options: Optional[str], disable_kv_cache_reuse: bool): +def main(ctx, model: str, log_level: str, **params): logger.set_level(log_level) - build_config = BuildConfig(max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - max_beam_width=max_beam_width, - max_seq_len=max_seq_len) + # TODO: unify LlmArgs parsing via Pydantic + build_config = BuildConfig(max_batch_size=params.get("max_batch_size"), + max_num_tokens=params.get("max_num_tokens"), + max_beam_width=params.get("max_beam_width"), + max_seq_len=params.get("max_seq_len")) kv_cache_config = KvCacheConfig( - free_gpu_memory_fraction=kv_cache_free_gpu_memory_fraction, - enable_block_reuse=not disable_kv_cache_reuse) + free_gpu_memory_fraction=params.get("kv_cache_free_gpu_memory_fraction", + 0.9), + enable_block_reuse=not params.get("disable_kv_cache_reuse", False)) llm_args = { "model": model, - "tokenizer": tokenizer, - "tensor_parallel_size": tp_size, - "pipeline_parallel_size": pp_size, - "moe_expert_parallel_size": ep_size, - "gpus_per_node": gpus_per_node, - "trust_remote_code": trust_remote_code, + "tokenizer": params.get("tokenizer"), + "tensor_parallel_size": params.get("tensor_parallel_size", 1), + "pipeline_parallel_size": params.get("pipeline_parallel_size", 1), + "moe_expert_parallel_size": params.get("moe_expert_parallel_size"), + "gpus_per_node": params.get("gpus_per_node"), + "trust_remote_code": params.get("trust_remote_code", False), "build_config": build_config, "kv_cache_config": kv_cache_config, } + extra_llm_api_options = params.get("extra_llm_api_options") if extra_llm_api_options is not None: llm_args = update_llm_args_with_extra_options(llm_args, extra_llm_api_options) profiler.start("trtllm init") + backend = params.get("backend", "pytorch") if backend == 'pytorch': llm = PyTorchLLM(**llm_args) elif backend == 'tensorrt': diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 7da0930264be..c02a89dc6a06 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -4,7 +4,7 @@ import signal # Added import import subprocess # nosec B404 import sys -from typing import Any, Dict, Mapping, Optional, Sequence +from typing import Any, Optional import click import torch @@ -17,6 +17,7 @@ from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm._torch.auto_deploy.llm import LLM as AutoDeployLLM from tensorrt_llm._utils import mpi_rank +from tensorrt_llm.commands.common_llm_options import common_llm_options from tensorrt_llm.executor.utils import LlmLauncherEnvs from tensorrt_llm.llmapi import (BuildConfig, CapacitySchedulerPolicy, DynamicBatchConfig, KvCacheConfig, @@ -26,7 +27,6 @@ parse_metadata_server_config_file) from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict from tensorrt_llm.llmapi.mpi_session import find_free_port -from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory from tensorrt_llm.logger import logger, severity_map from tensorrt_llm.serve import OpenAIDisaggServer, OpenAIServer @@ -202,116 +202,17 @@ def launch_mm_encoder_server( asyncio.run(server(host, port)) -class ChoiceWithAlias(click.Choice): - - def __init__(self, - choices: Sequence[str], - aliases: Mapping[str, str], - case_sensitive: bool = True) -> None: - super().__init__(choices, case_sensitive) - self.aliases = aliases - - def to_info_dict(self) -> Dict[str, Any]: - info_dict = super().to_info_dict() - info_dict["aliases"] = self.aliases - return info_dict - - def convert(self, value: Any, param: Optional["click.Parameter"], - ctx: Optional["click.Context"]) -> Any: - if value in self.aliases: - value = self.aliases[value] - return super().convert(value, param, ctx) - - @click.command("serve") @click.argument("model", type=str) -@click.option("--tokenizer", - type=str, - default=None, - help="Path | Name of the tokenizer." - "Specify this value only if using TensorRT engine as model.") @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( - "--backend", - type=ChoiceWithAlias(["pytorch", "tensorrt", "_autodeploy"], - {"trt": "tensorrt"}), - default="pytorch", - help="The backend to use to serve the model. Default is pytorch backend.") @click.option('--log_level', type=click.Choice(severity_map.keys()), default='info', help="The logging level.") -@click.option("--max_beam_width", - type=int, - default=BuildConfig.max_beam_width, - help="Maximum number of beams for beam search decoding.") -@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=BuildConfig.max_num_tokens, - help= - "Maximum number of batched input tokens after padding is removed in each batch." -) -@click.option( - "--max_seq_len", - type=int, - default=BuildConfig.max_seq_len, - help="Maximum total length of one request, including prompt and outputs. " - "If unspecified, the value is deduced from the model config.") -@click.option("--tp_size", type=int, default=1, help='Tensor parallelism size.') -@click.option("--pp_size", - type=int, - default=1, - help='Pipeline parallelism size.') -@click.option("--ep_size", - type=int, - default=None, - help="expert parallelism size") -@click.option("--cluster_size", - type=int, - default=None, - help="expert cluster parallelism size") -@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("--kv_cache_free_gpu_memory_fraction", - type=float, - default=0.9, - help="Free GPU memory fraction reserved for KV Cache, " - "after allocating model weights and buffers.") -@click.option( - "--num_postprocess_workers", - type=int, - default=0, - help="[Experimental] Number of workers to postprocess raw responses " - "to comply with OpenAI protocol.") -@click.option("--trust_remote_code", - is_flag=True, - default=False, - help="Flag for HF transformers.") -@click.option( - "--extra_llm_api_options", - type=str, - default=None, - help= - "Path to a YAML file that overwrites the parameters specified by trtllm-serve." -) -@click.option( - "--reasoning_parser", - type=click.Choice(ReasoningParserFactory.parsers.keys()), - default=None, - help="[Experimental] Specify the parser for reasoning models.", -) @click.option("--metadata_server_config_file", type=str, default=None, @@ -322,50 +223,40 @@ def convert(self, value: Any, param: Optional["click.Parameter"], default=None, help="Server role. Specify this value only if running in disaggregated mode." ) -@click.option( - "--fail_fast_on_attention_window_too_large", - is_flag=True, - default=False, - help= - "Exit with runtime error when attention window is too large to fit even a single sequence in the KV cache." -) -def serve( - model: str, tokenizer: Optional[str], host: str, port: int, - log_level: str, backend: str, max_beam_width: int, max_batch_size: int, - max_num_tokens: int, max_seq_len: int, tp_size: int, pp_size: int, - ep_size: Optional[int], cluster_size: Optional[int], - gpus_per_node: Optional[int], kv_cache_free_gpu_memory_fraction: float, - num_postprocess_workers: int, trust_remote_code: bool, - extra_llm_api_options: Optional[str], reasoning_parser: Optional[str], - metadata_server_config_file: Optional[str], server_role: Optional[str], - fail_fast_on_attention_window_too_large: bool): +@common_llm_options +def serve(model: str, host: str, port: int, log_level: str, + metadata_server_config_file: Optional[str], + server_role: Optional[str], **params): """Running an OpenAI API compatible server MODEL: model name | HF checkpoint path | TensorRT engine path """ logger.set_level(log_level) + # TODO: unify LlmArgs parsing via Pydantic llm_args, _ = get_llm_args( model=model, - tokenizer=tokenizer, - backend=backend, - max_beam_width=max_beam_width, - max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - max_seq_len=max_seq_len, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - moe_cluster_parallel_size=cluster_size, - gpus_per_node=gpus_per_node, - free_gpu_memory_fraction=kv_cache_free_gpu_memory_fraction, - num_postprocess_workers=num_postprocess_workers, - trust_remote_code=trust_remote_code, - reasoning_parser=reasoning_parser, - fail_fast_on_attention_window_too_large= - fail_fast_on_attention_window_too_large) + tokenizer=params.get("tokenizer"), + backend=params.get("backend", "pytorch"), + max_beam_width=params.get("max_beam_width"), + max_batch_size=params.get("max_batch_size"), + max_num_tokens=params.get("max_num_tokens"), + max_seq_len=params.get("max_seq_len"), + tensor_parallel_size=params.get("tensor_parallel_size", 1), + pipeline_parallel_size=params.get("pipeline_parallel_size", 1), + moe_expert_parallel_size=params.get("moe_expert_parallel_size"), + moe_cluster_parallel_size=params.get("moe_cluster_parallel_size"), + gpus_per_node=params.get("gpus_per_node"), + free_gpu_memory_fraction=params.get("kv_cache_free_gpu_memory_fraction", + 0.9), + num_postprocess_workers=params.get("num_postprocess_workers", 0), + trust_remote_code=params.get("trust_remote_code", False), + reasoning_parser=params.get("reasoning_parser"), + fail_fast_on_attention_window_too_large=params.get( + "fail_fast_on_attention_window_too_large", False)) llm_args_extra_dict = {} + extra_llm_api_options = params.get("extra_llm_api_options") if extra_llm_api_options is not None: with open(extra_llm_api_options, 'r') as f: llm_args_extra_dict = yaml.safe_load(f) @@ -395,42 +286,13 @@ def serve( 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") +@common_llm_options 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]): + metadata_server_config_file: Optional[str], **params): """Running an OpenAI API compatible server MODEL: model name | HF checkpoint path | TensorRT engine path @@ -438,18 +300,20 @@ def serve_encoder(model: str, host: str, port: int, log_level: str, 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) + llm_args, _ = get_llm_args( + model=model, + max_batch_size=params.get("max_batch_size"), + max_num_tokens=params.get("max_num_tokens"), + gpus_per_node=params.get("gpus_per_node"), + trust_remote_code=params.get("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) + extra_llm_api_options_dict = {} + extra_llm_api_options = params.get("extra_llm_api_options") + if extra_llm_api_options is not None: + with open(extra_llm_api_options, 'r') as f: + extra_llm_api_options_dict = yaml.safe_load(f) encoder_args = update_llm_args_with_extra_dict(llm_args, - encoder_args_extra_dict) + extra_llm_api_options_dict) metadata_server_cfg = parse_metadata_server_config_file( metadata_server_config_file) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 89a0d8d6193a..a4145269a493 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1130,8 +1130,9 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): free_gpu_memory_fraction: Optional[float] = Field( default=None, description= - "The fraction of GPU memory fraction that should be allocated for the KV cache. Default is 90%. If both `max_tokens` and `free_gpu_memory_fraction` are specified, memory corresponding to the minimum will be used." - ) + "The fraction of GPU memory fraction that should be allocated for the KV cache, after allocating model weights " + "and buffers. Default is 90%. If both `max_tokens` and `free_gpu_memory_fraction` are specified, memory " + "corresponding to the minimum will be used.") host_cache_size: Optional[int] = Field( default=None, description= @@ -1352,7 +1353,8 @@ class BaseLlmArgs(StrictBaseModel): tokenizer: Optional[Union[ str, Path, TokenizerBase, PreTrainedTokenizerBase]] = Field( description= - "The path to the tokenizer checkpoint or the tokenizer name from the Hugging Face Hub.", + "The path to the tokenizer checkpoint or the tokenizer name from the Hugging Face Hub. " + "Specify this value only if using a TensorRT engine as the model.", default=None) tokenizer_mode: Literal['auto', 'slow'] = Field( @@ -1389,7 +1391,8 @@ class BaseLlmArgs(StrictBaseModel): gpus_per_node: Optional[int] = Field( default=None, - description="The number of GPUs per node.", + description= + "The number of GPUs per node. If not provided, it will be detected automatically.", status="beta", validate_default=True) @@ -1484,21 +1487,27 @@ class BaseLlmArgs(StrictBaseModel): speculative_config: SpeculativeConfig = Field( default=None, description="Speculative decoding config.") - max_batch_size: Optional[int] = Field(default=None, - description="The maximum batch size.") + max_batch_size: Optional[int] = Field( + default=None, + description= + "The maximum number of requests that can be scheduled in one batch.") # generation constraints max_input_len: Optional[int] = Field( default=None, description="The maximum input length.") max_seq_len: Optional[int] = Field( - default=None, description="The maximum sequence length.") + default=None, + description= + "The maximum total length of one request, including prompt and outputs. " + "If unspecified, the value is deduced from the model config.") - max_beam_width: Optional[int] = Field(default=None, - description="The maximum beam width.") + max_beam_width: Optional[int] = Field( + default=None, + description="The maximum number of beams for beam search decoding.") max_num_tokens: Optional[int] = Field( - default=None, description="The maximum number of tokens.") + default=None, description="The maximum number of tokens in each batch.") gather_generation_logits: bool = Field( default=False, diff --git a/tests/unittest/llmapi/apps/_test_openai_mmencoder.py b/tests/unittest/llmapi/apps/_test_openai_mmencoder.py index 15a1f66cd501..54303357ee33 100644 --- a/tests/unittest/llmapi/apps/_test_openai_mmencoder.py +++ b/tests/unittest/llmapi/apps/_test_openai_mmencoder.py @@ -52,22 +52,22 @@ def model_name(): @pytest.fixture(scope="module", params=[True, False], ids=["extra_options", "no_extra_options"]) -def extra_encoder_options(request): +def extra_llm_api_options(request): return request.param @pytest.fixture(scope="module") -def temp_extra_encoder_options_file(request): +def temp_extra_llm_api_options_file(request): temp_dir = tempfile.gettempdir() - temp_file_path = os.path.join(temp_dir, "extra_encoder_options.yaml") + temp_file_path = os.path.join(temp_dir, "extra_llm_api_options.yaml") try: - extra_encoder_options_dict = { + extra_llm_api_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) + yaml.dump(extra_llm_api_options_dict, f) yield temp_file_path finally: @@ -76,13 +76,13 @@ def temp_extra_encoder_options_file(request): @pytest.fixture(scope="module") -def server(model_name: str, extra_encoder_options: bool, - temp_extra_encoder_options_file: str): +def server(model_name: str, extra_llm_api_options: bool, + temp_extra_llm_api_options_file: str): model_path = get_model_path(model_name) args = ["--max_batch_size", "8"] - if extra_encoder_options: + if extra_llm_api_options: args.extend( - ["--extra_encoder_options", temp_extra_encoder_options_file]) + ["--extra_llm_api_options", temp_extra_llm_api_options_file]) with RemoteMMEncoderServer(model_path, args) as remote_server: yield remote_server