Skip to content

vllm.model_executor.model_loader.base_loader

Classes:

Functions:

BaseModelLoader

Bases: ABC

Base class for model loaders.

Methods:

Source code in vllm/model_executor/model_loader/base_loader.py 
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
class BaseModelLoader(ABC):
    """Base class for model loaders."""

    def __init__(self, load_config: LoadConfig):
        self.load_config = load_config

    def get_external_weight_memory(self, vllm_config: VllmConfig) -> int:
        """Get weights memory from external process;
        0 when the weights are not external."""
        return 0

    @abstractmethod
    def download_model(self, model_config: ModelConfig) -> None:
        """Download a model so that it can be immediately loaded."""
        raise NotImplementedError

    @abstractmethod
    def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
        """Load weights into a model. This standalone API allows
        inplace weights loading for an already-initialized model"""
        raise NotImplementedError

    def create_model(
        self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
    ) -> nn.Module:
        """Create a model with the given configurations."""
        model = initialize_model(
            vllm_config=vllm_config,
            model_config=model_config,
            prefix=prefix,
        )
        log_online_quantization(vllm_config)
        log_model_inspection(model)
        return model

    def get_all_weights(
        self,
        model_config: ModelConfig,
        model: nn.Module,
    ) -> Generator[tuple[str, torch.Tensor], None, None]:
        """Iterate over all weights by checkpoint name, for filtered
        partial reloads (e.g. FT expert reload after scale_down)."""
        raise NotImplementedError(
            f"load_format={self.load_config.load_format!r} does not support "
            "iterating weights by checkpoint name"
        )

    @instrument(span_name="Load model")
    def load_model(
        self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
    ) -> nn.Module:
        """Load a model with the given configurations."""
        device_config = vllm_config.device_config
        load_config = vllm_config.load_config
        load_device = (
            device_config.device if load_config.device is None else load_config.device
        )
        target_device = torch.device(load_device)
        with set_default_torch_dtype(model_config.dtype):
            with target_device:
                time_before_load = time.perf_counter()
                model = self.create_model(
                    vllm_config=vllm_config,
                    model_config=model_config,
                    prefix=prefix,
                )
                time_after_load = time.perf_counter()
                logger.info_once(
                    "Initializing model took %.6f seconds",
                    time_after_load - time_before_load,
                )

            logger.debug("Loading weights on %s ...", load_device)
            self.load_weights(model, model_config)

            # Release transient tensors held by PyTorch's caching allocator
            # before weight post-processing and raw CUDA driver allocations.
            gc.collect()
            torch.accelerator.empty_cache()

            # Log peak GPU memory after loading weights. This is needed
            # to have test coverage on peak memory for online quantization.
            if current_platform.is_cuda_alike() or current_platform.is_xpu():
                peak_memory = torch.accelerator.max_memory_allocated()
                logger.debug_once(
                    "Peak GPU memory after loading weights: %s GiB",
                    format_gib(peak_memory),
                )

            # Process weights into kernel format. Note that when using online
            # quantization, weights are (typically) quantized as they are loaded.
            if _has_online_quant(model):
                finalize_layerwise_processing(model, model_config)

            process_weights_after_loading(model, model_config, target_device)
            log_online_quantization_time(vllm_config)

            gc.collect()
            torch.accelerator.empty_cache()

        return model.eval()

create_model(vllm_config, model_config, prefix='')

Create a model with the given configurations.

Source code in vllm/model_executor/model_loader/base_loader.py 
50
51
52
53
54
55
56
57
58
59
60
61
def create_model(
    self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
) -> nn.Module:
    """Create a model with the given configurations."""
    model = initialize_model(
        vllm_config=vllm_config,
        model_config=model_config,
        prefix=prefix,
    )
    log_online_quantization(vllm_config)
    log_model_inspection(model)
    return model

download_model(model_config) abstractmethod

Download a model so that it can be immediately loaded.

Source code in vllm/model_executor/model_loader/base_loader.py 
39
40
41
42
@abstractmethod
def download_model(self, model_config: ModelConfig) -> None:
    """Download a model so that it can be immediately loaded."""
    raise NotImplementedError

get_all_weights(model_config, model)

Iterate over all weights by checkpoint name, for filtered partial reloads (e.g. FT expert reload after scale_down).

