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 
 26
 27
 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
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

    @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)

            # 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)

        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 
48
49
50
51
52
53
54
55
56
57
58
59
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 
37
38
39
40
@abstractmethod
def download_model(self, model_config: ModelConfig) -> None:
    """Download a model so that it can be immediately loaded."""
    raise NotImplementedError

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 
32
33
34
35
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 
 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
@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)

        # 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)

    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 
42
43
44
45
46
@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

log_model_inspection(model)

Log model structure if VLLM_LOG_MODEL_INSPECTION=1.

Source code in vllm/model_executor/model_loader/base_loader.py 
108
109
110
111
112
113
114
115
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 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def log_online_quantization(vllm_config: VllmConfig) -> None:
    """Log the online-quantized layer count and types, when applicable."""
    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

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