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