| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import time |
| 4 | from abc import ABC, abstractmethod |
| 5 | |
| 6 | import torch |
| 7 | import torch.nn as nn |
| 8 | |
| 9 | import vllm.envs as envs |
| 10 | from vllm.config import ModelConfig, VllmConfig |
| 11 | from vllm.config.load import LoadConfig |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.model_executor.model_loader.reload import finalize_layerwise_processing |
| 14 | from vllm.model_executor.model_loader.utils import ( |
| 15 | initialize_model, |
| 16 | process_weights_after_loading, |
| 17 | ) |
| 18 | from vllm.platforms import current_platform |
| 19 | from vllm.tracing import instrument |
| 20 | from vllm.utils.mem_utils import format_gib |
| 21 | from vllm.utils.torch_utils import set_default_torch_dtype |
| 22 | |
| 23 | logger = init_logger(__name__) |
| 24 | |
| 25 | |
| 26 | class BaseModelLoader(ABC): |
| 27 | """Base class for model loaders.""" |
| 28 | |
| 29 | def __init__(self, load_config: LoadConfig): |
| 30 | self.load_config = load_config |
| 31 | |
| 32 | def get_external_weight_memory(self, vllm_config: VllmConfig) -> int: |
| 33 | """Get weights memory from external process; |
| 34 | 0 when the weights are not external.""" |
| 35 | return 0 |
| 36 | |
| 37 | @abstractmethod |
| 38 | def download_model(self, model_config: ModelConfig) -> None: |
| 39 | """Download a model so that it can be immediately loaded.""" |
| 40 | raise NotImplementedError |
| 41 | |
| 42 | @abstractmethod |
| 43 | def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: |
| 44 | """Load weights into a model. This standalone API allows |
| 45 | inplace weights loading for an already-initialized model""" |
| 46 | raise NotImplementedError |
| 47 | |
| 48 | def create_model( |
| 49 | self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = "" |
| 50 | ) -> nn.Module: |
| 51 | """Create a model with the given configurations.""" |
| 52 | model = initialize_model( |
| 53 | vllm_config=vllm_config, |
| 54 | model_config=model_config, |
| 55 | prefix=prefix, |
| 56 | ) |
| 57 | log_online_quantization(vllm_config) |
| 58 | log_model_inspection(model) |
| 59 | return model |
| 60 | |
| 61 | @instrument(span_name="Load model") |
| 62 | def load_model( |
| 63 | self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = "" |
| 64 | ) -> nn.Module: |
| 65 | """Load a model with the given configurations.""" |
| 66 | device_config = vllm_config.device_config |
| 67 | load_config = vllm_config.load_config |
| 68 | load_device = ( |
| 69 | device_config.device if load_config.device is None else load_config.device |
| 70 | ) |
| 71 | target_device = torch.device(load_device) |
| 72 | with set_default_torch_dtype(model_config.dtype): |
| 73 | with target_device: |
| 74 | time_before_load = time.perf_counter() |
| 75 | model = self.create_model( |
| 76 | vllm_config=vllm_config, |
| 77 | model_config=model_config, |
| 78 | prefix=prefix, |
| 79 | ) |
| 80 | time_after_load = time.perf_counter() |
| 81 | logger.info_once( |
| 82 | "Initializing model took %.6f seconds", |
| 83 | time_after_load - time_before_load, |
| 84 | ) |
| 85 | |
| 86 | logger.debug("Loading weights on %s ...", load_device) |
| 87 | self.load_weights(model, model_config) |
| 88 | |
| 89 | # Log peak GPU memory after loading weights. This is needed |
| 90 | # to have test coverage on peak memory for online quantization. |
| 91 | if current_platform.is_cuda_alike() or current_platform.is_xpu(): |
| 92 | peak_memory = torch.accelerator.max_memory_allocated() |
| 93 | logger.debug_once( |
| 94 | "Peak GPU memory after loading weights: %s GiB", |
| 95 | format_gib(peak_memory), |
| 96 | ) |
| 97 | |
| 98 | # Process weights into kernel format. Note that when using online |
| 99 | # quantization, weights are (typically) quantized as they are loaded. |
| 100 | if _has_online_quant(model): |
| 101 | finalize_layerwise_processing(model, model_config) |
| 102 | |
| 103 | process_weights_after_loading(model, model_config, target_device) |
| 104 | |
| 105 | return model.eval() |
| 106 | |
| 107 | |
| 108 | def log_model_inspection(model: nn.Module) -> None: |
| 109 | """Log model structure if VLLM_LOG_MODEL_INSPECTION=1.""" |
| 110 | if not envs.VLLM_LOG_MODEL_INSPECTION: |
| 111 | return |
| 112 | |
| 113 | from vllm.model_inspection import format_model_inspection |
| 114 | |
| 115 | logger.info("vLLM model structure:\n%s", format_model_inspection(model)) |
| 116 | |
| 117 | |
| 118 | def log_online_quantization(vllm_config: VllmConfig) -> None: |
| 119 | """Log the online-quantized layer count and types, when applicable.""" |
| 120 | from vllm.model_executor.layers.quantization.online.base import ( |
| 121 | OnlineQuantizationConfig, |
| 122 | ) |
| 123 | |
| 124 | quant_config = vllm_config.quant_config |
| 125 | online_quantization_config = getattr( |
| 126 | quant_config, "online_quantization_config", None |
| 127 | ) |
| 128 | if isinstance(online_quantization_config, OnlineQuantizationConfig): |
| 129 | quant_config = online_quantization_config |
| 130 | if not isinstance(quant_config, OnlineQuantizationConfig): |
| 131 | return |
| 132 | |
| 133 | logger.info( |
| 134 | "Quantized %d layers of types: %s", |
| 135 | len(quant_config.quantized_layers), |
| 136 | "; ".join(quant_config.quantized_layer_summaries), |
| 137 | ) |
| 138 | |
| 139 | |
| 140 | def _has_online_quant(model: nn.Module): |
| 141 | for module in model.modules(): |
| 142 | quant_method = getattr(module, "quant_method", None) |
| 143 | if getattr(quant_method, "uses_meta_device", False): |
| 144 | return True |
| 145 | |
| 146 | return False |