diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py index 3ac76d6809be..35d1c892d997 100644 --- a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod -from typing import Callable, List, Union +from collections.abc import Callable, Mapping +import torch from torch import nn from tensorrt_llm._torch.model_config import ModelConfig @@ -8,17 +9,56 @@ class BaseWeightMapper(ABC): + """Helper for loading weights to each child module. + + A typical weight loader function walks `model.named_modules()` and loads + weights for each child module. Although called `WeightMapper`, this class is + called at multiple locations of the walk to do: + - Weight dict preprocessing + - Weight dict key renaming + - Finding modules requiring special weight handling and calling weight processing + hooks before passing them to the module + - Finding really special modules and have complete custom methods to handle their + weight loading. + - Basic weight name prefix removing and weight copying onto PyTorch module + Parameter. + + Subclasses of `BaseWeightMapper` can be selected by checkpoint format and model + architecture through `AutoCheckpointMapper`. + + Abstract methods for subclasses to implement: + - map_weights + - apply_callbacks + """ def __init__(self): self._callbacks: list[Callable] = [] - self._mapping: dict = {} - self._skip_modules = [] - self._model: Union[nn.Module, DecoderModelForCausalLM] | None = None + # Mapping for modules that need special weight loading, like fusing + # several weights. + # It maps module names to the corresponding source names in the checkpoint, e.g. + # 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'], + # 'gate_up_proj': ['gate_proj', 'up_proj'] + # It will be initialized in `self.map_weights()` and queried in + # `self.does_require_special_handling()` + self._mapping: dict[str, list[str]] = {} + self._skip_modules: list[str] = [] + self._model: nn.Module | DecoderModelForCausalLM | None = None self._config: ModelConfig | None = None - def init_model_and_config(self, model: Union[nn.Module, - DecoderModelForCausalLM], - config: ModelConfig): + def init_model_and_config(self, model: nn.Module | DecoderModelForCausalLM, + config: ModelConfig) -> None: + """Bind this mapper to the LLM class instance and model config. + + Called once after the model is constructed and before weight loading + begins. It validates the model has the attributes needed by the mapper, + records the tensor parallel size, and calls `map_weights` so subclasses + can populate fused-module mappings. + + Args + - model: nn.Module | DecoderModelForCausalLM, LLM class instance whose + child modules will be loaded. + - config: ModelConfig, loaded model config used by mapper decisions. + """ self._model = model self._config = config @@ -38,46 +78,68 @@ def cleanup(self) -> None: @abstractmethod def map_weights(self) -> None: - """ - Maps weights from TRT-LLM to a source state dictionary (e.g., Hugging Face) + """Initialize mapping for modules that need special weight loading like weight fusion. + + This function is called inside `self.init_model_and_config()`. Derived classes implement + this function to initialize `self.mapping`, which maps special module names to the + corresponding source names in the checkpoint, e.g. + - 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'] + - 'gate_up_proj': ['gate_proj', 'up_proj'] """ @abstractmethod - def apply_callbacks(self, module: nn.Module, module_name: str, - module_names_breakdown: list[str], - weights: dict) -> list[dict]: - """ - Applies a series of transformation functions to an internal representation - of weights or to guide the mapping process. The exact behavior might depend - on the implementation (e.g., storing callbacks to be applied later). - - Args: - module: The module to apply the callbacks to - module_name: The specific module name (e.g., 'qkv_proj', 'gate_up_proj') - module_names_breakdown: List of module path components for building full paths - weights: The weights dictionary to process + def apply_callbacks( + self, module: nn.Module, module_name: str, + module_names_breakdown: list[str], + weights: Mapping[str, + torch.Tensor]) -> list[dict[str, torch.Tensor]]: + """Build processed weight dicts for a special child module. + + Used only when `self.does_require_special_handling()` is True, derived classes + implement this function to process raw `weights` before passing them into + the `module.load_weights()`. + + Example special module: qkv_proj that combines q_proj, k_proj and v_proj. + + Args + - module: nn.Module, must have `module.load_weights()` to receive the + returned weight dicts. + - module_name: str, final component of the child module name, such as + `qkv_proj` or `gate_up_proj`. + - module_names_breakdown: list[str], parent module path split on `.`, + needed for finding weights for this child module. + - weights: Mapping[str, torch.Tensor], full checkpoint weight dict. + + Returns + - module_weights: list[Mapping[str, torch.Tensor]], list of weight dicts + to pass to `module.load_weights()`. """ - def rename_by_params_map(self, params_map: dict[str, str], - weights: dict) -> dict: - """ - Rename weight keys according to regex pattern matching. - - Args: - pattern_mapping: A dictionary mapping regex patterns to replacement strings. The key is HF name pattern, and the value is corresponding TRT-LLM name pattern. - The patterns will be used to match keys in the weights dict and replace - them according to the replacement string, which can use regex backreferences. - Example: - HF name: vision_model.encoder.layers.1.self_attn.out_proj.{weight,bias} - TRT-LLM name: vision_model.encoder.layers.1.self_attn.o_proj.{weight,bias} - Then the pattern_mapping could be: - pattern_mapping = { - r'(.*?)out_proj(.*)': r'\1o_proj\2' - } - weights: A dictionary of weights (or ConsumableWeightsDict) - - Returns: - A dictionary of weights with renamed keys (preserves ConsumableWeightsDict if input was one) + def rename_by_params_map( + self, params_map: dict[str, str], + weights: Mapping[str, torch.Tensor]) -> Mapping[str, torch.Tensor]: + """Rename checkpoint keys in `weights` with string rules defined in `params_map`. + + This should be called at beginning of a weight loading function to rename + input weights. + + The basic implementation is regex replacement rule but derived classes of + `BaseWeightMapper` may change that. + + Regex rule example: `r'(.*?)out_proj(.*)' -> r'\\1o_proj\\2'` maps + `vision_model.encoder.layers.1.self_attn.out_proj.weight` to + `vision_model.encoder.layers.1.self_attn.o_proj.weight`. + + Args + - params_map: dict[str, str], regex pattern to replacement string. + Replacement strings may use regex backreferences. + - weights: Mapping[str, torch.Tensor], checkpoint/state-dict weight + tensors keyed by checkpoint name. + + Returns + - renamed_weights: Mapping[str, torch.Tensor], weight dict with renamed + keys and unchanged tensor values. If the input `weights` is a + `ConsumableWeightsDict`, the returned object preserved that type. """ import re @@ -114,49 +176,102 @@ def rename_by_params_map(self, params_map: dict[str, str], return ConsumableWeightsDict(renamed_weights) return renamed_weights - def preprocess_weights(self, weights: dict) -> dict: - """ - Preprocess weights before starting the loading process. + def preprocess_weights( + self, weights: Mapping[str, + torch.Tensor]) -> Mapping[str, torch.Tensor]: + """Rewrite a full checkpoint weight dict before module walking starts. + + If simple string renaming rules used in `rename_by_params_map` does not satisfy the + need, call this function for a more custom rewrite of the checkpoint weight dict. + + Args + - weights: Mapping[str, torch.Tensor], full checkpoint/state-dict weight + tensors keyed by checkpoint name. + + Returns + - weights: Mapping[str, torch.Tensor], preprocessed weight dict to pass to + the child-module loader. """ ... def handle_manual_copy(self, module_name: str, - module_weights: dict, + module_weights: Mapping[str, torch.Tensor], n: str, p: nn.Parameter, allow_partial_loading: bool = False) -> None: + """Copy one parameter for a module that has no `load_weights` method. + + Args + - module_name: str, final component of the child module name. + - module_weights: Mapping[str, torch.Tensor], tensors for this child module + with the module prefix removed. + - n: str, parameter name inside the child module, such as `weight` or + `bias`. + - p: nn.Parameter, destination parameter to update. + - allow_partial_loading: bool, if `True`, skip missing parameters; if + `False`, assert that `n` exists in `module_weights`. + """ if not allow_partial_loading: assert n in module_weights if n in module_weights: p.data.copy_(module_weights[n][:]) def does_require_special_handling(self, module_name: str) -> bool: + """ + Whether a module requires special weight loading like fusing weights. + Examples include module 'qkv_proj' fuses weights 'q_proj', 'k_proj' and 'v_proj'. + """ return module_name in self.mapping def is_special_instance_module(self, module: nn.Module) -> bool: + """If the module is special enough for a complete custom weight handling hook. + + If true, the weight loader function should call + `self.handle_special_instance_module()` next to handle that special module. + """ return False def handle_special_instance_module( self, module: nn.Module, module_name: str, - module_weights: dict, + module_weights: Mapping[str, torch.Tensor], allow_partial_loading: bool = False) -> None: + """Load weights for a special module that needs custom behavior. + + Only call this if `self.is_special_instance_module()` returns true. + Subclasses opt into this path by overriding `is_special_instance_module` and + `is_special_instance_module`. This hook is for special modules who cannot be + handled by `does_require_special_handling()` and `apply_callbacks()`. + + Args + - module: nn.Module, special module to load weights. + - module_name: str, final component of the child module name. + - module_weights: Mapping[str, torch.Tensor], tensors for this child module + with the module prefix removed. + - allow_partial_loading: bool, if `True`, the subclass should tolerate + missing tensors when its custom loading logic supports that mode. + """ raise NotImplementedError() @property - def skip_modules(self) -> List[str]: + def skip_modules(self) -> list[str]: return self._skip_modules - def add_skip_modules(self, value: List[str]) -> None: + def add_skip_modules(self, value: list[str]) -> None: self._skip_modules.extend(value) def should_skip_module(self, module_name: str) -> bool: return any(skip_module in module_name for skip_module in self._skip_modules) - def filter_weights(self, prefix: str, weights: dict) -> dict: + def filter_weights( + self, prefix: str, + weights: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """ + Return only weights that start with the prefix (and with the prefix removed) + """ result = {} for k, v in weights.items(): if k.startswith(prefix): @@ -165,7 +280,16 @@ def filter_weights(self, prefix: str, weights: dict) -> dict: return result @property - def mapping(self) -> dict: + def mapping(self) -> dict[str, list[str]]: + """Return the mapping for modules that need special weight loading. + + It maps module names to the corresponding source names in the checkpoint, e.g. + 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'], + 'gate_up_proj': ['gate_proj', 'up_proj'] + Those two modules fuse several GEMM weights together for a single GEMM call. + + Returns the mapping as dict[str, list[str]] + """ return self._mapping @property @@ -175,7 +299,7 @@ def config(self) -> ModelConfig: return self._config @property - def model(self) -> Union[nn.Module, DecoderModelForCausalLM]: + def model(self) -> nn.Module | DecoderModelForCausalLM: if self._model is None: raise RuntimeError("Weight mapper is not initialized") return self._model diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index a4e8f898548b..a49ea9cf22c3 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -225,7 +225,13 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping, use_consolidated: bool = False, - **kwargs) -> dict[str, Any]: + **kwargs) -> dict[str, torch.Tensor]: + """ + Load the model weights as a dict of name -> torch.Tensor. + + If CPU memory is large enough, prefetch all files in parallel to warm up OS file cache. + Then build and return a `ConsumableWeightsDict` by loading all files in parallel. + """ weight_files = glob.glob(f"{checkpoint_dir}/*.safetensors") # Some model checkpoint directories contain not only the sharded safetensors, but one # consolidated tensor. In the presence of both, we favor the former unless specified explicitly, as there really is no need @@ -310,7 +316,7 @@ def _load_weights_in_parallel(self, weight_files: List[str], load_func, return ConsumableWeightsDict(weights) @staticmethod - def _load_safetensors_file(file): + def _load_safetensors_file(file) -> dict[str, torch.Tensor]: logger.info(f"Start to load safetensor file {file}") return safetensors.torch.load_file(file) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py index 0828d10f6e66..2982b328be31 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping + import torch from torch import nn @@ -22,16 +24,21 @@ def map_weights(self) -> None: 'gate_up_proj': ['gate_proj', 'up_proj'] }) - def apply_callbacks(self, module: nn.Module, module_name: str, - module_names_breakdown: list[str], - weights: dict) -> list[dict]: + def apply_callbacks( + self, module: nn.Module, module_name: str, + module_names_breakdown: list[str], + weights: Mapping[str, + torch.Tensor]) -> list[dict[str, torch.Tensor]]: module_weights = [] - for new_name in self._mapping[module_name]: - fw = self.filter_weights( - '.'.join(module_names_breakdown + [new_name]), weights) + # For each raw checkpoint module name needed by `module`: + for source_name in self._mapping[module_name]: + # Find the tensors under this checkpoint module and remove + # their module path prefixes + fw: dict[str, torch.Tensor] = self.filter_weights( + '.'.join(module_names_breakdown + [source_name]), weights) for callback in self._callbacks: - fw = callback(module, new_name, fw) + fw = callback(module, source_name, fw) module_weights.append(fw) return module_weights @@ -65,8 +72,17 @@ def _num_kv_heads(self) -> int: return config.num_key_value_heads return config.num_attention_heads - def _duplicate_kv_weights(self, module: nn.Module, new_name: str, - weights: dict): + def _duplicate_kv_weights( + self, module: nn.Module, new_name: str, + weights: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """ + + If `new_name` is `k_proj` or `v_proj`, duplicate "weight" and "bias" weights + in `weights` if needed so that they can be sliced later to match sliced query heads, + allowing tensor parallelism to work on attention. + + Returned the potentially processed `weights`. + """ if new_name in ['k_proj', 'v_proj']: num_kv_heads = self._num_kv_heads @@ -87,26 +103,48 @@ def _duplicate_kv_weights(self, module: nn.Module, new_name: str, return weights def _duplicate_kv(self, weight: torch.Tensor, num_kv_heads: int, - tensor_parallel_size: int): + tensor_parallel_size: int) -> torch.Tensor: + """ + Duplicate K/V proj weight to match `tensor_parallel_size`. + + It expands kv proj weight [num_kv_heads*kv_head_dim, hidden_size] into + [num_kv_heads*rep*kv_head_dim, hidden_size], where rep = tp_size / num_kv_heads. + As an example, if rep == 2, the weight will be converted from logically: + [head_0_proj_matrix, head_1_proj_matrix, ...] to + [head_0_proj_matrix, head_0_proj_matrix, head_1_proj_matrix, head_1_proj_matrix, ...] + so that there are `tensor_parallel_size` proj matrices total, to be divided later + as one proj matrix per tp rank. + + Return the expanded weight tensor. + """ if num_kv_heads >= tensor_parallel_size: + # assume `num_kv_heads` is divisible by tp_size assert num_kv_heads % tensor_parallel_size == 0 return weight + # assume tp_size is divisible by `num_kv_heads` assert tensor_parallel_size % num_kv_heads == 0 reps = tensor_parallel_size // num_kv_heads # bias if weight.ndim == 1: + # From Yijing: this might be a bug, repeat_interleave does: + # repeat_interleave([1, 2, 3], 2) -> [1, 1, 2, 2, 3, 3] + # this is not how we usually slice tensors for tensor parallelism. + # This bug is not triggered yet since most if not all of LLM attention modules + # have no bias. return weight.repeat_interleave(reps) # weight and scale assert weight.shape[0] % num_kv_heads == 0 - size_per_kv_head = weight.shape[0] // num_kv_heads + size_per_kv_head = weight.shape[0] // num_kv_heads # aka, kv_head_dim + # Build a view of the tensor with shape [num_kv_heads, reps, kv_head_dim, hidden_size] weight = weight.reshape(num_kv_heads, size_per_kv_head, -1)[:, None, :, :].expand(num_kv_heads, reps, size_per_kv_head, weight.shape[1]) + # Return a new tensor of shape [num_kv_heads*reps*kv_head_dim, hidden_size] return weight.reshape(num_kv_heads * reps * size_per_kv_head, -1).clone().detach() diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 94d290d4f192..c65cc4ae0f6d 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -7,8 +7,8 @@ import os import time from dataclasses import dataclass -from typing import (Any, Dict, Generic, List, Literal, Optional, Tuple, Type, - TypeVar, Union) +from typing import (Any, Dict, Generic, Iterator, List, Literal, Optional, + Tuple, Type, TypeVar, Union) import torch from torch import nn @@ -35,13 +35,22 @@ @contextlib.contextmanager -def timing(message: str): +def timing(message: str) -> Iterator[None]: start = time.time() yield end = time.time() print(f"{message} -- {(end-start):.2f}s") +@contextlib.contextmanager +def timing_metric(metric_name: str, metric_dict: dict[str, + float]) -> Iterator[None]: + start = time.perf_counter() + yield + metric_dict[metric_name] = (metric_dict.get(metric_name, 0.0) + + time.perf_counter() - start) + + @dataclass class EagerFusionConfig: PRE_MOE_FUSION: bool = False @@ -754,11 +763,40 @@ def forward( ) def load_weights(self, - weights: Dict, + weights: dict[str, torch.Tensor], weight_mapper: Optional["BaseWeightMapper"] = None, - skip_modules: List[str] = [], - params_map: Optional[Dict[str, str]] = None, + skip_modules: list[str] = [], + params_map: dict[str, str] | None = None, allow_partial_loading: bool = False): + """Load checkpoint weights into this model. + + Basic function for an LLM class to load weights from a dict of weight + tensors. + The function walks the model's named modules, select matching tensors + from `weights`, and either call each module's `load_weights` method + if available, or copy tensors into parameters directly. + If `weight_mapper` is not None, uses it to perform custom weight mapping + before loading weights to each module; otherwise, perform hardcoded + weight fusion for some modules during loading. + + Args: + weights: dict[str, Tensor], mapping from checkpoint/state-dict keys + to tensors. If the key string does not match LLM class's child + module names, you need to use `params_map` or `weight_mapper` to + remap them. + weight_mapper: Optional mapper initialized for this model and + checkpoint format. When provided, it controls model-specific key + filtering, fused-module mappings, special module handling, and + manual parameter copies. + skip_modules: list[str], skip modules which contain these substrings. + This is used for LLM classes who have some child modules (e.g. + speculative decoding modules) that should be loaded from a + different function later. + params_map: Optional regex replacement map applied before loading + to rename checkpoint keys into the model's expected key space. + allow_partial_loading: if true, accept `weights` as an incomplete + weight dict and update only the parameters present. + """ # TODO smor- this solution is a temporary solution to load weights while we are still using # the old checkpoint format loading process. Once checkpoint format is unified # this method will be removed. @@ -1210,69 +1248,77 @@ def _load_weights_impl_v2(model: Union[nn.Module, DecoderModelForCausalLM], def load_single_module(name, module): torch.cuda.set_device(device_id) - if len(module._parameters) > 0: - if weight_mapper.should_skip_module(name): - return + if len(module._parameters) == 0 or weight_mapper.should_skip_module( + name): + return - names = name.split('.') + names = name.split('.') - # Special case: ConfigurableMoE.backend (TRTLLMGenFusedMoE) - # Currently saved MoE weights don't include 'backend' in their names. - # After MoE refactoring, ConfigurableMoE now has a backend submodule, - # and weights loading is done in the backend, so module name includes '.backend'. - # We need to use parent module name (without .backend) to match saved weight names. - # After MoE refactoring is fully complete, all paths will follow this branch. - if names[-1] == "backend" and isinstance(module, MoE): - name = '.'.join(names[:-1]) - names = name.split('.') + # Special case: ConfigurableMoE.backend (TRTLLMGenFusedMoE) + # Currently saved MoE weights don't include 'backend' in their names. + # After MoE refactoring, ConfigurableMoE now has a backend submodule, + # and weights loading is done in the backend, so module name includes '.backend'. + # We need to use parent module name (without .backend) to match saved weight names. + # After MoE refactoring is fully complete, all paths will follow this branch. + if names[-1] == "backend" and isinstance(module, MoE): + name = '.'.join(names[:-1]) + names = name.split('.') - module_names_breakdown, module_name = names[:-1], names[-1] + module_names_breakdown, module_name = names[:-1], names[-1] - if weight_mapper.does_require_special_handling(module_name): - module_weights = weight_mapper.apply_callbacks( + # check if the module has non-default weight loading, like fusing some weight + # tensors together. + if weight_mapper.does_require_special_handling(module_name): + # Process the weights, e.g. duplicating kv heads to match query heads after + # slicing for tensor parallelism. + module_weights: list[dict[ + str, torch.Tensor]] = weight_mapper.apply_callbacks( module, module_name, module_names_breakdown, weights) - module.load_weights(weights=module_weights, - allow_partial_loading=allow_partial_loading) - - # Mark consumed source weights (e.g., q_proj, k_proj, v_proj for qkv_proj) - if hasattr(weights, 'mark_consumed'): - for src_name in weight_mapper._mapping.get(module_name, []): - prefix = '.'.join(module_names_breakdown + [src_name]) - weights.mark_consumed(prefix) - else: - module_weights = weight_mapper.filter_weights(name, weights) - # Note: module_weights may be empty after filtering (e.g., in streaming weight updates) - if module_weights: - if weight_mapper.is_special_instance_module(module): - weight_mapper.handle_special_instance_module( - module, - module_name, - module_weights, - allow_partial_loading=allow_partial_loading) - elif hasattr(module, 'load_weights'): - if "linear_attn.conv1d" in name: - module_weights['weight'] = module_weights[ - 'weight'].squeeze(dim=1) - args = inspect.getfullargspec(module.load_weights).args - if "allow_partial_loading" not in args: - assert not allow_partial_loading, "allow_partial_loading is not supported for this model" - module.load_weights(weights=[module_weights]) - else: - module.load_weights( - weights=[module_weights], - allow_partial_loading=allow_partial_loading) - else: - for n, p in module.named_parameters(recurse=False): - weight_mapper.handle_manual_copy( - module_name, - module_weights, - n, - p, + # Call module's custom `load_weights()` to process weight, e.g. fusing + # several GEMM matrices together + module.load_weights(weights=module_weights, allow_partial_loading=allow_partial_loading) - # Mark consumed weights - if hasattr(weights, 'mark_consumed'): - weights.mark_consumed(name) + # Mark consumed source weights (e.g., q_proj, k_proj, v_proj for qkv_proj) + if hasattr(weights, 'mark_consumed'): + for src_name in weight_mapper.mapping.get(module_name, []): + prefix = '.'.join(module_names_breakdown + [src_name]) + weights.mark_consumed(prefix) + return + module_weights: dict[str, torch.Tensor] = weight_mapper.filter_weights( + name, weights) + # Note: module_weights may be empty after filtering (e.g., in streaming weight updates) + if not module_weights: + return + if weight_mapper.is_special_instance_module(module): + weight_mapper.handle_special_instance_module( + module, + module_name, + module_weights, + allow_partial_loading=allow_partial_loading) + elif hasattr(module, 'load_weights'): + if "linear_attn.conv1d" in name: + module_weights['weight'] = module_weights['weight'].squeeze( + dim=1) + args = inspect.getfullargspec(module.load_weights).args + if "allow_partial_loading" not in args: + assert not allow_partial_loading, "allow_partial_loading is not supported for this model" + module.load_weights(weights=[module_weights]) + else: + module.load_weights(weights=[module_weights], + allow_partial_loading=allow_partial_loading) + else: + for n, p in module.named_parameters(recurse=False): + weight_mapper.handle_manual_copy( + module_name, + module_weights, + n, + p, + allow_partial_loading=allow_partial_loading) + + # Mark consumed weights + if hasattr(weights, 'mark_consumed'): + weights.mark_consumed(name) if os.environ.get("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", "False") in ["True", "true", "1", "yes", "y"]: diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py b/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py index 115891c860c0..b622effe7c47 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py @@ -1,7 +1,7 @@ import gc import os import threading -from contextlib import nullcontext +from contextlib import AbstractContextManager, nullcontext from multiprocessing import resource_tracker, shared_memory from typing import Callable, Dict, List, Optional, Tuple @@ -1035,14 +1035,16 @@ def __exit__(self, exc_type, exc_val, exc_tb): def maybe_create_moe_load_balancer( - model_config, mapping: Optional[Mapping]) -> Optional[MoeLoadBalancer]: + model_config, mapping: Optional[Mapping] +) -> AbstractContextManager[MoeLoadBalancer | None]: ep_rank = model_config.mapping.moe_ep_rank ep_size = model_config.mapping.moe_ep_size model_arch = model_config.pretrained_config.architectures[0] using_ep = mapping and mapping.moe_ep_size > 1 in_supported_model_arch = model_arch in moe_model_arch_list using_smart_router = mapping and mapping.moe_cluster_size > 1 - moe_load_balancer = nullcontext() + moe_load_balancer: AbstractContextManager[MoeLoadBalancer + | None] = nullcontext() if in_supported_model_arch and using_ep and not using_smart_router and model_config.moe_load_balancer is not None: model_config.moe_load_balancer.setup(ep_rank=ep_rank, ep_size=ep_size) if model_config.moe_load_balancer.layer_updates_per_iter > 0: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 50aa317e9e43..10afde69aab8 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -407,6 +407,7 @@ def __init__( model_weights_memory_tag=model_weights_memory_tag, model_weights_restore_mode=model_weights_restore_mode, ) + # Open checkpoint and load the LLM module object. self.model, moe_load_balancer = self.model_loader.load( checkpoint_dir=model_path, checkpoint_loader=checkpoint_loader) if isinstance(moe_load_balancer, MoeLoadBalancer): diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index f2284e06bb46..83cb2e3ff3f0 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -6,6 +6,7 @@ import os import traceback import warnings +from enum import Enum from typing import Callable, Optional, Tuple import torch @@ -37,7 +38,7 @@ from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader from ..models.modeling_utils import (MODEL_CLASS_MAPPING, DecoderModelForCausalLM, MetaInitMode, - timing) + timing_metric) from ..modules.fused_moe.moe_load_balancer import ( MoeLoadBalancer, maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode @@ -304,6 +305,16 @@ def _apply_to_buffers_only(model: torch.nn.Module, fn): module._buffers[key] = fn(buf) +class ModelLoaderMetricNames(Enum): + TOTAL_MODEL_LOADING_SECONDS = "total_model_loading_seconds" + CHECKPOINT_PREPARATION_SECONDS = "checkpoint_preparation_seconds" + WEIGHT_POPULATION_SECONDS = "weight_population_seconds" + DRAFT_CHECKPOINT_PREPARATION_SECONDS = ( + "draft_checkpoint_preparation_seconds") + DRAFT_WEIGHT_POPULATION_SECONDS = "draft_weight_population_seconds" + POST_LOAD_PROCESSING_SECONDS = "post_load_processing_seconds" + + class ModelLoader: """ Handles the loading, configuration, and weight initialization of a PyTorch model. @@ -359,6 +370,13 @@ def __init__(self, self.weight_mapper = None self._weight_pool_proxy = None self._gms_backend = None + # Mostly weight loading and processing time metrics, updated when load() is called. + self._metrics: dict[str, float] = {} + + @property + def metrics(self) -> dict[str, float]: + """Return weight loading and processing time metrics.""" + return self._metrics @staticmethod def load_config_and_apply_defaults( @@ -477,7 +495,7 @@ def load( self, checkpoint_dir: str, checkpoint_loader: BaseCheckpointLoader, - ): + ) -> tuple[DecoderModelForCausalLM, MoeLoadBalancer | None]: """ Loads the model, its weights, and applies necessary configurations. @@ -488,12 +506,15 @@ def load( Returns: The loaded and initialized PyTorch model. """ + self._metrics = {} config = self._load_and_validate_config(checkpoint_dir, checkpoint_loader) load_format = self.llm_args.load_format - with timing("Model init total"), maybe_create_moe_load_balancer( - config, self.mapping) as moe_load_balancer: + with timing_metric( + ModelLoaderMetricNames.TOTAL_MODEL_LOADING_SECONDS.value, + self._metrics), maybe_create_moe_load_balancer( + config, self.mapping) as moe_load_balancer: try: # config will be modified in-place for some models, like Qwen2 config_copy = copy.deepcopy(config) @@ -639,12 +660,13 @@ def init_meta_tensor(t: torch.Tensor): load_weights_kwargs[ "prepare_post_transform_receiver"] = self._setup_aliases - if hasattr(model, 'llm_checkpoint_dir'): + _checkpoint_dir = getattr(model, "llm_checkpoint_dir", + checkpoint_dir) + with timing_metric( + ModelLoaderMetricNames.CHECKPOINT_PREPARATION_SECONDS. + value, self._metrics): weights = checkpoint_loader.load_weights( - model.llm_checkpoint_dir, **load_weights_kwargs) - else: - weights = checkpoint_loader.load_weights( - checkpoint_dir, **load_weights_kwargs) + _checkpoint_dir, **load_weights_kwargs) # When MX P2P succeeds, weights are already in model params. # A non-empty dict contains size-mismatched tensors that @@ -654,14 +676,22 @@ def init_meta_tensor(t: torch.Tensor): model, config) if weights: - self._call_load_weights(model.load_weights, weights, - self.weight_mapper) + with timing_metric( + ModelLoaderMetricNames.WEIGHT_POPULATION_SECONDS. + value, self._metrics): + self._call_load_weights(model.load_weights, weights, + self.weight_mapper) + torch.cuda.synchronize() if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights( ): - weights = checkpoint_loader.load_weights( - self.spec_config.speculative_model, - mapping=self.mapping) + with timing_metric( + ModelLoaderMetricNames. + DRAFT_CHECKPOINT_PREPARATION_SECONDS.value, + self._metrics): + weights = checkpoint_loader.load_weights( + self.spec_config.speculative_model, + mapping=self.mapping) draft_model_arch = model.draft_config.pretrained_config.architectures[ 0] @@ -670,8 +700,13 @@ def init_meta_tensor(t: torch.Tensor): draft_weight_mapper.init_model_and_config( model.draft_model, model.draft_config) - self._call_load_weights(model.load_draft_weights, weights, - draft_weight_mapper) + with timing_metric( + ModelLoaderMetricNames. + DRAFT_WEIGHT_POPULATION_SECONDS.value, + self._metrics): + self._call_load_weights(model.load_draft_weights, + weights, draft_weight_mapper) + torch.cuda.synchronize() elif load_format == LoadFormat.GMS: # GPU Memory Service path: weight tensors live in a @@ -765,8 +800,12 @@ def init_meta_tensor_in_pool(t: torch.Tensor): if qualification.qualified: load_weights_kwargs[ "prepare_post_transform_receiver"] = self._setup_aliases - weights = checkpoint_loader.load_weights( - weight_source, **load_weights_kwargs) + with timing_metric( + ModelLoaderMetricNames. + CHECKPOINT_PREPARATION_SECONDS.value, + self._metrics): + weights = checkpoint_loader.load_weights( + weight_source, **load_weights_kwargs) # `weights` may be: # - non-empty dict: standard mapping pipeline runs @@ -785,9 +824,14 @@ def init_meta_tensor_in_pool(t: torch.Tensor): if weights: self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( model, config) - self._call_load_weights(model.load_weights, - weights, - self.weight_mapper) + with timing_metric( + ModelLoaderMetricNames. + WEIGHT_POPULATION_SECONDS.value, + self._metrics): + self._call_load_weights( + model.load_weights, weights, + self.weight_mapper) + torch.cuda.synchronize() elif not weights_preloaded: raise RuntimeError( f"GMS RW: checkpoint loader " @@ -799,9 +843,13 @@ def init_meta_tensor_in_pool(t: torch.Tensor): if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights( ): - draft_weights = checkpoint_loader.load_weights( - self.spec_config.speculative_model, - mapping=self.mapping) + with timing_metric( + ModelLoaderMetricNames. + DRAFT_CHECKPOINT_PREPARATION_SECONDS. + value, self._metrics): + draft_weights = checkpoint_loader.load_weights( + self.spec_config.speculative_model, + mapping=self.mapping) draft_model_arch = model.draft_config.pretrained_config.architectures[ 0] @@ -811,55 +859,68 @@ def init_meta_tensor_in_pool(t: torch.Tensor): draft_weight_mapper.init_model_and_config( model.draft_model, model.draft_config) - self._call_load_weights( - model.load_draft_weights, draft_weights, - draft_weight_mapper) - - # Run post_load hooks INSIDE the pool so any - # tensors they create or rebind (fused QKV, - # quantization scales, derived aliases) land in - # GMS and become part of the committed layout - # that RO peers receive. Closes hhzhang16's - # narrow-scope and commit-ordering concerns. - checkpoint_loader.post_load_apply( - model, weights_preloaded=weights_preloaded) - - mx_staged_receiver_path = self._should_run_mx_staged_receiver_path( - checkpoint_loader, - model, - weights_preloaded=weights_preloaded, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) - if mx_staged_receiver_path: - self._setup_aliases(model) - self._mark_weights_transformed(model) - self._walk_cache_state(model) - else: - self._walk_full_post_load(model) - - # Defensive last-mile sweep: catches strays from - # C++ ops that bypassed the active torch - # allocator (e.g. native cudaMalloc). - gms_backend.move_untracked_params(model) - # Safe with active GMS mappings: GMSBackend.connect() - # installed upstream's patch_empty_cache(), which makes - # torch.cuda.empty_cache() skip GMS-backed VMM regions. - # Frees the non-GMS originals dropped by - # move_untracked_params (tensor.data was rebound to - # GMS-backed replacements above) before commit so the - # cached size doesn't show as live in memory accounting. - torch.cuda.empty_cache() - - self._post_load_publish( - checkpoint_loader, - model, - checkpoint_dir=checkpoint_dir, - weights_preloaded=weights_preloaded, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) + with timing_metric( + ModelLoaderMetricNames. + DRAFT_WEIGHT_POPULATION_SECONDS.value, + self._metrics): + self._call_load_weights( + model.load_draft_weights, draft_weights, + draft_weight_mapper) + torch.cuda.synchronize() + + with timing_metric( + ModelLoaderMetricNames. + POST_LOAD_PROCESSING_SECONDS.value, + self._metrics): + # Run post_load hooks INSIDE the pool so any + # tensors they create or rebind (fused QKV, + # quantization scales, derived aliases) land in + # GMS and become part of the committed layout + # that RO peers receive. Closes hhzhang16's + # narrow-scope and commit-ordering concerns. + checkpoint_loader.post_load_apply( + model, weights_preloaded=weights_preloaded) + + mx_staged_receiver_path = self._should_run_mx_staged_receiver_path( + checkpoint_loader, + model, + weights_preloaded=weights_preloaded, + speculative_mode=speculative_mode, + loads_draft_weights=loads_draft_weights) + if mx_staged_receiver_path: + self._setup_aliases(model) + self._mark_weights_transformed(model) + self._walk_cache_state(model) + else: + self._walk_full_post_load(model) + + # Defensive last-mile sweep: catches strays from + # C++ ops that bypassed the active torch + # allocator (e.g. native cudaMalloc). + gms_backend.move_untracked_params(model) + # Safe with active GMS mappings: GMSBackend.connect() + # installed upstream's patch_empty_cache(), which makes + # torch.cuda.empty_cache() skip GMS-backed VMM regions. + # Frees the non-GMS originals dropped by + # move_untracked_params (tensor.data was rebound to + # GMS-backed replacements above) before commit so the + # cached size doesn't show as live in memory accounting. + torch.cuda.empty_cache() + + self._post_load_publish( + checkpoint_loader, + model, + checkpoint_dir=checkpoint_dir, + weights_preloaded=weights_preloaded, + speculative_mode=speculative_mode, + loads_draft_weights=loads_draft_weights) # Pool closed. Commit the post-post_load layout. - gms_backend.finalize_write(model) + with timing_metric( + ModelLoaderMetricNames. + POST_LOAD_PROCESSING_SECONDS.value, + self._metrics): + gms_backend.finalize_write(model) gms_post_load_handled = True logger.info( "LoadFormat.GMS (RW): loaded and committed weights via %s", @@ -891,26 +952,30 @@ def init_meta_tensor_in_pool(t: torch.Tensor): # tensors without re-running one-shot transforms. # 6. `post_load_publish`: any receiver-side # publish (no-op via the receiver guard). - checkpoint_loader.post_load_apply( - model, weights_preloaded=True) + with timing_metric( + ModelLoaderMetricNames. + POST_LOAD_PROCESSING_SECONDS.value, + self._metrics): + checkpoint_loader.post_load_apply( + model, weights_preloaded=True) - self._setup_aliases(model) + self._setup_aliases(model) - # Pre-materialize compatibility gate. GMS has no - # disk-fallback path, so a mismatch raises under STRICT - # rather than falling back. - self._check_gms_source_identity(gms_backend) + # Pre-materialize compatibility gate. GMS has no + # disk-fallback path, so a mismatch raises under STRICT + # rather than falling back. + self._check_gms_source_identity(gms_backend) - gms_backend.materialize_module(model) - self._walk_cache_state(model) + gms_backend.materialize_module(model) + self._walk_cache_state(model) - self._post_load_publish( - checkpoint_loader, - model, - checkpoint_dir=checkpoint_dir, - weights_preloaded=True, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) + self._post_load_publish( + checkpoint_loader, + model, + checkpoint_dir=checkpoint_dir, + weights_preloaded=True, + speculative_mode=speculative_mode, + loads_draft_weights=loads_draft_weights) gms_post_load_handled = True logger.info("LoadFormat.GMS (RO): materialized weights") else: @@ -942,56 +1007,60 @@ def init_meta_tensor_in_pool(t: torch.Tensor): raise NotImplementedError( f"No load support for load format: {load_format}") - if not gms_post_load_handled: - checkpoint_loader.post_load_apply( - model, weights_preloaded=weights_preloaded) - mx_staged_receiver_path = self._should_run_mx_staged_receiver_path( - checkpoint_loader, - model, - weights_preloaded=weights_preloaded, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) - if mx_staged_receiver_path: - self._setup_aliases(model) - self._mark_weights_transformed(model) - self._walk_cache_state(model) - else: - self._walk_full_post_load(model) - self._post_load_publish(checkpoint_loader, - model, - checkpoint_dir=checkpoint_dir, - weights_preloaded=weights_preloaded, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) - - # TODO(GMS-MOE-LB): when the (MoE, GMS) combination is enabled, - # `register_weight_slots_after_to_cuda` and `finalize_model` - # must run INSIDE `mem_pool_scope` and BEFORE `finalize_write` - # so MoE allocations become part of the committed layout that - # RO peers receive. Today they run outside the pool and after - # commit, which would silently produce a broken MoE routing - # state on RO peers — that combination is REJECTED at config - # time by `TorchLlmArgs.validate_gms_moe_compat` in - # `tensorrt_llm/llmapi/llm_args.py`. When implementing the - # follow-up, drop the validator gate AFTER moving these calls - # into the pool scope. - if isinstance(moe_load_balancer, MoeLoadBalancer): - moe_load_balancer.register_weight_slots_after_to_cuda() - logger.info("moe_load_balancer finalizing model...") - moe_load_balancer.finalize_model() - logger.info("moe_load_balancer finalize model done") - - torch.cuda.current_stream().synchronize() - # Reclaim segments freed during per-module post_load_weights (e.g. - # MegaMoE _transform_main_weights releases ~5-6 GiB of redundant - # weight Parameters via `.data = empty(0)` that PyTorch's caching - # allocator otherwise holds onto). Returning them to the driver - # gives downstream stages (KV cache estimation, attention workspace - # alloc, autotuner warmup symmetric-fabric setup) full visibility - # of free HBM. Single one-shot call after all modules are - # finalized; per-layer empty_cache here is unsafe because it can - # perturb NVLink barrier synchronization in multi-rank DG init. - torch.cuda.empty_cache() + with timing_metric( + ModelLoaderMetricNames.POST_LOAD_PROCESSING_SECONDS.value, + self._metrics): + if not gms_post_load_handled: + checkpoint_loader.post_load_apply( + model, weights_preloaded=weights_preloaded) + mx_staged_receiver_path = self._should_run_mx_staged_receiver_path( + checkpoint_loader, + model, + weights_preloaded=weights_preloaded, + speculative_mode=speculative_mode, + loads_draft_weights=loads_draft_weights) + if mx_staged_receiver_path: + self._setup_aliases(model) + self._mark_weights_transformed(model) + self._walk_cache_state(model) + else: + self._walk_full_post_load(model) + self._post_load_publish( + checkpoint_loader, + model, + checkpoint_dir=checkpoint_dir, + weights_preloaded=weights_preloaded, + speculative_mode=speculative_mode, + loads_draft_weights=loads_draft_weights) + + # TODO(GMS-MOE-LB): when the (MoE, GMS) combination is enabled, + # `register_weight_slots_after_to_cuda` and `finalize_model` + # must run INSIDE `mem_pool_scope` and BEFORE `finalize_write` + # so MoE allocations become part of the committed layout that + # RO peers receive. Today they run outside the pool and after + # commit, which would silently produce a broken MoE routing + # state on RO peers — that combination is REJECTED at config + # time by `TorchLlmArgs.validate_gms_moe_compat` in + # `tensorrt_llm/llmapi/llm_args.py`. When implementing the + # follow-up, drop the validator gate AFTER moving these calls + # into the pool scope. + if isinstance(moe_load_balancer, MoeLoadBalancer): + moe_load_balancer.register_weight_slots_after_to_cuda() + logger.info("moe_load_balancer finalizing model...") + moe_load_balancer.finalize_model() + logger.info("moe_load_balancer finalize model done") + + torch.cuda.synchronize() + # Reclaim segments freed during per-module post_load_weights (e.g. + # MegaMoE _transform_main_weights releases ~5-6 GiB of redundant + # weight Parameters via `.data = empty(0)` that PyTorch's caching + # allocator otherwise holds onto). Returning them to the driver + # gives downstream stages (KV cache estimation, attention workspace + # alloc, autotuner warmup symmetric-fabric setup) full visibility + # of free HBM. Single one-shot call after all modules are + # finalized; per-layer empty_cache here is unsafe because it can + # perturb NVLink barrier synchronization in multi-rank DG init. + torch.cuda.empty_cache() return model, moe_load_balancer diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index 3f59228b2b00..d1daff05e76f 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -335,6 +335,7 @@ def latency_command( logger.info("Setting up latency benchmark.") llm = get_llm(runtime_config, kwargs) + startup_metrics = None if options.report_json is None else llm.startup_metrics ignore_eos = True if runtime_config.decoding_config.decoding_mode == SpeculativeDecodingMode.NONE else False eos_id = tokenizer.eos_token_id if not ignore_eos else -1 @@ -390,8 +391,13 @@ def latency_command( # For multimodal models, we need to update the metadata with the correct input lengths metadata = update_metadata_for_multimodal(metadata, statistics) - report_utility = ReportUtility(statistics, metadata, runtime_config, - logger, kwargs, True) + report_utility = ReportUtility(statistics, + metadata, + runtime_config, + logger, + kwargs, + True, + startup_metrics=startup_metrics) # Generate reports for statistics, output tokens, and request info. generate_json_report(options.report_json, report_utility.get_statistics_dict) diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index b31a4608a5c3..4205e848f07a 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -468,6 +468,7 @@ def throughput_command( "max_num_tokens", runtime_config.settings_config.max_num_tokens) llm = get_llm(runtime_config, kwargs) + startup_metrics = None if options.report_json is None else llm.startup_metrics sampler_args = { "end_id": options.eos_id, @@ -525,8 +526,13 @@ def throughput_command( # For multimodal models, we need to update the metadata with the correct input lengths metadata = update_metadata_for_multimodal(metadata, statistics) - report_utility = ReportUtility(statistics, metadata, runtime_config, - logger, kwargs, options.streaming) + report_utility = ReportUtility(statistics, + metadata, + runtime_config, + logger, + kwargs, + options.streaming, + startup_metrics=startup_metrics) # Generate reports for statistics, output tokens, and request info. generate_json_report(options.report_json, report_utility.get_statistics_dict) diff --git a/tensorrt_llm/bench/dataclasses/reporting.py b/tensorrt_llm/bench/dataclasses/reporting.py index 3c92667f2444..60aa3508aa6b 100755 --- a/tensorrt_llm/bench/dataclasses/reporting.py +++ b/tensorrt_llm/bench/dataclasses/reporting.py @@ -263,7 +263,8 @@ def __init__(self, rt_cfg: RuntimeConfig, logger: Logger, kwargs: Dict[str, Any], - streaming: bool = False) -> None: + streaming: bool = False, + startup_metrics: Optional[Dict[str, Any]] = None) -> None: """Initialize the ReportingController. Args: @@ -272,6 +273,8 @@ def __init__(self, rt_cfg (RuntimeConfig): Configuration for the run. logger (Logger): A logger for logging. streaming (bool, optional): Streaming benchmark used. Defaults to False. + startup_metrics (Dict[str, Any], optional): Metrics captured while + initializing the model. """ self.dataset_metadata = dataset_metadata self.rt_cfg = rt_cfg @@ -282,6 +285,7 @@ def __init__(self, self.get_max_draft_len(), self.rt_cfg.settings_config.max_batch_size) self.streaming = streaming + self.startup_metrics = startup_metrics def _query_gpu_info(self) -> Dict[str, Any]: """Query first GPU info (all GPUs must be identical for TRT-LLM).""" @@ -402,6 +406,8 @@ def get_statistics_dict(self) -> Dict[str, Any]: "version": self.rt_cfg.sw_version, }, } + if self.startup_metrics: + stats_dict["startup_metrics"] = self.startup_metrics # Machine / GPU details - query only first GPU (all GPUs must be identical) stats_dict["machine"] = self._query_gpu_info() diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 5e1413a4f0da..abf5843600f6 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -983,6 +983,25 @@ def get_data_transceiver_state(self) -> bytes: return b"" return self.engine.kv_cache_transceiver.get_data_transceiver_state() + def get_startup_metrics(self) -> dict: + """Return rank-local startup metrics for the PyTorch backend.""" + if not self._is_pytorch_backend or self.engine is None: + return {} + + startup_metrics = {} + model_engine = getattr(self.engine, "model_engine", None) + model_loader = getattr(model_engine, "model_loader", None) + if model_loader is not None: + startup_metrics["model_loader"] = dict(model_loader.metrics) + + draft_model_engine = getattr(self.engine, "draft_model_engine", None) + draft_model_loader = getattr(draft_model_engine, "model_loader", None) + if draft_model_loader is not None: + startup_metrics["draft_model_loader"] = dict( + draft_model_loader.metrics) + + return startup_metrics + @staticmethod def _stats_serializer(stats) -> str: # Per-rank path: stats is ("per_rank_dict", {..., "rank": N}). diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index 98afe8daa44f..affd24137b72 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -471,6 +471,10 @@ def get_cache_transceiver(self) -> Optional["KvCacheTransceiver"]: def get_data_transceiver_state(self) -> bytes: return b"" + def get_startup_metrics(self) -> dict: + """Return metrics captured while initializing the executor.""" + return {} + @staticmethod def _create_ray_executor( worker_kwargs: Dict, diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index d1baa0792785..eb6e7cfbe9dd 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -958,6 +958,19 @@ def get_data_transceiver_state(self) -> bytes: logger.error(f"Error fetching data transceiver state via RPC: {e}") raise + def get_startup_metrics(self) -> dict: + """Get rank-0 startup metrics from the worker runtime via RPC.""" + if self.rpc_client is None: + logger.warning( + "RPC client not initialized, cannot get startup metrics") + return {} + try: + metrics = self.rpc_client.get_startup_metrics().remote() + return metrics if isinstance(metrics, dict) else {} + except RPCError as e: + logger.warning(f"Error fetching startup metrics via RPC: {e}") + return {} + def aget_stats(self, timeout: float) -> IterationResult: """Get iteration statistics from the runtime via RPC (async). diff --git a/tensorrt_llm/executor/ray_executor.py b/tensorrt_llm/executor/ray_executor.py index accd9efc1632..f480937dddd3 100644 --- a/tensorrt_llm/executor/ray_executor.py +++ b/tensorrt_llm/executor/ray_executor.py @@ -256,6 +256,11 @@ async def collective_rpc_async( target_ranks=target_ranks) return await asyncio.gather(*refs) + def get_startup_metrics(self) -> dict: + """Get startup metrics from rank 0.""" + metrics = self.collective_rpc("get_startup_metrics", target_ranks=0) + return metrics[0] if metrics and isinstance(metrics[0], dict) else {} + def submit(self, request: "GenerationRequest") -> "GenerationResult": """ Low-level API to the executor. Return a "future" GenerationResult diff --git a/tensorrt_llm/executor/rpc_proxy.py b/tensorrt_llm/executor/rpc_proxy.py index 0763e09ad032..43a809625b04 100644 --- a/tensorrt_llm/executor/rpc_proxy.py +++ b/tensorrt_llm/executor/rpc_proxy.py @@ -129,6 +129,15 @@ def get_kv_cache_capacity(self) -> dict: logger.debug(f"Error fetching kv cache capacity via RPC: {e}") return {} + def get_startup_metrics(self) -> dict: + """Get rank-0 startup metrics from the worker runtime via RPC.""" + try: + metrics = self.rpc_client.get_startup_metrics().remote() + return metrics if isinstance(metrics, dict) else {} + except RPCError as e: + logger.warning(f"Error fetching startup metrics via RPC: {e}") + return {} + def aget_stats(self, timeout: float) -> IterationResult: """Get iteration statistics from the runtime via RPC (async). diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 577717586748..867b7cc27ce5 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -163,7 +163,9 @@ def __init__(self, self._executor_cls = kwargs.pop("executor_cls", GenerationExecutor) self._orchestrator_type = kwargs.get("orchestrator_type", None) self._llm_id = None - self._disaggregated_params: Optional[dict] = None + self._disaggregated_params: dict | None = None + # dict containing startup metrics like weight loading time + self._startup_metrics: dict | None = None log_level = logger.level logger.set_level("info") # force display the backend @@ -347,6 +349,15 @@ def disaggregated_params(self) -> dict: ) if self._executor else {} return self._disaggregated_params + @property + @set_api_status("beta") + def startup_metrics(self) -> dict: + """Return the cached startup metrics reported by worker rank 0.""" + if self._startup_metrics is None: + self._startup_metrics = self._executor.get_startup_metrics( + ) if self._executor else {} + return self._startup_metrics + @staticmethod def _is_token_id_list(value: Any) -> bool: return isinstance(value, list) and all( diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index ad1070f2d8ea..159c4eb6a5fb 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -2363,6 +2363,9 @@ async def update_weights(self, return JSONResponse(content={"status": "success"}) async def get_server_info(self) -> JSONResponse: + # Note: calling self.generator.disaggregated_params and startup_metrics below + # may trigger an RPC sync call, blocking the server event loop. Since this server_info + # is usually called only once before accepting requests, it's not a big concern. content = {"disaggregated_params": self.generator.disaggregated_params} args = getattr(self.generator, "args", None) if args is not None: @@ -2377,6 +2380,8 @@ async def get_server_info(self) -> JSONResponse: if kv_cache_config.tokens_per_block is not None: content[ "tokens_per_block"] = kv_cache_config.tokens_per_block + content["startup_metrics"] = getattr(self.generator, "startup_metrics", + {}) return JSONResponse(content=content) async def openai_image_generation(self, request: ImageGenerationRequest, diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index 85d9eaf95797..c1c494833f28 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -18,6 +18,7 @@ from tensorrt_llm.executor import GenerationResultBase, RequestError from tensorrt_llm.llmapi import (KvCacheConfig, KvCacheRetentionConfig, LookaheadDecodingConfig, RequestOutput) +from tensorrt_llm.llmapi.llm import BaseLLM from tensorrt_llm.llmapi.llm_args import DynamicBatchConfig, SchedulerConfig from tensorrt_llm.llmapi.llm_utils import _ParallelConfig from tensorrt_llm.llmapi.tokenizer import (TokenizerBase, TransformersTokenizer, @@ -1293,6 +1294,41 @@ async def is_disconnected(self): return True +def test_startup_metrics_are_cached_and_reported_by_server_info() -> None: + + class FakeExecutor: + + def __init__(self): + self.calls = 0 + + def get_startup_metrics(self): + self.calls += 1 + return {"model_loader": {"total_model_loading_seconds": 1.5}} + + executor = FakeExecutor() + generator = object.__new__(BaseLLM) + generator._executor = executor + generator._startup_metrics = None + generator._disaggregated_params = {} + + server = object.__new__(OpenAIServer) + server.generator = generator + first_response = asyncio.run(server.get_server_info()) + second_response = asyncio.run(server.get_server_info()) + + expected = { + "disaggregated_params": {}, + "startup_metrics": { + "model_loader": { + "total_model_loading_seconds": 1.5 + } + }, + } + assert json.loads(first_response.body) == expected + assert json.loads(second_response.body) == expected + assert executor.calls == 1 + + def test_openai_completion_list_prompt_stream_reuses_stream_metadata() -> None: async def run_request():