__all__ = [ 'LlmArgs', 'LlmBuildStats', 'ModelLoader', '_ModelRuntimeContext', '_ModelInfo', '_ParallelConfig', '_ModelFormatKind', 'BatchingType', 'ExecutorConfig', 'SchedulerConfig', 'KvCacheConfig', 'ContextChunkingPolicy', 'CapacitySchedulerPolicy', 'BuildConfig', 'BuildCacheConfig', 'QuantConfig', 'CachedModelLoader', 'ConfigArbitrateError', '_ConfigArbitrator', ] import copy import json import os import shutil import tempfile import time from argparse import Namespace from dataclasses import asdict, dataclass, field, fields from enum import Enum from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Tuple, Union import tensorrt as trt import torch from tqdm import tqdm from transformers import PreTrainedTokenizerBase from .._utils import mpi_barrier, mpi_broadcast, mpi_rank, release_gc from ..auto_parallel import AutoParallelConfig, infer_cluster_config from ..bindings.executor import (BatchingType, CapacitySchedulerPolicy, ContextChunkingPolicy, DecodingConfig, ExecutorConfig, KvCacheConfig, PeftCacheConfig, SchedulerConfig) from ..builder import BuildConfig, Engine, EngineConfig, build from ..logger import logger from ..mapping import Mapping from ..models import MODEL_MAP from ..models.modeling_utils import (PretrainedConfig, QuantAlgo, QuantConfig, TopModelMixin) from ..module import Module from .build_cache import (BuildCache, BuildCacheConfig, CachedStage, get_build_cache_config_from_env) from .mpi_session import MPINodeState, MpiSession from .tokenizer import TokenizerBase, TransformersTokenizer, tokenizer_factory # TODO[chunweiy]: move the following symbols back to utils scope, and remove the following import from .utils import (GpuArch, download_hf_model, download_hf_pretrained_config, file_with_glob_exists, file_with_suffix_exists, print_colored, print_traceback_on_error) @dataclass class _ParallelConfig: ''' The model distribution configs for LLM. ''' tp_size: int = 1 pp_size: int = 1 auto_parallel: bool = False _world_size: int = field(default=1, init=False) _devices: Optional[List[int]] = field(default=None, init=False) @property def devices(self) -> List[int]: if self._devices is None: return list(range(self.world_size)) return self._devices @devices.setter def devices(self, devices: List[int]): if len(devices) != self.world_size: raise ValueError( f"devices {devices} should have the same length as world_size {self.world_size}" ) self._devices = devices @property def world_size(self) -> bool: if self.auto_parallel: if self.tp_size > 1 or self.pp_size > 1: raise RuntimeError( "manually TP and PP are not supported in auto parallel mode." ) return self._world_size if self._world_size > 1: raise RuntimeError( "world_size > 1 is only supported in auto parallel mode.") return self.tp_size * self.pp_size @world_size.setter def world_size(self, world_size: int): if self.auto_parallel: self._world_size = world_size elif (not self.auto_parallel ) and world_size != self.tp_size * self.pp_size: raise ValueError( f"world_size {world_size} should be equal to tp_size * pp_size {self.tp_size * self.pp_size} " "in non-auto_parallel mode.\n" "For non-auto-parallel mode, the world_size is not needed to set" ) @property def is_multi_gpu(self) -> bool: return self.world_size > 1 class _ModelFormatKind(Enum): HF = 0 TLLM_CKPT = 1 TLLM_ENGINE = 2 @dataclass class _ModelInfo: dtype: Optional[str] = None architecture: Optional[str] = None @property def model_name(self) -> str: if self.architecture is None: raise RuntimeError("The architecture is not set yet.") return self.architecture @classmethod def from_pretrained_config(cls, config: PretrainedConfig): return cls(dtype=config.dtype, architecture=config.architecture) @classmethod def from_builder_config_json(cls, config: dict): if 'version' in config: # The Dict format is { 'builder_config':..., 'plugin_config':...} dtype = config['plugin_config']['gpt_attention_plugin'] else: dtype = config['pretrained_config']['dtype'] return cls(dtype=dtype, architecture=config['builder_config']['name']) @classmethod def from_module(cls, module: Module): raise NotImplementedError() @dataclass class LlmArgs: ''' The arguments for constructing a LLM instance. Parameters: model (str or Path): The model name or a local model directory. Note that if the value could be both a model name or a local model directory, the local model directory will be prioritized. parallel_config (_ParallelConfig): The parallel configuration for the model. Default is an empty _ParallelConfig instance. tokenizer (str, Path, TokenizerBase, PreTrainedTokenizerBase, optional): The name or path of a HuggingFace Transformers tokenizer, or the loaded tokenizer. Default is None. skip_tokenizer_init (bool): If true, skip initialization of tokenizer and detokenizer. LLM.generate and LLM.generate_async will accept prompt token ids as input only. tokenizer_revision (str, optional): The revision of the tokenizer to use. Default is None. dtype (str, default="auto"): The data type for the model weights and activations. Can be "float16", "bfloat16", "float32", or "auto". If "auto", the data type will be automatically inferred from the source model. If the source data type is "float32", it will be converted to "float16". revision (str, optional): The revision of the model to use. Default is None. build_config (BuildConfig, default=BuildConfig()): The build configuration for the model. Default is an empty BuildConfig instance. quant_config (QuantConfig, default=QuantConfig()): The quantization configuration for the model. Default is an empty QuantConfig instance. embedding_parallel_mode (str, default="SHARDING_ALONG_VOCAB"): The parallel mode for embeddings. share_embedding_table (bool, default=False): Whether to share the embedding table. kv_cache_config (KvCacheConfig, optional): The key-value cache configuration for the model. Default is None. peft_cache_config (PeftCacheConfig, optional): The PEFT cache configuration for the model. Default is None. decoding_config (DecodingConfig, optional): The decoding configuration for the model. Default is None. logits_post_processor_map (Dict[str, Callable], optional): A map of logit post-processing functions. Default is None. scheduler_config (SchedulerConfig, default=SchedulerConfig()): The scheduler configuration for the model. Default is an empty SchedulerConfig instance. normalize_log_probs (bool, default=False): Whether to normalize log probabilities for the model. iter_stats_max_iterations (int, optional): The maximum number of iterations for iteration statistics. Default is None. request_stats_max_iterations (int, optional): The maximum number of iterations for request statistics. Default is None. batching_type (BatchingType, optional): The batching type for the model. Default is None. enable_build_cache (bool or BuildCacheConfig, optional): Whether to enable build caching for the model. Default is None. enable_tqdm (bool, default=False): Whether to display a progress bar during model building. ''' model: Union[str, Path] parallel_config: _ParallelConfig = field(default_factory=_ParallelConfig) tokenizer: Optional[Union[str, Path, TokenizerBase, PreTrainedTokenizerBase]] = None skip_tokenizer_init: bool = False tokenizer_revision: Optional[str] = None dtype: str = "auto" revision: Optional[str] = None # BuildConfig is introduced to give users a familiar interface to configure the model building. build_config: Optional[BuildConfig] = None quant_config: QuantConfig = field(default_factory=QuantConfig) # A handful of options from PretrainedConfig embedding_parallel_mode: str = 'SHARDING_ALONG_VOCAB' share_embedding_table: bool = False # Several options from ExecutorConfig, expanded here for less hierarchy kv_cache_config: Optional[KvCacheConfig] = None peft_cache_config: Optional[PeftCacheConfig] = None # TODO[enweiz]: this might affect medusa, and could be removed in the future for API consistency decoding_config: Optional[DecodingConfig] = None logits_post_processor_map: Optional[Dict[str, Callable]] = None scheduler_config: SchedulerConfig = field(default_factory=SchedulerConfig) normalize_log_probs: bool = False iter_stats_max_iterations: Optional[int] = None request_stats_max_iterations: Optional[int] = None batching_type: Optional[BatchingType] = None # Once set, the model will reuse the build_cache enable_build_cache: Union[BuildCacheConfig, bool] = False # Display the model building progress bar enable_tqdm: bool = False def __post_init__(self): # NOTE: this is only for the compatibility with the old API, and will be removed in the future # chunked context is disabled by default, and it is recommended to keep it enabled. # The underlying implementation might disable it if it is not supported. self.enable_chunked_context: bool = False if self.skip_tokenizer_init: self.tokenizer = None else: self.tokenizer = tokenizer_factory(self.tokenizer) self._engine_config: Optional[EngineConfig] = None self.auto_parallel_config = AutoParallelConfig( sharded_io_allowlist=[ "past_key_value_\\d+", "present_key_value_\\d*", ], same_buffer_io={ "past_key_value_(\\d+)": "present_key_value_\\1", }, **infer_cluster_config(), ) self.kv_cache_config = self.kv_cache_config or KvCacheConfig() # This is used to hold th options for convert_checkpoint self._convert_checkpoint_options = {} @classmethod def from_kwargs(cls, **kwargs) -> "LlmArgs": LlmArgs._check_executor_config_options_consistency() parallel_config = _ParallelConfig( tp_size=kwargs.pop('tensor_parallel_size', 1), pp_size=kwargs.pop('pipeline_parallel_size', 1), auto_parallel=kwargs.pop('auto_parallel', False), ) # world_size is only used for auto_parallel mode world_size = kwargs.pop('world_size', 1) if parallel_config.auto_parallel: parallel_config.world_size = world_size if devices := kwargs.pop('devices', None): parallel_config.devices = devices ret = cls(parallel_config=parallel_config, **kwargs) ret.setup() return ret @staticmethod def _check_executor_config_options_consistency(): # max_beam_width is not included since vague behavior due to lacking the support for dynamic beam width during # generation black_list = set(["max_beam_width"]) executor_config_attrs = set(attr for attr in dir(ExecutorConfig) if not attr.startswith('_') and callable(getattr(ExecutorConfig, attr))) executor_config_attrs -= black_list llm_args_attr = set([f.name for f in fields(LlmArgs)]) # NOTE: When cpp ExecutorConfig add new options, please add the new options into `_LlmArgs` with docs as well # ASK chunweiy for help if you are not sure about the new options. assert executor_config_attrs.issubset( llm_args_attr ), f"New options found in underlying ExecutorConfig: {llm_args_attr - executor_config_attrs}" def setup(self): ''' This method will setup the configs right before building the model. It will check the consistency of the configs and arbitrate the conflicts. ''' assert isinstance(self.model, (str, Path)), f"Invalid model: {self.model}" self._check_model_or_model_dir() self._setup_embedding_parallel_mode() if self.enable_build_cache: self.enable_build_cache = BuildCacheConfig() if isinstance( self.enable_build_cache, bool) else self.enable_build_cache if not isinstance(self.enable_build_cache, BuildCacheConfig): raise ValueError( f"Invalid build_cache_config: {self.enable_build_cache}") if self.is_local_model: # Load parallel_config from the engine. self.model_format = ModelLoader.get_model_format(self.model_dir) if self.model_format is _ModelFormatKind.TLLM_ENGINE: if self.build_config is not None: logger.warning( "The build_config is ignored for model format of TLLM_ENGINE." ) self._load_config_from_engine(Path(self.model_dir)) # Load parallel_config from the checkpoint. elif self.model_format is _ModelFormatKind.TLLM_CKPT: self._load_config_from_ckpt(Path(self.model_dir)) else: self.model_format = _ModelFormatKind.HF self.build_config = self.build_config or BuildConfig() self._config_arbitrator = _ConfigArbitrator() if self.build_config_mutable: if not self.build_config.max_num_tokens: self.build_config.max_num_tokens = 2048 if not GpuArch.is_post_ampere(): self._config_arbitrator.setup("pre-ampere not supported", config_name="plugin_config", use_paged_context_fmha=False) self._setup_enable_chunked_context() self._setup_enable_streaming_llm() self._setup_quant_config() if self.build_config.max_beam_width > 1: self._config_arbitrator.claim_func( "beam_search (beam_width > 1)", config_name="kv_cache_config", enable_block_reuse=False) else: self._setup_build_config_into_config_arbitrator() self._setup_kv_cache_config() self._config_arbitrator(plugin_config=self.build_config.plugin_config, kv_cache_config=self.kv_cache_config, build_config=self.build_config) def _check_model_or_model_dir(self): if not self.model: raise ValueError("model should be provided.") assert isinstance(self.model, (str, Path)), f"Invalid model: {self.model}" model_dir = Path(self.model) if model_dir.exists() and model_dir.is_dir(): self.model = model_dir @property def is_local_model(self) -> bool: return isinstance(self.model, Path) @property def is_hub_model(self) -> bool: return not self.is_local_model @property def model_dir(self) -> Path: assert self.is_local_model return self.model @property def build_config_mutable(self) -> bool: return self.model_format is not _ModelFormatKind.TLLM_ENGINE def _update_plugin_config(self, key: str, value: Any): setattr(self.build_config.plugin_config, key, value) def _load_config_from_engine(self, engine_dir: Path): engine_config = EngineConfig.from_json_file(engine_dir / "config.json") self._pretrained_config = engine_config.pretrained_config self.build_config = engine_config.build_config # load and check parallel_config mapping = self._pretrained_config.mapping if self.parallel_config.tp_size not in (1, mapping.tp_size): raise ValueError( f"tp_size {self.parallel_config.tp_size} is not consistent with the engine's tp_size {mapping.tp_size}" ) if self.parallel_config.pp_size not in (1, mapping.pp_size): raise ValueError( f"pp_size {self.parallel_config.pp_size} is not consistent with the engine's pp_size {mapping.pp_size}" ) self.parallel_config = _ParallelConfig( tp_size=mapping.tp_size, pp_size=mapping.pp_size, ) def _load_config_from_ckpt(self, ckpt_dir: Path): pretrained_config = PretrainedConfig.from_json_file(ckpt_dir / "config.json") tp_size = pretrained_config.mapping.tp_size pp_size = pretrained_config.mapping.pp_size world_size = pretrained_config.mapping.world_size # load parallel_config if self.parallel_config.tp_size != 1 and self.parallel_config.tp_size != tp_size: raise ValueError( f"tp_size {self.parallel_config.tp_size} is not consistent with the checkpoint's tp_size {tp_size}" ) if self.parallel_config.pp_size != 1 and self.parallel_config.pp_size != pp_size: raise ValueError( f"pp_size {self.parallel_config.pp_size} is not consistent with the checkpoint's pp_size {pp_size}" ) if (self.parallel_config.auto_parallel and self.parallel_config.world_size != 1 and world_size != 1): raise ValueError( f"auto parallel with world_size {self.parallel_config.world_size} does not support checkpoint with " "world_size {world_size} > 1") if not self.parallel_config.auto_parallel: self.parallel_config = _ParallelConfig( tp_size=tp_size, pp_size=pp_size, ) def _setup_embedding_parallel_mode(self): if self.embedding_parallel_mode == 'NONE': self._convert_checkpoint_options['use_parallel_embedding'] = False elif self.embedding_parallel_mode == 'SHARDING_ALONG_VOCAB': self._convert_checkpoint_options['use_parallel_embedding'] = True self._convert_checkpoint_options['embedding_sharding_dim'] = 0 elif self.embedding_parallel_mode == 'SHARDING_ALONG_HIDDEN': self._convert_checkpoint_options['use_parallel_embedding'] = True self._convert_checkpoint_options['embedding_sharding_dim'] = 1 else: raise ValueError( f"Invalid embedding_parallel_mode: {self.llm_args.embedding_parallel_mode}" ) self._convert_checkpoint_options[ 'share_embedding_table'] = self.share_embedding_table def _setup_build_config_into_config_arbitrator(self): # Setup the ConfigArbitrator with the plugin_config, the runtime configs such as KvCacheConfig should not be # conflict with it. build_config = asdict(self.build_config) del build_config['plugin_config'] self._config_arbitrator.setup("BuildConfig is readonly", config_name="build_config", **build_config) plugin_config = asdict(self.build_config.plugin_config) self._config_arbitrator.setup("PluginConfig is readonly", config_name="plugin_config", **plugin_config) def _setup_enable_chunked_context(self): def fallback(): logger.warning( f"Disabling chunked context due to configuration conflict.") self.enable_chunked_context = False if self.enable_chunked_context: if self.build_config_mutable: self._config_arbitrator.claim_perf("chunked_context", config_name="plugin_config", use_paged_context_fmha=True, fallback=fallback) def _setup_enable_streaming_llm(self): if self.build_config.plugin_config.streamingllm: self._validate_kv_cache_config() self._config_arbitrator.claim_func("streamingllm", config_name="plugin_config", streamingllm=True, use_paged_context_fmha=False) self._config_arbitrator.claim_func("streamingllm", config_name="kv_cache_config", enable_block_reuse=False) def _validate_kv_cache_config(self): if self.kv_cache_config is None: raise ValueError("KvCacheConfig is required for streaming LLM.") if self.kv_cache_config.max_attention_window is None: raise ValueError( "KvCacheConfig.max_attention_window should be set for streaming LLM." ) if self.kv_cache_config.max_attention_window <= 0: raise ValueError( "KvCacheConfig.max_attention_window should be greater than 0.") if self.kv_cache_config.sink_token_length is None: raise ValueError( "KvCacheConfig.sink_token_length should be set for streaming LLM." ) if self.kv_cache_config.sink_token_length <= 0: raise ValueError( "KvCacheConfig.sink_token_length should be greater than 0.") def _setup_kv_cache_config(self): assert self.kv_cache_config is not None if not GpuArch.is_post_ampere(): self._config_arbitrator.setup("pre-ampere not supported", config_name="kv_cache_config", enable_block_reuse=False) if self.kv_cache_config.enable_block_reuse: self._config_arbitrator.claim_func("enable_block_reuse", config_name="kv_cache_config", enable_block_reuse=True) self._config_arbitrator.claim_func("enable_block_reuse", config_name="plugin_config", use_paged_context_fmha=True) def _setup_quant_config(self): if self.quant_config.quant_algo is QuantAlgo.FP8: self._config_arbitrator.claim_func("fp8_quant", config_name="plugin_config", use_paged_context_fmha=False) def __setstate__(self, state): self.__dict__.update(state) def __getstate__(self): state = self.__dict__.copy() del state['_config_arbitrator'] return state class ConfigArbitrateError(Exception): ''' Exception raised when there is a conflict in configurations. ''' def __init__(self, message): super().__init__(message) class _ConfigArbitrator: ''' The ConfigArbitrator will arbitrate the options from different sources and raise errors if there are conflicts. ''' def __init__(self): # Dict of configs, the format is {config_name: {option: value}} self.virtual_configs: Dict[str, Dict[str, Any]] = {} # The claims for functionalities, the format is {config_name: [(func_name, {option: value})]} self.func_claims: Dict[str, List[Tuple[str, dict]]] = {} # The claims for performances, the format is {perf_name: [(config_name, {option: value}, fallback)]}, # the fallback is a callback function to be called when the performance is abandoned. self.perf_claims: Dict[str, List[Tuple[str, dict, Optional[Callable[[], None]]]]] = {} # Track where the option settings came from, this will be used for messages when encountered conflicts. # The format is {config_name: {option: error_information}} self.option_sources: Dict[str, Dict[str, str]] = {} def __call__(self, **configs) -> None: ''' Args: configs: name to config instance for each config need to be arbitrated. ''' self._arbitrate() # Apply the successfully arbitrated virtual configs to the real configs for name, config in configs.items(): if name in self.virtual_configs: virtual_config = self.virtual_configs[name] for option, value in virtual_config.items(): setattr(config, option, value) def setup(self, info: str, config_name: str, **kwargs): ''' Setup with some pre-defined configs comes from environment such as GPU arch. ''' config = self.virtual_configs.setdefault(config_name, {}) option_sources = self.option_sources.setdefault(config_name, {}) for option, value in kwargs.items(): assert config.get(option, value) == value config[option] = value option_sources[option] = info def claim_func(self, func: str, config_name: str, **options): ''' Claim a functionality demanding with configs and options. The functionality should be fulfilled, or errors will be raised. ''' claims = self.func_claims.setdefault(config_name, []) claims.append((func, options)) def claim_perf(self, perf: str, config_name: str, fallback: Optional[Callable[[], None]] = None, **options): ''' Claim a performance demanding for configs and options. The performance could be abandoned if the demanding is not available.''' claims = self.perf_claims.setdefault(perf, []) claims.append((config_name, options, fallback)) def _arbitrate(self): ''' Arbitrate the configs for all the functionalities and performances. ''' # Resolve functionality claims for config_name, funcs in self.func_claims.items(): virtual_config = self.virtual_configs.setdefault(config_name, {}) option_sources = self.option_sources.setdefault(config_name, {}) for func, options in funcs: for option, value in options.items(): if option in virtual_config: if virtual_config[option] != value: existing_func = option_sources[option] raise ConfigArbitrateError( f"Cannot set '{option}' to be '{value}' when enabling '{func}', " f"since '{existing_func}' has set it to be '{virtual_config[option]}'." ) else: virtual_config[option] = value # Track where the setting came from option_sources[option] = func # copy for restore # Resolve performance claims for perf, options in self.perf_claims.items(): option_sources = copy.copy(self.option_sources) virtual_configs = copy.copy(self.virtual_configs) restore = False for config_name, options, fallback in options: virtual_config = virtual_configs.setdefault(config_name, {}) option_source = option_sources.setdefault(config_name, {}) for option, value in options.items(): if option in virtual_config and virtual_config[ option] != value: logger.warning( f"Ignoring performance claim '{perf}' for option '{option}' due to conflict." ) restore = True else: virtual_config[option] = value option_source[option] = perf if restore: break if restore: if fallback: fallback() break if not restore: self.option_sources = option_sources self.virtual_configs = virtual_configs @dataclass class _ModelRuntimeContext: ''' _ModelRuntimeContext holds the minimum runtime resources for running a model. It could be a runtime cache in MPI nodes. ''' engine_buffer: Optional[trt.IHostMemory] = None tokenizer: Optional[TokenizerBase] = None # engine_config is only used for saving the engine to disk engine_config: Optional[Union[dict, EngineConfig]] = None mapping: Optional[Mapping] = None model_info: Optional[_ModelInfo] = None # This is only used when build-cache is enabled engine_path: Optional[str] = None @property def engine(self) -> trt.IHostMemory: assert self.engine_buffer is not None return self.engine_buffer @property def model_arch(self) -> str: # "LlaMACausalForLM" or "OPTForCausalLM" and so on return self.engine_config.pretrained_config['architecture'] class ModelLoader: ''' The ModelLoader is used to build an end-to-end model for a single-gpu. It accepts model name or a local model dir, and will download the model if necessary. ''' def __init__(self, llm_args: LlmArgs, tokenizer: Optional[TokenizerBase], workspace: Optional[str | tempfile.TemporaryDirectory] = None, llm_build_stats: Optional["LlmBuildStats"] = None): self.llm_args = llm_args self.tokenizer = tokenizer self._workspace = workspace or tempfile.TemporaryDirectory() self.llm_build_stats = llm_build_stats or LlmBuildStats() assert self.llm_args.build_config self.build_config = self.llm_args.build_config self.convert_checkpoint_options = self.llm_args._convert_checkpoint_options self.rank = mpi_rank() if llm_args.parallel_config.is_multi_gpu else 0 if llm_args.parallel_config.is_multi_gpu and not llm_args.parallel_config.auto_parallel: self.mapping = Mapping( tp_size=llm_args.parallel_config.tp_size, pp_size=llm_args.parallel_config.pp_size, rank=self.rank, world_size=llm_args.parallel_config.world_size, ) else: self.mapping = Mapping() self._build_pipeline = [] # For model from hub, the _model_dir is None, and will updated once downloaded self._model_dir: Optional[ Path] = self.llm_args.model_dir if self.llm_args.is_local_model else None self._model_info: Optional[_ModelInfo] = None self._model_name = self.llm_args.model self._model_format = self.llm_args.model_format self.auto_parallel_config = AutoParallelConfig( world_size=llm_args.parallel_config.world_size if llm_args. parallel_config.auto_parallel else 1) default_config = self.llm_args.auto_parallel_config self.auto_parallel_config.set_defaults( cluster_key=default_config.cluster_key, cluster_info=default_config.cluster_info, same_buffer_io=default_config.same_buffer_io, sharded_io_allowlist=default_config.sharded_io_allowlist, ) self._gather_build_steps() def _gather_build_steps(self): # Prepare the model processing pipeline if isinstance(self.llm_args.model, Module): # Build engine from user provided model self._build_pipeline.append( ("Build TensorRT-LLM engine", self._build_engine_from_inmemory_model)) return if self.llm_args.is_hub_model and self._model_format is not _ModelFormatKind.TLLM_ENGINE: # Download HF model if necessary if self.llm_args.model is None: raise ValueError( "Either model_dir or model should be provided to ModelConfig." ) self._build_pipeline.append( ("Downloading HF model", self._download_hf_model)) if self._model_format is _ModelFormatKind.HF: # HF -> TRT checkpoints -> engine self._build_pipeline.append( ("Loading HF model to memory", self._load_model_from_hf)) self._build_pipeline.append( ("Building TRT-LLM engine", self._build_engine)) elif self._model_format is _ModelFormatKind.TLLM_CKPT: # TRT checkpoints -> engine self._build_pipeline.append(("Loading TRT checkpoints to memory", self._load_model_from_ckpt)) self._build_pipeline.append( ("Build TRT-LLM engine", self._build_engine)) elif self._model_format is _ModelFormatKind.TLLM_ENGINE: # Nothing need to do pass else: raise ValueError(f"Unknown model format {self._model_format}") class BuildPipeline: def __init__(self, enable_tqdm: bool, labels: List[str], step_handlers: List[Callable], llm_build_stats: "LlmBuildStats"): assert len(labels) == len(step_handlers) self.labels = labels self.step_handlers = step_handlers self.llm_build_stats = llm_build_stats self.to_log = mpi_rank() == 0 self.counter = 0 self.progress_bar = tqdm( total=len(labels)) if enable_tqdm and self.to_log else None def __call__(self): start_time = time.time() for i in range(len(self.labels)): self.step_forward() if self.to_log: if self.progress_bar: self.progress_bar.close() else: overall_latency = time.time() - start_time print_colored("Loading model done.\n", 'bold_green') print_colored( 'Total latency: {:.3f}s\n'.format(overall_latency), 'grey') def step_forward(self): n_steps = len(self.labels) label = self.labels[self.counter] # display step information if self.to_log: if self.progress_bar: self.progress_bar.set_description(self.labels[self.counter]) else: print_colored("Loading Model: ") print_colored(f"[{self.counter+1}/{n_steps}]\t", 'bold_green') print_colored(f"{label}\n") # execute the step start_time = time.time() self.step_handlers[self.counter]() if self.progress_bar: self.progress_bar.update(1) latency = time.time() - start_time if self.to_log and not self.progress_bar: print_colored("Time: {:.3f}s\n".format(latency), 'grey') self.llm_build_stats.build_steps_info.append((label, latency)) self.counter += 1 def __call__(self, engine_dir: Optional[Path] = None) -> Path: ''' The engine_dir is the path to save the built engine. ''' if self.llm_args.model_format is _ModelFormatKind.TLLM_ENGINE: return self.llm_args.model_dir if self.llm_args.parallel_config.is_multi_gpu: torch.cuda.set_device(self.rank) len(self._build_pipeline) to_log = self.rank == 0 pipeline = ModelLoader.BuildPipeline( self.llm_args.enable_tqdm, [label for label, _ in self._build_pipeline], [handler for _, handler in self._build_pipeline], llm_build_stats=self.llm_build_stats, ) pipeline() if not hasattr(self, '_engine_config'): raise RuntimeError("config is not loaded.") config = self._engine_config assert engine_dir runtime_context = _ModelRuntimeContext( tokenizer=self.tokenizer, engine_buffer=self._engine_buffer, engine_config=config, mapping=self.mapping, model_info=self._model_info, ) ModelLoader.save(runtime_context, self.llm_args.model_dir, engine_dir) return engine_dir def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): for attr_name in dir(self): if not callable(getattr( self, attr_name)) and not attr_name.startswith("__"): if attr_name not in ('model_format', 'workspace'): setattr(self, attr_name, None) release_gc() @property def workspace(self) -> str: return self._workspace @property def model_format(self) -> _ModelFormatKind: return self._model_format @staticmethod def save( model: _ModelRuntimeContext, model_dir: str, engine_dir: str, ): ''' Save the built engine on a single GPU to the given path. ''' mapping = model.mapping rank = mapping.rank def copy_hf_tokenizer_data_to_engine_dir(): # Copy the HF tokenizer stuff to the engine dir so that we can use the engine dir as a standalone model dir # supports end-to-end task. # This is only for HF model for now, not available for users' customized tokenizers. import shutil for name in os.listdir(model_dir): src = os.path.join(model_dir, name) dst = os.path.join(engine_dir, name) if name.startswith('tokenizer'): src = os.path.realpath(src) if os.path.islink(src) else src if os.path.isdir(src): shutil.copytree(src, dst, dirs_exist_ok=True) else: shutil.copy2(src, dst) engine = Engine(config=model.engine_config, engine=model.engine) engine.save(engine_dir) if rank == 0: copy_hf_tokenizer_data_to_engine_dir() @staticmethod def get_model_format(model_dir: str) -> _ModelFormatKind: ''' Get the format of the model. ''' # TODO: migrate to detect version field in config.json after TRTLLM-256 finished if Path.exists( Path(model_dir) / 'config.json') and file_with_glob_exists( model_dir, 'rank*.safetensors'): return _ModelFormatKind.TLLM_CKPT if (Path.exists(Path(model_dir) / 'config.json') and (file_with_suffix_exists(model_dir, '.bin') or file_with_suffix_exists(model_dir, '.safetensors'))): return _ModelFormatKind.HF if Path.exists( Path(model_dir) / 'config.json') and file_with_suffix_exists( model_dir, '.engine'): return _ModelFormatKind.TLLM_ENGINE raise ValueError(f"Unknown model format for {model_dir}") def _download_hf_model(self): ''' Download HF model from third-party model hub like www.modelscope.cn or huggingface. ''' model_dir = None # Only the rank0 are allowed to download model if mpi_rank() == 0: assert self._workspace is not None assert isinstance(self.llm_args.model, str) # this will download only once when multiple MPI processes are running model_dir = download_hf_model(self.llm_args.model, revision=self.llm_args.revision) # Make all the processes got the same model_dir self._model_dir = mpi_broadcast(model_dir, root=0) self.llm_args.model = Path(self._model_dir) # mark as a local model print_colored(f"Downloaded model to {self._model_dir}\n", 'grey') def _load_model_from_hf(self): ''' Load a TRT-LLM model from a HF model. ''' assert self._model_dir is not None import transformers hf_config = transformers.AutoConfig.from_pretrained(self._model_dir) architecture = hf_config.architectures[0] if architecture not in MODEL_MAP: raise KeyError(f"Unsupported model architecture: {architecture}") model_cls = MODEL_MAP[architecture] if TopModelMixin.__name__ in model_cls.from_hugging_face.__qualname__: raise NotImplementedError( f"Unsupported model architecture in HLAPI: {architecture}") if self.llm_args.quant_config.quant_mode.has_any_quant(): assert self.workspace is not None checkpoint_dir = f"{self.workspace}/quantized-checkpoint" if self.rank == 0: model_cls.quantize( self._model_dir, checkpoint_dir, dtype=self.llm_args.dtype, mapping=self.mapping, quant_config=self.llm_args.quant_config, ) if self.llm_args.parallel_config.is_multi_gpu: mpi_barrier() self.model = model_cls.from_checkpoint(checkpoint_dir, rank=self.mapping.rank) else: self.model = model_cls.from_hugging_face( str(self._model_dir), dtype=self.llm_args.dtype, mapping=self.mapping, quant_config=self.llm_args.quant_config, load_model_on_cpu= True, # TODO:TRTLLM-195 to enhance the weights loading memory usage and chose best location **self.convert_checkpoint_options, ) self.pretrained_config = self.model.config self._model_info = _ModelInfo.from_pretrained_config( self.pretrained_config) def _load_model_from_ckpt(self): ''' Load a TRT-LLM model from checkpoint. ''' self.pretrained_config = PretrainedConfig.from_json_file( os.path.join(self._model_dir, 'config.json')) self.pretrained_config.mapping = self.mapping architecture = self.pretrained_config.architecture assert architecture in MODEL_MAP, \ f"Unsupported model architecture: {architecture}" model_cls = MODEL_MAP[architecture] self.model = model_cls.from_checkpoint(self._model_dir, config=self.pretrained_config) self._model_info = _ModelInfo.from_pretrained_config( self.pretrained_config) # load embedding sharing related options self.convert_checkpoint_options[ 'share_embedding_table'] = self.pretrained_config.share_embedding_table self.convert_checkpoint_options[ 'use_parallel_embedding'] = self.pretrained_config.use_parallel_embedding def _build_engine_from_inmemory_model(self): assert isinstance(self.llm_args.model, Module) self._model_info = _ModelInfo.from_module(self.model) def _build_engine(self): self.build_config.update(auto_parallel_config=self.auto_parallel_config) if self.auto_parallel_config.enabled: self.model.config.mapping.rank = self.rank assert self.model is not None, "model is loaded yet." assert isinstance( self.build_config, BuildConfig), f"build_config is not set yet: {self.build_config}" engine = build(self.model, self.build_config) self._engine_buffer = engine.engine self._engine_config = engine.config self.mapping = self.model.config.mapping # delete the model explicitly to free all the build-time resources self.model = None def _save_engine_for_runtime(self): ''' Persist the engine to disk for the cpp runtime. Currently, the cpp runtime can accept an engine path, that requires the engine should always be saved to disk. This explicit saving will be removed in the future when the cpp runtime can accept the engine buffer directly. But this is necessary for a build cache, but it can be optimized to async IO. ''' if self.build_cache_enabled: self._model_dir = self.engine_cache_stage.cache_dir self._model_format = _ModelFormatKind.TLLM_ENGINE return def _load_engine_buffer(self): # Load engine buffer from disk engine = Engine.from_dir(self._model_dir) self._engine_buffer = engine.engine self._engine_config = engine.config @staticmethod def load_extra_build_configs_from_engine( model_dir: str) -> Optional[Namespace]: ''' Load the extra build configs from the engine directory, return None if model isn't an engine. ''' if ModelLoader.get_model_format( model_dir) is not _ModelFormatKind.TLLM_ENGINE: return None with open(Path(model_dir) / "config.json", "r") as f: engine_config = json.load(f) build_config = engine_config['build_config'] build_config.pop("plugin_config") return Namespace(**build_config) @staticmethod def load_hf_tokenizer(model_dir) -> Optional[TransformersTokenizer]: try: return TransformersTokenizer.from_pretrained(model_dir, legacy=False, padding_side='left', truncation_side='left', trust_remote_code=True, use_fast=True) except: return None class CachedModelLoader: ''' The CacheModelLoader is used to build the model in both single or multi-gpu, with cache might be enabled. ''' def __init__( self, llm_args: LlmArgs, llm_build_stats: "LlmBuildStats", mpi_session: Optional[MpiSession] = None, workspace: Optional[str] = None, ): self.llm_args = llm_args self.mpi_session = mpi_session self._workspace = workspace or tempfile.TemporaryDirectory() self.llm_build_stats = llm_build_stats # This is used for build cache. To compute the cache key, a local HF model is required, it could be download # from HF model hub, so this helps to hold the path. self._hf_model_dir: Optional[Path] = None @property def workspace(self) -> Path: return Path(self._workspace.name) if isinstance( self._workspace, tempfile.TemporaryDirectory) else Path( self._workspace) def __call__(self) -> Path: if self.llm_args.model_format is _ModelFormatKind.TLLM_ENGINE: # do nothing for engine input return self.llm_args.model_dir self.engine_cache_stage: Optional[CachedStage] = None if self.build_cache_enabled: if self.llm_args.is_hub_model: # This will download the config.json from HF model hub, this helps to create a PretrainedConfig for # cache key. self._hf_model_dir = download_hf_pretrained_config( self.llm_args.model, revision=self.llm_args.revision) elif self.llm_args.is_local_model: self._hf_model_dir = self.llm_args.model_dir if self.llm_args.model_format is _ModelFormatKind.HF else None self.engine_cache_stage = self._get_engine_cache_stage() if self.engine_cache_stage.cache_hitted(): print_colored( f"Reusing cached engine in {self.engine_cache_stage.get_engine_path()}\n\n", 'grey') self.llm_build_stats.cache_hitted = True self.llm_args.model = self.engine_cache_stage.get_engine_path() self.llm_build_stats.engine_dir = self.llm_args.model_dir return self.llm_build_stats.engine_dir return self._build_model() def get_engine_dir(self) -> Path: if self.llm_args.model_format is _ModelFormatKind.TLLM_ENGINE: return self.llm_args.model_dir # generate a new path for writing the engine if self.build_cache_enabled: cache_stage = self._get_engine_cache_stage() return cache_stage.get_engine_path() return self.workspace / "tmp.engine" @property def build_cache_enabled(self) -> bool: _enable_build_cache, _ = get_build_cache_config_from_env() return (self.llm_args.enable_build_cache or _enable_build_cache) and ( self.llm_args.model_format is _ModelFormatKind.HF) def _get_engine_cache_stage(self) -> CachedStage: ''' Get the cache stage for engine building. ''' build_cache = BuildCache(self.llm_args.enable_build_cache) assert self._hf_model_dir is not None, "HF model dir is required for cache key." dummy_build_config = CachedModelLoader.get_final_build_config( self.llm_args, self._hf_model_dir) return build_cache.get_engine_building_cache_stage( build_config=dummy_build_config, model_path=self._hf_model_dir, # for PretrainedConfig parallel_config=self.llm_args.parallel_config, # Other configs affecting the engine building quant_config=self.llm_args.quant_config) @staticmethod def get_final_build_config(llm_args: LlmArgs, model_dir: Path) -> BuildConfig: ''' Get the build_config for cache key. The tricky part is that, the build_config will be altered in `build()`, but we need a final version of build_config before `build()` is called for cache key. Args: llm_args: The LlmArgs for building the model. model_dir: The path to the local HF model. ''' # This is only needed by BuildCache for cache key # The build() doesn't need the real model instance to get a updated BuildConig. What is really needed is the # dtype. That's why the model will be downloaded from HF if necessary to get the accurate dtype. import transformers hf_config = transformers.AutoConfig.from_pretrained(model_dir) architecture = hf_config.architectures[0] if architecture not in MODEL_MAP: raise KeyError(f"Unsupported model architecture: {architecture}") model_cls = MODEL_MAP[architecture] config_cls = model_cls.config_class pretrained_config = config_cls.from_hugging_face( model_dir, mapping=Mapping(world_size=llm_args.parallel_config.world_size, tp_size=llm_args.parallel_config.tp_size, pp_size=llm_args.parallel_config.pp_size), quant_config=llm_args.quant_config, dtype=llm_args.dtype) @dataclass class DummyModel: # This is only used for getting the updated BuildConfig from build() without actually loading the whole # pretrained model to save overhead and memory. config: PretrainedConfig # dry_run to get the updated build_config for cache key. The build_config is modified within build(), so using # a build_config before build() is not correct for cache key, so we need to get the build_config after build() # in dry_run mode. dummy_model = DummyModel(pretrained_config) dummy_build_config = copy.copy(llm_args.build_config) dummy_build_config.dry_run = True updated_build_config = build(dummy_model, dummy_build_config, return_build_config=True) return updated_build_config def _build_model(self) -> Path: model_format = self.llm_args.model_format def build_task(): if model_format is not _ModelFormatKind.TLLM_ENGINE: if self.llm_args.parallel_config.is_multi_gpu: assert self.mpi_session # The engine_dir:Path will be stored to MPINodeState.state build_infos = self.mpi_session.submit_sync( CachedModelLoader._node_build_task, llm_args=self.llm_args, tokenizer=self.llm_args. tokenizer, # TODO[chunweiy]: Use llm_args directly dtype=self.llm_args.dtype, engine_dir=self.get_engine_dir()) self.llm_build_stats.build_steps_info = build_infos[0] else: # single-gpu with ModelLoader(self.llm_args, tokenizer=self.llm_args.tokenizer, workspace=self.workspace.name, llm_build_stats=self.llm_build_stats ) as model_loader: model_loader(self.get_engine_dir()) release_gc() if self.build_cache_enabled: with self.engine_cache_stage.write_guard(): build_task() return self.get_engine_dir() else: build_task() return self.get_engine_dir() @print_traceback_on_error @staticmethod def _node_build_task( llm_args: LlmArgs, tokenizer: Optional[TokenizerBase] = None, dtype: str = 'auto', engine_dir: Optional[Path] = None, ): if MPINodeState.is_initialized(): raise RuntimeError("The MPI node is already initialized.") with ModelLoader( llm_args, tokenizer=tokenizer, ) as model_loader: model_loader(engine_dir=engine_dir) return model_loader.llm_build_stats.build_steps_info def save(self, engine_dir: Path): # copy the engine directory to the target directory shutil.copytree(self.get_engine_dir(), engine_dir) @dataclass class LlmBuildStats: ''' LlmBuildStats is the statistics for the LLM model building. ''' # Whether the cache is hitted for the engine cache_hitted: bool = False model_from_hf_hub: bool = False local_model_dir: Optional[Path] = None # The path to the trt-llm engine engine_dir: Optional[Path] = None # The build steps information, including the step name and the latency in seconds. build_steps_info: List[Tuple[str, float]] = field(default_factory=list)