38 lines
1.5 KiB
Python
38 lines
1.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from vllm.config import LoadConfig, ModelConfig, VllmConfig
|
|
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
|
|
from vllm.model_executor.model_loader.utils import (
|
|
initialize_model, process_weights_after_loading, set_default_torch_dtype)
|
|
from vllm.model_executor.model_loader.weight_utils import (
|
|
initialize_dummy_weights)
|
|
|
|
|
|
class DummyModelLoader(BaseModelLoader):
|
|
"""Model loader that will set model weights to random values."""
|
|
|
|
def __init__(self, load_config: LoadConfig):
|
|
super().__init__(load_config)
|
|
if load_config.model_loader_extra_config:
|
|
raise ValueError(f"Model loader extra config is not supported for "
|
|
f"load format {load_config.load_format}")
|
|
|
|
def download_model(self, model_config: ModelConfig) -> None:
|
|
pass # Nothing to download
|
|
|
|
def load_model(self, vllm_config: VllmConfig) -> nn.Module:
|
|
device_config = vllm_config.device_config
|
|
model_config = vllm_config.model_config
|
|
target_device = torch.device(device_config.device)
|
|
with set_default_torch_dtype(model_config.dtype):
|
|
with target_device:
|
|
model = initialize_model(vllm_config=vllm_config)
|
|
# NOTE(woosuk): For accurate performance evaluation, we assign
|
|
# random values to the weights.
|
|
initialize_dummy_weights(model)
|
|
|
|
process_weights_after_loading(model, model_config, target_device)
|
|
return model.eval()
|