first commit

vllm_br/model_executor/model_loader/__init__.py

@@ -0,0 +1,17 @@
+################################################################################
+# 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.
+#
+################################################################################
+
+from . import default_loader  # noqa: F401

vllm_br/model_executor/model_loader/default_loader.py

@@ -0,0 +1,83 @@
+################################################################################
+# 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

vllm_br/model_executor/model_loader/utils.py

@@ -0,0 +1,62 @@
+################################################################################
+# 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 torch
+from torch import nn
+
+from vllm.attention import Attention
+from vllm.config import ModelConfig
+from vllm.model_executor.layers.linear import QKVCrossParallelLinear
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizeMethodBase)
+
+
+def process_weights_after_loading(model: nn.Module, model_config: ModelConfig,
+                                  target_device: torch.device) -> None:
+    for _, module in model.named_modules():
+        if isinstance(module, QKVCrossParallelLinear):
+            # NOTE(Isotr0py): special case for cross QKV layer because
+            # q and kv proj aren't registered as submodules intentionally
+            module.process_weights_after_loading()
+            torch.supa.empty_cache()
+
+            continue
+        quant_method = getattr(module, "quant_method", None)
+        if isinstance(quant_method, QuantizeMethodBase):
+            # When quant methods need to process weights after loading
+            # (for repacking, quantizing, etc), they expect parameters
+            # to be on the global target device. This scope is for the
+            # case where cpu offloading is used, where we will move the
+            # parameters onto device for processing and back off after.
+
+            # with device_loading_context(module, target_device):
+            quant_method.weight_type = model_config.weight_type
+            quant_method.use_ds_mla = model_config.use_ds_mla
+            quant_method.use_ds_mla_sparse = model_config.use_ds_mla_sparse
+            quant_method.process_weights_after_loading(module)
+
+            torch.supa.empty_cache()
+
+    # Currently only used by MLA.
+    # NOTE: This intentionally happens after other modules so we can easily
+    # decompress the weights for MLA.
+    for _, module in model.named_modules():
+        if isinstance(module, Attention) and \
+            hasattr(module, "process_weights_after_loading"):
+            # TODO(lucas): see if there is a way to unify the signatures
+            # of process_weights_after_loading
+            module.process_weights_after_loading(model_config.dtype)
+            torch.supa.empty_cache()