Source code for tensorrt_llm.hlapi.llm_utils

__all__ = [
    'LlmArgs',
    'LlmBuildStats',
    'ModelLoader',
    '_ModelRuntimeContext',
    '_ModelInfo',
    '_ParallelConfig',
    '_ModelFormatKind',
    'BatchingType',
    'ExecutorConfig',
    'SchedulerConfig',
    'KvCacheConfig',
    'ContextChunkingPolicy',
    'CapacitySchedulerPolicy',
    'BuildConfig',
    'BuildCacheConfig',
    'QuantConfig',
    'CalibConfig',
    '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, Literal, Optional, Tuple, Union

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.automodel import MODEL_MAP, AutoConfig, AutoModelForCausalLM
from ..models.modeling_utils import PretrainedConfig, QuantAlgo, QuantConfig
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,
                    get_directory_size_in_gb, print_colored,
                    print_traceback_on_error, set_docstring)


@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


[docs] @dataclass(slots=True) class CalibConfig: """ Calibration configuration. Args: device (Literal['cuda', 'cpu'], default='cuda'): The device to run calibration. calib_dataset (str, default='cnn_dailymail'): The name or local path of calibration dataset. calib_batches (int, default=512): The number of batches that the calibration runs. calib_batch_size (int, default=1): The batch size that the calibration runs. calib_max_seq_length (int, default=512): The maximum sequence length that the calibration runs. random_seed (int, default=1234): The random seed used for calibration. tokenizer_max_seq_length (int, default=2048): The maximum sequence length to initialize tokenizer for calibration. """ device: Literal['cuda', 'cpu'] = 'cuda' calib_dataset: str = 'cnn_dailymail' calib_batches: int = 512 calib_batch_size: int = 1 calib_max_seq_length: int = 512 random_seed: int = 1234 tokenizer_max_seq_length: int = 2048
[docs] @classmethod def from_dict(cls, config: dict): return cls(**config)
[docs] def to_dict(self): return asdict(self)
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() LLMARGS_STAET_DOCSTRING = "The arguments for constructing a LLM instance.\n\nParameters:\n" # The arguments locate in LLM class's explicit arg-list, these will not be included in LLM class's apidocs. LLMARGS_EXPLICIT_ARGS_DOCSTRING = r""" 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. load_format (Literal['auto', 'dummy'], default='auto'): The format of the model weights to load. * 'auto' will try to load the weights from the provided checkpoint. * 'dummy' will initialize the weights with random values, which is mainly for profiling. """ # The arguments locate in LLM class's kwargs, and will be concatenated to LLM class's apidocs. # The parallel_config is replaced by {auto_parallel, pipeline_parallel_size} arguments, the tensor_parallel_size is # already in the LLM class's apidocs, so it is not included here. LLMARGS_REMAINING_ARGS_DOCSTRING = r""" auto_parallel (bool, default=False): Enable auto parallel mode. pipeline_parallel_size (int, default=1): The pipeline parallel size. enable_lora (bool, default=False): Enable LoRA adapters. max_lora_rank (int, default=None): Maximum LoRA rank. If specified, it overrides `build_config.lora_config.max_lora_rank`. max_loras (int, default=4): Maximum number of LoRA adapters to be stored in GPU memory. max_cpu_loras (int, default=4): Maximum number of LoRA adapters to be stored in CPU memory. 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. calib_config (CalibConfig, default=CalibConfig()): The calibration configuration for the model. 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. trust_remote_code (bool, default=False): Whether to trust remote code when downloading model and tokenizer from Hugging Face. """ @set_docstring(LLMARGS_STAET_DOCSTRING + LLMARGS_EXPLICIT_ARGS_DOCSTRING + LLMARGS_REMAINING_ARGS_DOCSTRING) @dataclass class LlmArgs: 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 load_format: Literal['auto', 'dummy'] = 'auto' # LoRA arguments enable_lora: bool = False max_lora_rank: Optional[int] = None max_loras: int = 4 max_cpu_loras: int = 4 # BuildConfig is introduced to give users a familiar interface to configure the model building. build_config: Optional[BuildConfig] = None fast_build: Optional[bool] = False quant_config: QuantConfig = field(default_factory=QuantConfig) calib_config: CalibConfig = field(default_factory=CalibConfig) # 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 trust_remote_code: 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 # TODO[chunweiy]: Enable this option in the future # Currently we want HLAPI to be consistent with the lower APIs in the model building, thus disable this to avoid # magics. self.perform_config_arbitration = False if self.skip_tokenizer_init: self.tokenizer = None else: self.tokenizer = tokenizer_factory(self.tokenizer) if torch.cuda.get_device_properties(0).major < 8: if self.dtype == 'auto': self.dtype = 'float16' if self.dtype == 'bfloat16': raise RuntimeError("Pre SM 80 GPUs do not support bfloat16") 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() # TODO(xiweny): remove the checker when manage weights support all data types if self.fast_build and (self.quant_config.quant_algo is QuantAlgo.FP8 or self.quant_config.quant_algo is None): self._update_plugin_config("manage_weights", True) if self.enable_lora: self.build_config.plugin_config.lora_plugin = 'auto' if self.max_lora_rank is not None: self.build_config.lora_config.max_lora_rank = self.max_lora_rank if self.perform_config_arbitration: self._perform_config_arbitration() def _perform_config_arbitration(self): ''' Arbitrate the configurations for the model building. The configs between different functional or performance features might be conflicted, and this method will arbitrate the conflicts and raise errors if necessary. ''' 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) self._config_arbitrator = None 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, f"model_dir is only available for local model, {self.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 any(i <= 0 for i in self.kv_cache_config.max_attention_window): raise ValueError( "Elements in 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() if '_config_arbitrator' in state: 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: Optional[Engine] = 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 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, workspace: Optional[str | tempfile.TemporaryDirectory] = None, llm_build_stats: Optional["LlmBuildStats"] = None): self.llm_args = llm_args 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) 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() assert engine_dir runtime_context = _ModelRuntimeContext( engine=self._engine, 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) model.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. ''' if not (Path(model_dir) / 'config.json').exists(): raise ValueError( f"Failed to infer model format because no config.json exists in {model_dir}" ) with open(Path(model_dir) / 'config.json') as f: config = json.load(f) try: if 'pretrained_config' in config and 'build_config' in config: model_format = _ModelFormatKind.TLLM_ENGINE EngineConfig.from_json_file(Path(model_dir) / 'config.json') elif 'architecture' in config and 'dtype' in config: model_format = _ModelFormatKind.TLLM_CKPT PretrainedConfig.from_checkpoint(model_dir) else: model_format = _ModelFormatKind.HF AutoConfig.from_hugging_face(model_dir) except Exception as e: raise ValueError( f"Inferred model format {model_format}, but failed to load config.json: {e}" ) else: return model_format 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) print_colored(f"Downloaded model to {model_dir}\n", 'grey') # 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 assert self.llm_args.is_local_model def _load_model_from_hf(self): ''' Load a TRT-LLM model from a HF model. ''' assert self._model_dir is not None model_cls = AutoModelForCausalLM.get_trtllm_model_class( self._model_dir, self.llm_args.trust_remote_code) if self.llm_args.load_format == 'dummy': config = model_cls.config_class.from_hugging_face( str(self._model_dir), dtype=self.llm_args.dtype, mapping=self.mapping, quant_config=self.llm_args.quant_config, **self.convert_checkpoint_options, ) self.model = model_cls(config) elif self.llm_args.quant_config.requires_calibration: 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, **self.llm_args.calib_config.to_dict(), trust_remote_code=self.llm_args.trust_remote_code, ) 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 trust_remote_code=self.llm_args.trust_remote_code, **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 #TODO: TRTLLM-1091, change the architecture in the checkpoint to TRT-LLM one, not HF one. architecture = self.pretrained_config.architecture assert architecture in MODEL_MAP, \ f"Unsupported model architecture: {architecture}" model_cls = MODEL_MAP[architecture] if self.llm_args.load_format == 'dummy': self.model = model_cls(self.pretrained_config) else: 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): assert isinstance( self.build_config, BuildConfig), f"build_config is not set yet: {self.build_config}" # avoid the original build_config is modified, avoid the side effect copied_build_config = copy.deepcopy(self.build_config) copied_build_config.update( auto_parallel_config=self.auto_parallel_config) copied_build_config.update_kv_cache_type(self._model_info.architecture) if self.auto_parallel_config.enabled: self.model.config.mapping.rank = self.rank assert self.model is not None, "model is loaded yet." self._engine = build(self.model, copied_build_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 self._engine = Engine.from_dir(self._model_dir) @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, trust_remote_code) -> Optional[TransformersTokenizer]: try: return TransformersTokenizer.from_pretrained( model_dir, legacy=False, padding_side='left', truncation_side='left', trust_remote_code=trust_remote_code, use_fast=True) except Exception as e: logger.error(f"Failed to load tokenizer from {model_dir}: {e}") 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.is_cached(): self.llm_build_stats.cache_hitted = True print_colored( f"Reusing cached engine in {self.engine_cache_stage.get_engine_path()}\n\n", 'grey') 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 ) and not self.llm_args.parallel_config.auto_parallel 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." def serialize(d) -> str: dic = asdict(d) if not isinstance( d, PretrainedConfig) else d.to_dict() return json.dumps(dic, sort_keys=True) parallel_config = self.llm_args.parallel_config force_rebuild = False if parallel_config.auto_parallel: force_rebuild = True if self.llm_args.model_format is not _ModelFormatKind.HF: force_rebuild = True return build_cache.get_engine_building_cache_stage( build_config=self.llm_args.build_config, model_path=self._hf_model_dir, force_rebuild=force_rebuild, # Other configs affecting the engine building parallel_config=serialize(parallel_config), pretrained_config=serialize(self.get_pretrained_config()), quant_config=serialize(self.llm_args.quant_config), ) def get_pretrained_config(self) -> PretrainedConfig: ''' Get the PretrainedConfig for cache key. NOTE, this is not the HF model's config, but the TRT-LLM's config. We use this as a generic information for HF and other models. ''' assert self._hf_model_dir is not None return AutoConfig.from_hugging_face( self._hf_model_dir, mapping=Mapping(world_size=self.llm_args.parallel_config.world_size, tp_size=self.llm_args.parallel_config.tp_size, pp_size=self.llm_args.parallel_config.pp_size), quant_config=self.llm_args.quant_config, dtype=self.llm_args.dtype) def _build_model(self) -> Path: model_format = self.llm_args.model_format def build_task(engine_dir: Path): if model_format is not _ModelFormatKind.TLLM_ENGINE: model_loader_kwargs = { 'llm_args': self.llm_args, 'workspace': str(self.workspace), 'llm_build_stats': self.llm_build_stats, } 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, engine_dir=engine_dir, **model_loader_kwargs) self.llm_build_stats.build_steps_info = build_infos[0] else: # single-gpu with ModelLoader(**model_loader_kwargs) as model_loader: model_loader(engine_dir=engine_dir) release_gc() has_storage = True if self.build_cache_enabled: try: # TODO[chunweiy]: Cover the case when the model is from HF model hub. if self.llm_args.is_local_model: # This is not perfect, but will make build-cache much more robust. has_storage = self.engine_cache_stage.parent.free_storage_in_gb( ) >= get_directory_size_in_gb(self.llm_args.model_dir) except ValueError: has_storage = False except Exception as e: logger.error(e) has_storage = False if has_storage: with self.engine_cache_stage.write_guard() as engine_dir: build_task(engine_dir) self.llm_build_stats.cache_hitted = True else: print_colored( "The cache directory is too small, build-cache is disabled.\n", 'grey') self.llm_build_stats.cache_hitted = False self.llm_build_stats.cache_info = "The cache root directory is too small." if not (has_storage and self.build_cache_enabled): build_task(self.get_engine_dir()) return self.get_engine_dir() @print_traceback_on_error @staticmethod def _node_build_task( llm_args: LlmArgs, workspace: Optional[str | tempfile.TemporaryDirectory] = None, llm_build_stats: Optional['LlmBuildStats'] = None, engine_dir: Optional[Path] = None, ): if MPINodeState.is_initialized(): raise RuntimeError("The MPI node is already initialized.") with ModelLoader(llm_args, workspace=workspace, llm_build_stats=llm_build_stats) 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 hit for the engine cache_hitted: bool = False cache_info: Optional[str] = None 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)