325 lines
13 KiB
Python
325 lines
13 KiB
Python
|
#
|
||
|
|
# Copyright (c) 2025 Huawei Technologies 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 gc
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
from copy import deepcopy
|
||
|
|
from typing import List, Optional, Tuple
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
from vllm.config import LoadConfig, ModelConfig, VllmConfig
|
||
|
|
from vllm.logger import logger
|
||
|
|
from vllm.model_executor.model_loader import register_model_loader
|
||
|
|
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
|
||
|
|
from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
|
||
|
|
from vllm.model_executor.model_loader.utils import (
|
||
|
|
initialize_model, process_weights_after_loading, set_default_torch_dtype)
|
||
|
|
|
||
|
|
from .interaction.elastic import ElasticServer
|
||
|
|
from .load import elastic_load
|
||
|
|
from .utils import find_free_port, is_valid_path_prefix
|
||
|
|
|
||
|
|
|
||
|
|
@register_model_loader("netloader")
|
||
|
|
class ModelNetLoaderElastic(BaseModelLoader):
|
||
|
|
"""
|
||
|
|
A model loader that uses elastic loading for loading weights.
|
||
|
|
"""
|
||
|
|
source: Optional[List[dict]]
|
||
|
|
model_path: Optional[str]
|
||
|
|
listen_port: Optional[int]
|
||
|
|
int8_cache: str
|
||
|
|
int8_cache_name: Optional[List[str]]
|
||
|
|
output_prefix: Optional[str]
|
||
|
|
|
||
|
|
def __init__(self, load_config: LoadConfig):
|
||
|
|
"""
|
||
|
|
Initializes the ModelNetLoaderElastic with configuration.
|
||
|
|
|
||
|
|
Parameters:
|
||
|
|
- load_config: Configuration for loading the model.
|
||
|
|
"""
|
||
|
|
super().__init__(load_config)
|
||
|
|
|
||
|
|
config = None
|
||
|
|
|
||
|
|
# Try to read config file at first
|
||
|
|
extra = load_config.model_loader_extra_config
|
||
|
|
if extra and "CONFIG_FILE" in extra:
|
||
|
|
try:
|
||
|
|
logger.info(
|
||
|
|
f"Reading configs in file {load_config.model_loader_extra_config['CONFIG_FILE']} ..."
|
||
|
|
)
|
||
|
|
with open(extra["CONFIG_FILE"], 'r') as f:
|
||
|
|
config = json.load(f)
|
||
|
|
except FileNotFoundError:
|
||
|
|
logger.error("CONFIG_FILE not found")
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
logger.error("CONFIG_FILE is not a valid JSON file")
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(
|
||
|
|
f"Unexpected error while reading CONFIG_FILE: {e}")
|
||
|
|
|
||
|
|
if config is None and extra:
|
||
|
|
logger.info("Reading configs in model_loader_extra_config ...")
|
||
|
|
config = extra
|
||
|
|
config = config or {}
|
||
|
|
|
||
|
|
for key, attr, checker, caster, default in [
|
||
|
|
("SOURCE", "source", lambda v: isinstance(v, list), lambda v: v,
|
||
|
|
None),
|
||
|
|
("MODEL", "model_path", lambda v: isinstance(v, str), lambda v: v,
|
||
|
|
None),
|
||
|
|
("LISTEN_PORT", "listen_port", lambda v: isinstance(v, int) or
|
||
|
|
(isinstance(v, str) and v.isdigit()), lambda v: int(v), None),
|
||
|
|
("INT8_CACHE", "int8_cache", lambda v: isinstance(v, str) and v.
|
||
|
|
lower() in ['hbm', 'dram', 'no'], lambda v: v.lower(), 'no'),
|
||
|
|
("INT8_CACHE_NAME", "int8_cache_name",
|
||
|
|
lambda v: isinstance(v, list), lambda v: v, None),
|
||
|
|
("OUTPUT_PREFIX", "output_prefix",
|
||
|
|
lambda v: isinstance(v, str) and is_valid_path_prefix(v),
|
||
|
|
lambda v: v, None),
|
||
|
|
]:
|
||
|
|
v = config.get(key, default)
|
||
|
|
if not checker(v):
|
||
|
|
v = default
|
||
|
|
else:
|
||
|
|
v = caster(v)
|
||
|
|
setattr(self, attr, v)
|
||
|
|
|
||
|
|
logger.info(
|
||
|
|
"Initializing elastic Netloader with config: "
|
||
|
|
"MODEL=%s, LISTEN_PORT=%s,"
|
||
|
|
"SOURCE=%s, INT8_CACHE=%s, INT8_CACHE_NAME=%s,"
|
||
|
|
"OUTPUT_PREFIX=%s)",
|
||
|
|
self.model_path,
|
||
|
|
self.listen_port,
|
||
|
|
self.source,
|
||
|
|
self.int8_cache,
|
||
|
|
self.int8_cache_name,
|
||
|
|
self.output_prefix,
|
||
|
|
)
|
||
|
|
|
||
|
|
def load_model(self, vllm_config: VllmConfig,
|
||
|
|
model_config: ModelConfig) -> nn.Module:
|
||
|
|
"""
|
||
|
|
Loads the model using the specified configuration.
|
||
|
|
|
||
|
|
Parameters:
|
||
|
|
- vllm_config: Configuration for the VLLM.
|
||
|
|
- model_config: Configuration for the model.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
- The loaded model.
|
||
|
|
"""
|
||
|
|
|
||
|
|
device_config = vllm_config.device_config
|
||
|
|
parallel_config = vllm_config.parallel_config
|
||
|
|
|
||
|
|
need_process_weights_after_loading = False
|
||
|
|
|
||
|
|
if self.model_path is None:
|
||
|
|
self.model_path = model_config.model
|
||
|
|
logger.info(f"model_path is set to {self.model_path}")
|
||
|
|
|
||
|
|
device_id = torch.distributed.get_rank()
|
||
|
|
|
||
|
|
if (self.source is None or not isinstance(self.source, list)
|
||
|
|
or device_id not in [
|
||
|
|
one_device["device_id"] for one_device in self.source if
|
||
|
|
isinstance(one_device, dict) and "device_id" in one_device
|
||
|
|
]):
|
||
|
|
logger.warning(
|
||
|
|
"Did not get valid source info, use DefaultModelLoader")
|
||
|
|
model, need_process_weights_after_loading = self.revert_to_default(
|
||
|
|
model_config, vllm_config, device_config)
|
||
|
|
|
||
|
|
else:
|
||
|
|
target_device = torch.device(device_config.device)
|
||
|
|
|
||
|
|
vllm_config_backup = deepcopy(vllm_config)
|
||
|
|
model_config_backup = deepcopy(model_config)
|
||
|
|
|
||
|
|
with set_default_torch_dtype(model_config.dtype):
|
||
|
|
with target_device:
|
||
|
|
model = initialize_model(vllm_config=vllm_config,
|
||
|
|
model_config=model_config)
|
||
|
|
|
||
|
|
start_elastic_load = time.perf_counter()
|
||
|
|
model = elastic_load(
|
||
|
|
model=model,
|
||
|
|
device_id=device_id,
|
||
|
|
model_path=self.model_path,
|
||
|
|
sources=self.source,
|
||
|
|
tp=parallel_config.tensor_parallel_size,
|
||
|
|
pp=parallel_config.pipeline_parallel_size,
|
||
|
|
)
|
||
|
|
end_elastic_load = time.perf_counter()
|
||
|
|
logger.info(
|
||
|
|
f"Elastic load time: {end_elastic_load - start_elastic_load}, rank: {device_id}"
|
||
|
|
)
|
||
|
|
need_process_weights_after_loading = True
|
||
|
|
|
||
|
|
if model is None:
|
||
|
|
logger.warning(
|
||
|
|
"Netloader elastic loading fails, use load format DefaultModelLoader"
|
||
|
|
)
|
||
|
|
|
||
|
|
vllm_config = vllm_config_backup
|
||
|
|
model_config = model_config_backup
|
||
|
|
|
||
|
|
del model
|
||
|
|
gc.collect()
|
||
|
|
if device_config.device_type == 'npu':
|
||
|
|
logger.info("Empty NPU cache")
|
||
|
|
torch.npu.empty_cache()
|
||
|
|
elif device_config.device_type == 'cuda':
|
||
|
|
logger.info("Empty CUDA cache")
|
||
|
|
torch.cuda.empty_cache()
|
||
|
|
|
||
|
|
model, need_process_weights_after_loading = self.revert_to_default(
|
||
|
|
model_config, vllm_config, device_config)
|
||
|
|
|
||
|
|
start_elastic_server = time.perf_counter()
|
||
|
|
# start elastic server
|
||
|
|
if model is not None and (
|
||
|
|
(self.listen_port and self.listen_port in range(1024, 65535)) or
|
||
|
|
(self.listen_port is None)):
|
||
|
|
from vllm.utils import get_ip
|
||
|
|
driver_ip = get_ip()
|
||
|
|
|
||
|
|
if driver_ip == '0.0.0.0':
|
||
|
|
logger.error(
|
||
|
|
"Driver IP is not set, skip to start Netloader server")
|
||
|
|
else:
|
||
|
|
if self.listen_port is None:
|
||
|
|
self.listen_port = find_free_port()
|
||
|
|
else:
|
||
|
|
self.listen_port += device_id
|
||
|
|
|
||
|
|
logger.info(
|
||
|
|
f"Start elastic Netloader server, rank: {device_id}, listen port: {driver_ip}:{self.listen_port}"
|
||
|
|
)
|
||
|
|
|
||
|
|
if self.output_prefix is not None:
|
||
|
|
try:
|
||
|
|
with open(self.output_prefix + str(device_id) + '.txt',
|
||
|
|
'w') as file:
|
||
|
|
file.write(f"{driver_ip}:{self.listen_port}")
|
||
|
|
logger.info(
|
||
|
|
f"Successfully wrote server address to file: {self.output_prefix + str(device_id)}"
|
||
|
|
)
|
||
|
|
except FileNotFoundError:
|
||
|
|
logger.error(
|
||
|
|
f"File path {self.output_prefix + str(device_id)} does not exist."
|
||
|
|
)
|
||
|
|
except PermissionError:
|
||
|
|
logger.error(
|
||
|
|
f"No permission to write to file {self.output_prefix + str(device_id)}."
|
||
|
|
)
|
||
|
|
except IOError as e:
|
||
|
|
logger.error(
|
||
|
|
f"I/O error occurred while writing to file {self.output_prefix + str(device_id)}: {e}"
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Unknown error: {e}")
|
||
|
|
|
||
|
|
try:
|
||
|
|
assert isinstance(
|
||
|
|
self.listen_port, int
|
||
|
|
), f"listen port should be int but get {self.listen_port}"
|
||
|
|
|
||
|
|
elastic_server = ElasticServer(
|
||
|
|
driver_ip, self.listen_port, model, device_id,
|
||
|
|
self.model_path, parallel_config.tensor_parallel_size,
|
||
|
|
parallel_config.pipeline_parallel_size,
|
||
|
|
self.int8_cache, self.int8_cache_name)
|
||
|
|
elastic_server.start()
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(
|
||
|
|
f"Failed to start Netloader server for rank: {device_id}, details: {e}"
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
logger.info("Skip to start Netloader server")
|
||
|
|
|
||
|
|
end_elastic_server = time.perf_counter()
|
||
|
|
logger.info(
|
||
|
|
f"Elastic server start time: {end_elastic_server - start_elastic_server}, rank: {device_id}"
|
||
|
|
)
|
||
|
|
|
||
|
|
if need_process_weights_after_loading:
|
||
|
|
process_weights_after_loading(model, model_config,
|
||
|
|
torch.device(device_config.device))
|
||
|
|
|
||
|
|
if model is None:
|
||
|
|
logger.error("NetLoader elastic loads model fails")
|
||
|
|
return None
|
||
|
|
|
||
|
|
return model.eval()
|
||
|
|
|
||
|
|
def revert_to_default(self, model_config, vllm_config,
|
||
|
|
device_config) -> Tuple[nn.Module, bool]:
|
||
|
|
"""
|
||
|
|
Reverts to the default model loading logic when elastic loading fails or is not applicable.
|
||
|
|
|
||
|
|
This method resets the loader's extra config and load format to defaults,
|
||
|
|
then delegates model loading to a DefaultModelLoader.
|
||
|
|
If quantization is enabled, it will load the model and then run the
|
||
|
|
processing of weights (i.e. applying quantization adjustments) before returning.
|
||
|
|
|
||
|
|
Parameters:
|
||
|
|
- model_config: Configuration describing model architecture, quantization, etc.
|
||
|
|
- vllm_config: Configuration for vLLM (device, parallelism, dtype, etc).
|
||
|
|
- device_config: Configuration for the target device (device type, device id, etc).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
- A tuple (model, need_process_weights_after_loading):
|
||
|
|
* model: The loaded `nn.Module` under default loading logic.
|
||
|
|
* need_process_weights_after_loading: A boolean flag indicating whether
|
||
|
|
weights post-processing (e.g. quantization adjustments) still needs to be applied.
|
||
|
|
"""
|
||
|
|
self.load_config.model_loader_extra_config = {}
|
||
|
|
self.load_config.load_format = "auto"
|
||
|
|
default_model_loader = DefaultModelLoader(self.load_config)
|
||
|
|
|
||
|
|
if model_config.quantization is None:
|
||
|
|
model = default_model_loader.load_model(vllm_config=vllm_config,
|
||
|
|
model_config=model_config)
|
||
|
|
need_process_weights_after_loading = False
|
||
|
|
else:
|
||
|
|
logger.warning(
|
||
|
|
"Quantization is set, netloader use DefaultModelLoader with process_weights_after_loading "
|
||
|
|
)
|
||
|
|
need_process_weights_after_loading = True
|
||
|
|
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,
|
||
|
|
model_config=model_config)
|
||
|
|
default_model_loader.load_weights(model, model_config)
|
||
|
|
model = model.eval()
|
||
|
|
|
||
|
|
return model, need_process_weights_after_loading
|
||
|
|
|
||
|
|
def download_model(self, model_config: ModelConfig) -> None:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def load_weights(self, model: nn.Module,
|
||
|
|
model_config: ModelConfig) -> None:
|
||
|
|
pass
|