Source code in vllm/model_executor/model_loader/base_loader.py 
63
64
65
66
67
68
69
70
71
72
73
def get_all_weights(
    self,
    model_config: ModelConfig,
    model: nn.Module,
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Iterate over all weights by checkpoint name, for filtered
    partial reloads (e.g. FT expert reload after scale_down)."""
    raise NotImplementedError(
        f"load_format={self.load_config.load_format!r} does not support "
        "iterating weights by checkpoint name"
    )

get_external_weight_memory(vllm_config)

Get weights memory from external process; 0 when the weights are not external.

Source code in vllm/model_executor/model_loader/base_loader.py 
34
35
36
37
def get_external_weight_memory(self, vllm_config: VllmConfig) -> int:
    """Get weights memory from external process;
    0 when the weights are not external."""
    return 0

load_model(vllm_config, model_config, prefix='')

Load a model with the given configurations.

Source code in vllm/model_executor/model_loader/base_loader.py 
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
@instrument(span_name="Load model")
def load_model(
    self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
) -> nn.Module:
    """Load a model with the given configurations."""
    device_config = vllm_config.device_config
    load_config = vllm_config.load_config
    load_device = (
        device_config.device if load_config.device is None else load_config.device
    )
    target_device = torch.device(load_device)
    with set_default_torch_dtype(model_config.dtype):
        with target_device:
            time_before_load = time.perf_counter()
            model = self.create_model(
                vllm_config=vllm_config,
                model_config=model_config,
                prefix=prefix,
            )
            time_after_load = time.perf_counter()
            logger.info_once(
                "Initializing model took %.6f seconds",
                time_after_load - time_before_load,
            )

        logger.debug("Loading weights on %s ...", load_device)
        self.load_weights(model, model_config)

        # Release transient tensors held by PyTorch's caching allocator
        # before weight post-processing and raw CUDA driver allocations.
        gc.collect()
        torch.accelerator.empty_cache()

        # Log peak GPU memory after loading weights. This is needed
        # to have test coverage on peak memory for online quantization.
        if current_platform.is_cuda_alike() or current_platform.is_xpu():
            peak_memory = torch.accelerator.max_memory_allocated()
            logger.debug_once(
                "Peak GPU memory after loading weights: %s GiB",
                format_gib(peak_memory),
            )

        # Process weights into kernel format. Note that when using online
        # quantization, weights are (typically) quantized as they are loaded.
        if _has_online_quant(model):
            finalize_layerwise_processing(model, model_config)

        process_weights_after_loading(model, model_config, target_device)
        log_online_quantization_time(vllm_config)

        gc.collect()
        torch.accelerator.empty_cache()

    return model.eval()

load_weights(model, model_config) abstractmethod

Load weights into a model. This standalone API allows inplace weights loading for an already-initialized model

Source code in vllm/model_executor/model_loader/base_loader.py 
44
45
46
47
48
@abstractmethod
def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
    """Load weights into a model. This standalone API allows
    inplace weights loading for an already-initialized model"""
    raise NotImplementedError

_get_online_quantization_config(vllm_config)

Return the online quantization config when one is configured.

Source code in vllm/model_executor/model_loader/base_loader.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
def _get_online_quantization_config(vllm_config: VllmConfig):
    """Return the online quantization config when one is configured."""
    from vllm.model_executor.layers.quantization.online.base import (
        OnlineQuantizationConfig,
    )

    quant_config = vllm_config.quant_config
    online_quantization_config = getattr(
        quant_config, "online_quantization_config", None
    )
    if isinstance(online_quantization_config, OnlineQuantizationConfig):
        quant_config = online_quantization_config
    if not isinstance(quant_config, OnlineQuantizationConfig):
        return None
    return quant_config

log_model_inspection(model)

Log model structure if VLLM_LOG_MODEL_INSPECTION=1.

Source code in vllm/model_executor/model_loader/base_loader.py 
131
132
133
134
135
136
137
138
def log_model_inspection(model: nn.Module) -> None:
    """Log model structure if VLLM_LOG_MODEL_INSPECTION=1."""
    if not envs.VLLM_LOG_MODEL_INSPECTION:
        return

    from vllm.model_inspection import format_model_inspection

    logger.info("vLLM model structure:\n%s", format_model_inspection(model))

log_online_quantization(vllm_config)

Log the online-quantized layer count and types, when applicable.

Source code in vllm/model_executor/model_loader/base_loader.py 
158
159
160
161
162
163
164
165
166
167
168
def log_online_quantization(vllm_config: VllmConfig) -> None:
    """Log the online-quantized layer count and types, when applicable."""
    quant_config = _get_online_quantization_config(vllm_config)
    if quant_config is None:
        return

    logger.info(
        "Quantizing %d layers of types: %s",
        len(quant_config.quantized_layers),
        "; ".join(quant_config.quantized_layer_summaries),
    )

log_online_quantization_time(vllm_config)

Log online quantization processing time, when applicable.

Source code in vllm/model_executor/model_loader/base_loader.py 
171
172
173
174
175
176
177
178
179
180
def log_online_quantization_time(vllm_config: VllmConfig) -> None:
    """Log online quantization processing time, when applicable."""
    quant_config = _get_online_quantization_config(vllm_config)
    if quant_config is None:
        return

    logger.info(
        "Online quantization as part of model loading took %.2f seconds",
        quant_config.online_quantization_time,
    )