334 lines
13 KiB
Python
334 lines
13 KiB
Python
#
|
|
# Copyright (c) 2026 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 os
|
|
import time
|
|
from copy import copy
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.nn import Module
|
|
from vllm.config import ModelConfig, VllmConfig
|
|
from vllm.config.load import LoadConfig
|
|
from vllm.distributed import get_tensor_model_parallel_rank
|
|
from vllm.distributed.parallel_state import get_ep_group, get_pp_group
|
|
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.utils import (
|
|
initialize_model,
|
|
process_weights_after_loading,
|
|
)
|
|
from vllm.utils.torch_utils import set_default_torch_dtype
|
|
|
|
from vllm_ascend.model_loader.rfork.rfork_worker import RForkWorker
|
|
|
|
|
|
def _is_mtp_hf_config(hf_config: object | None) -> bool:
|
|
if hf_config is None:
|
|
return False
|
|
|
|
model_type = getattr(hf_config, "model_type", None)
|
|
if isinstance(model_type, str) and model_type.lower().endswith("_mtp"):
|
|
return True
|
|
|
|
architectures = getattr(hf_config, "architectures", None)
|
|
if isinstance(architectures, str):
|
|
architectures = [architectures]
|
|
if not isinstance(architectures, (list, tuple)):
|
|
return False
|
|
|
|
return any(isinstance(architecture, str) and architecture.endswith("MTPModel") for architecture in architectures)
|
|
|
|
|
|
def _is_draft_model_config(model_config: object | None) -> bool:
|
|
if model_config is None:
|
|
return False
|
|
if getattr(model_config, "runner_type", None) == "draft":
|
|
return True
|
|
|
|
return any(
|
|
_is_mtp_hf_config(getattr(model_config, hf_config_attr, None))
|
|
for hf_config_attr in ("hf_config", "hf_text_config")
|
|
)
|
|
|
|
|
|
def _is_draft_model(vllm_config: VllmConfig, model_config: ModelConfig | None = None) -> bool:
|
|
return (
|
|
_is_draft_model_config(model_config)
|
|
or _is_draft_model_config(getattr(vllm_config, "model_config", None))
|
|
or _is_draft_model_config(getattr(vllm_config, "scheduler_config", None))
|
|
)
|
|
|
|
|
|
def _get_rfork_worker_attr(vllm_config: VllmConfig, model_config: ModelConfig) -> str:
|
|
return "rfork_draft_worker" if _is_draft_model(vllm_config, model_config) else "rfork_worker"
|
|
|
|
|
|
def _get_ep_rank(vllm_config: VllmConfig) -> int | None:
|
|
parallel_config = vllm_config.parallel_config
|
|
if not parallel_config.enable_expert_parallel or getattr(parallel_config, "is_moe_model", None) is False:
|
|
return None
|
|
|
|
try:
|
|
return get_ep_group().rank_in_group
|
|
except AssertionError as e:
|
|
raise RuntimeError("Expert parallelism is enabled, but the EP group is not initialized.") from e
|
|
|
|
|
|
def _get_pp_rank(vllm_config: VllmConfig) -> int | None:
|
|
if getattr(vllm_config.parallel_config, "pipeline_parallel_size", 1) <= 1:
|
|
return None
|
|
|
|
try:
|
|
return get_pp_group().rank_in_group
|
|
except AssertionError as e:
|
|
raise RuntimeError("Pipeline parallelism is enabled, but the PP group is not initialized.") from e
|
|
|
|
|
|
def _make_fallback_load_config(load_config: LoadConfig) -> LoadConfig:
|
|
fallback_load_config = copy(load_config)
|
|
fallback_load_config.load_format = "auto"
|
|
fallback_load_config.model_loader_extra_config = {}
|
|
return fallback_load_config
|
|
|
|
|
|
def _is_layer_sharding_enabled(vllm_config: VllmConfig) -> bool:
|
|
additional_config = getattr(vllm_config, "additional_config", None) or {}
|
|
return bool(additional_config.get("layer_sharding"))
|
|
|
|
|
|
def _is_dynamic_eplb_enabled(vllm_config: VllmConfig) -> bool:
|
|
parallel_config = getattr(vllm_config, "parallel_config", None)
|
|
if bool(getattr(parallel_config, "enable_eplb", False)):
|
|
return True
|
|
|
|
additional_config = getattr(vllm_config, "additional_config", None) or {}
|
|
eplb_config = additional_config.get("eplb_config", {})
|
|
if not isinstance(eplb_config, dict):
|
|
return False
|
|
return bool(eplb_config.get("dynamic_eplb") or eplb_config.get("expert_map_record_path"))
|
|
|
|
|
|
@register_model_loader("rfork")
|
|
class RForkModelLoader(BaseModelLoader):
|
|
def __init__(self, load_config: LoadConfig):
|
|
super().__init__(load_config)
|
|
config = load_config.model_loader_extra_config
|
|
if config is None:
|
|
config = {}
|
|
elif not isinstance(config, dict):
|
|
err_msg = "RFork requires --model-loader-extra-config to be a JSON object."
|
|
logger.error(err_msg)
|
|
raise RuntimeError(err_msg)
|
|
|
|
def _get_extra_config(key: str, default: str = "") -> str:
|
|
value = config.get(key)
|
|
if value is None or not isinstance(value, str):
|
|
value = os.environ.get(key.upper())
|
|
return value if isinstance(value, str) and value else default
|
|
|
|
def _get_extra_config_float(key: str, default: float) -> float:
|
|
value = config.get(key)
|
|
if value is None or isinstance(value, bool) or not isinstance(value, (int, float, str)):
|
|
value = os.environ.get(key.upper())
|
|
parsed_value = default
|
|
if isinstance(value, (int, float)):
|
|
parsed_value = float(value)
|
|
elif isinstance(value, str) and value:
|
|
try:
|
|
parsed_value = float(value)
|
|
except ValueError:
|
|
return default
|
|
|
|
if parsed_value <= 0:
|
|
return default
|
|
|
|
return parsed_value
|
|
|
|
self.model_url = _get_extra_config("model_url", "")
|
|
self.model_deploy_strategy_name = _get_extra_config("model_deploy_strategy_name", "")
|
|
self.scheduler_url = _get_extra_config("rfork_scheduler_url", "")
|
|
self.seed_timeout_sec = _get_extra_config_float("rfork_seed_timeout_sec", 5.0)
|
|
self.seed_key_separator = _get_extra_config("rfork_seed_key_separator", "$")
|
|
|
|
logger.info(
|
|
"Initializing rfork with config: "
|
|
"MODEL_URL=%s, MODEL_DEPLOY_STRATEGY_NAME=%s, "
|
|
"SCHEDULER_URL=%s, SEED_TIMEOUT_SEC=%s, "
|
|
"SEED_KEY_SEPARATOR=%s",
|
|
self.model_url,
|
|
self.model_deploy_strategy_name,
|
|
self.scheduler_url,
|
|
self.seed_timeout_sec,
|
|
self.seed_key_separator,
|
|
)
|
|
|
|
def download_model(self, model_config: ModelConfig) -> None:
|
|
raise NotImplementedError
|
|
|
|
def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
|
|
raise NotImplementedError
|
|
|
|
def _ensure_rfork_worker(self, vllm_config: VllmConfig, model_config: ModelConfig) -> RForkWorker:
|
|
worker_attr = _get_rfork_worker_attr(vllm_config, model_config)
|
|
rfork_worker = getattr(self.load_config, worker_attr, None)
|
|
if rfork_worker is None:
|
|
kv_transfer_config = vllm_config.kv_transfer_config
|
|
disaggregation_mode = "kv_both" if kv_transfer_config is None else str(kv_transfer_config.kv_role)
|
|
is_draft_model = _is_draft_model(vllm_config, model_config)
|
|
device_id = torch.distributed.get_rank()
|
|
pp_rank = _get_pp_rank(vllm_config)
|
|
ep_rank = _get_ep_rank(vllm_config)
|
|
rfork_worker = RForkWorker(
|
|
disaggregation_mode=disaggregation_mode,
|
|
node_rank=vllm_config.parallel_config.node_rank,
|
|
tp_rank=get_tensor_model_parallel_rank(),
|
|
device_id=device_id,
|
|
scheduler_url=self.scheduler_url,
|
|
model_url=self.model_url,
|
|
model_deploy_strategy_name=self.model_deploy_strategy_name,
|
|
seed_timeout_sec=self.seed_timeout_sec,
|
|
seed_key_separator=self.seed_key_separator,
|
|
is_draft_model=is_draft_model,
|
|
pp_rank=pp_rank,
|
|
ep_rank=ep_rank,
|
|
)
|
|
setattr(self.load_config, worker_attr, rfork_worker)
|
|
logger.info(
|
|
"RFork worker initialized, load_format=rfork, is_draft_model=%s, worker_attr=%s",
|
|
is_draft_model,
|
|
worker_attr,
|
|
)
|
|
return rfork_worker
|
|
|
|
def _requires_processed_layout_transfer(self, model_config: ModelConfig) -> bool:
|
|
return getattr(model_config, "quantization", None) is not None
|
|
|
|
def load_model(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
model_config: ModelConfig,
|
|
prefix: str = "",
|
|
) -> Module | None:
|
|
device_config = vllm_config.device_config
|
|
load_config = self.load_config
|
|
load_device = device_config.device if load_config.device is None else load_config.device
|
|
target_device = torch.device(load_device)
|
|
|
|
with set_default_torch_dtype(model_config.dtype):
|
|
need_del = False
|
|
bypass_reason = None
|
|
if _is_layer_sharding_enabled(vllm_config):
|
|
bypass_reason = "additional_config.layer_sharding"
|
|
elif _is_dynamic_eplb_enabled(vllm_config):
|
|
bypass_reason = "dynamic EPLB"
|
|
|
|
if bypass_reason is not None:
|
|
logger.warning(
|
|
"RFork transfer is disabled when %s is enabled; using the default model loader.",
|
|
bypass_reason,
|
|
)
|
|
fallback_load_config = _make_fallback_load_config(self.load_config)
|
|
|
|
from vllm.model_executor.model_loader import get_model
|
|
|
|
try:
|
|
return get_model(
|
|
vllm_config=vllm_config,
|
|
model_config=model_config,
|
|
load_config=fallback_load_config,
|
|
prefix=prefix,
|
|
)
|
|
except Exception:
|
|
logger.exception("RFork disabled for %s, but default loader failed.", bypass_reason)
|
|
raise
|
|
|
|
rfork_worker = self._ensure_rfork_worker(vllm_config, model_config)
|
|
processed_layout_transfer = self._requires_processed_layout_transfer(model_config)
|
|
try:
|
|
if not rfork_worker.is_seed_available():
|
|
raise RuntimeError("seed is not available.")
|
|
|
|
with target_device:
|
|
model = initialize_model(
|
|
vllm_config=vllm_config,
|
|
model_config=model_config,
|
|
prefix=prefix,
|
|
)
|
|
need_del = True
|
|
|
|
if processed_layout_transfer:
|
|
logger.info("RFork uses post-load tensor layout transfer for quantized model.")
|
|
process_weights_after_loading(model, model_config, target_device)
|
|
|
|
weight_load_start_time = time.perf_counter()
|
|
if not rfork_worker.pre_transfer(model):
|
|
raise RuntimeError("pre_transfer failed.")
|
|
if not rfork_worker.transfer(model):
|
|
raise RuntimeError("transfer failed.")
|
|
if not rfork_worker.post_transfer():
|
|
raise RuntimeError("post_transfer failed.")
|
|
logger.info(
|
|
"Loading model weights took %.2f seconds",
|
|
time.perf_counter() - weight_load_start_time,
|
|
)
|
|
|
|
rfork_worker.start_seed_service(model)
|
|
if not processed_layout_transfer:
|
|
process_weights_after_loading(model, model_config, target_device)
|
|
|
|
return model.eval()
|
|
except Exception as e:
|
|
logger.warning("RFork transfer failed: %s, clean up and fall back to default loader", e)
|
|
|
|
rfork_worker.post_transfer()
|
|
rfork_worker.reset_transfer_state()
|
|
|
|
if need_del:
|
|
del model
|
|
gc.collect()
|
|
torch.npu.empty_cache()
|
|
for _ in range(3):
|
|
gc.collect()
|
|
torch.npu.empty_cache()
|
|
|
|
fallback_load_config = _make_fallback_load_config(self.load_config)
|
|
|
|
from vllm.model_executor.model_loader import get_model
|
|
|
|
try:
|
|
model = get_model(
|
|
vllm_config=vllm_config,
|
|
model_config=model_config,
|
|
load_config=fallback_load_config,
|
|
prefix=prefix,
|
|
)
|
|
except Exception:
|
|
logger.exception("RFork fallback default loader failed.")
|
|
raise
|
|
|
|
try:
|
|
rfork_worker.reset_transfer_state()
|
|
rfork_worker.start_seed_service(model)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"Fallback model loaded, but start_seed_service failed: %s",
|
|
e,
|
|
)
|
|
return model
|