84 lines
3.3 KiB
Python
84 lines
3.3 KiB
Python
################################################################################
|
|
# Copyright(c)2020-2025 Shanghai Biren Technology Co., Ltd. All rights reserved.
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
################################################################################
|
|
|
|
import time
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from vllm.config import ModelConfig, VllmConfig
|
|
from vllm.logger import logger
|
|
from vllm.model_executor.model_loader import DefaultModelLoader
|
|
from vllm.model_executor.model_loader.utils import (initialize_model,
|
|
set_default_torch_dtype)
|
|
from .utils import process_weights_after_loading
|
|
|
|
|
|
def load_model(self, vllm_config: VllmConfig,
|
|
model_config: ModelConfig) -> nn.Module:
|
|
device_config = vllm_config.device_config
|
|
target_device = torch.device(device_config.device)
|
|
with set_default_torch_dtype(model_config.dtype):
|
|
model = initialize_model(vllm_config=vllm_config,
|
|
model_config=model_config)
|
|
# NOTE: on SUPA, with device context may not take effect, mamully to device
|
|
# model = model.to(target_device)
|
|
|
|
# NOTE: move moe weight to cpu, reduce device memory usage, more layers can be moved to cpu if necessary
|
|
moe_packed_weights = [
|
|
"mlp.experts.w13_weight_packed",
|
|
"mlp.experts.w2_weight_packed",
|
|
"mlp.gate_up_proj",
|
|
"mlp.down_proj",
|
|
"mlp.experts",
|
|
"self_attn.qkv_proj",
|
|
"self_attn.o_proj",
|
|
]
|
|
for name, module in model.named_parameters():
|
|
if any(s in name for s in moe_packed_weights):
|
|
module.data = module.to("cpu")
|
|
else:
|
|
module.data = module.to(target_device)
|
|
|
|
torch.supa.empty_cache()
|
|
|
|
weights_to_load = {name for name, _ in model.named_parameters()}
|
|
loaded_weights = model.load_weights(
|
|
self.get_all_weights(model_config, model))
|
|
|
|
torch.supa.empty_cache()
|
|
|
|
self.counter_after_loading_weights = time.perf_counter()
|
|
logger.info(
|
|
"Loading weights took %.2f seconds",
|
|
self.counter_after_loading_weights -
|
|
self.counter_before_loading_weights)
|
|
# We only enable strict check for non-quantized models
|
|
# that have loaded weights tracking currently.
|
|
if model_config.quantization is None and loaded_weights is not None:
|
|
weights_not_loaded = weights_to_load - loaded_weights
|
|
if weights_not_loaded:
|
|
raise ValueError("Following weights were not initialized from "
|
|
f"checkpoint: {weights_not_loaded}")
|
|
|
|
process_weights_after_loading(model, model_config, target_device)
|
|
|
|
torch.cuda.empty_cache()
|
|
return model.eval()
|
|
|
|
|
|
DefaultModelLoader.load_model = load_model
|