wehub-resource-sync
285 lines
11 KiB
Python
285 lines
11 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import ray
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
|
|
|
|
from swift.rlhf_trainers.args_mixin import VllmArguments
|
|
from swift.utils.logger import get_logger
|
|
|
|
if TYPE_CHECKING:
|
|
from ..resource_pool import ResourcePool
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class RolloutMode(str, Enum):
|
|
HYBRID = 'hybrid'
|
|
STANDALONE = 'standalone'
|
|
|
|
|
|
@dataclass
|
|
class VllmEngineConfig(VllmArguments):
|
|
model: str = ''
|
|
sleep_level: int = 0
|
|
vllm_enable_lora: bool = False
|
|
trust_remote_code: bool = True
|
|
dtype: str = 'auto'
|
|
load_format: str = 'auto'
|
|
|
|
# override
|
|
vllm_enable_prefix_caching: bool = True
|
|
|
|
def __post_init__(self):
|
|
VllmArguments.__post_init__(self)
|
|
|
|
@classmethod
|
|
def from_rollout_cfg(cls, rollout_cfg: Dict[str, Any], *, sleep_level: int = 0) -> 'VllmEngineConfig':
|
|
"""Build from merged rollout config dict."""
|
|
cfg = rollout_cfg or {}
|
|
known_fields = {f.name for f in cls.__dataclass_fields__.values()}
|
|
|
|
kwargs: Dict[str, Any] = {}
|
|
for key, val in cfg.items():
|
|
if key in known_fields and val is not None:
|
|
kwargs[key] = val
|
|
|
|
if sleep_level > 0:
|
|
kwargs['sleep_level'] = sleep_level
|
|
|
|
if cfg.get('tuner_type', 'full') == 'lora':
|
|
kwargs.setdefault('vllm_enable_lora', True)
|
|
|
|
return cls(**kwargs)
|
|
|
|
def to_launch_kwargs(self, rollout_mode: str, template_kwargs: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
|
kw: Dict[str, Any] = {
|
|
'model_id': self.model,
|
|
'dtype': self.dtype,
|
|
'rollout_mode': rollout_mode,
|
|
'tensor_parallel_size': self.vllm_tensor_parallel_size,
|
|
'gpu_memory_utilization': self.vllm_gpu_memory_utilization,
|
|
'max_num_seqs': self.vllm_max_num_seqs,
|
|
'enforce_eager': self.vllm_enforce_eager,
|
|
'trust_remote_code': self.trust_remote_code,
|
|
'load_format': self.load_format,
|
|
'enable_sleep_mode': (self.sleep_level > 0),
|
|
'enable_lora': self.vllm_enable_lora,
|
|
'max_lora_rank': self.vllm_max_lora_rank,
|
|
'data_parallel_size': self.vllm_data_parallel_size,
|
|
'enable_prefix_caching': self.vllm_enable_prefix_caching,
|
|
}
|
|
if self.vllm_max_model_len is not None:
|
|
kw['max_model_len'] = self.vllm_max_model_len
|
|
if template_kwargs:
|
|
kw['template_kwargs'] = template_kwargs
|
|
extra = self.vllm_engine_kwargs
|
|
if isinstance(extra, dict):
|
|
kw.update(extra)
|
|
return kw
|
|
|
|
|
|
class RolloutReplica:
|
|
"""One vLLM rollout replica on Ray (single or multi-node)."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: VllmEngineConfig,
|
|
mode: RolloutMode = RolloutMode.HYBRID,
|
|
replica_rank: int = 0,
|
|
template_kwargs: Optional[Dict[str, Any]] = None,
|
|
actor_name_prefix: str = 'swift_rollout_server',
|
|
) -> None:
|
|
self.config = config
|
|
self.mode = mode
|
|
self.replica_rank = replica_rank
|
|
self.template_kwargs = template_kwargs
|
|
self.actor_name_prefix = actor_name_prefix
|
|
self._servers: List[Any] = []
|
|
|
|
@classmethod
|
|
def create_replicas(
|
|
cls,
|
|
rollout_cfg: Dict[str, Any],
|
|
rollout_gpus: int,
|
|
pool: 'ResourcePool',
|
|
is_hybrid: bool,
|
|
sleep_level: int = 0,
|
|
template_kwargs: Optional[Dict[str, Any]] = None,
|
|
actor_name_prefix: str = 'swift_rollout_server',
|
|
) -> List['RolloutReplica']:
|
|
"""Factory: create all rollout replicas from pipeline config.
|
|
|
|
Uses two-phase initialization for parallelism:
|
|
Phase 1 — spawn all Ray actors (fast, non-blocking)
|
|
Phase 2 — launch_server on all replicas concurrently
|
|
"""
|
|
config = VllmEngineConfig.from_rollout_cfg(rollout_cfg, sleep_level=sleep_level)
|
|
world_size_per_replica = config.vllm_tensor_parallel_size * config.vllm_data_parallel_size
|
|
if world_size_per_replica > rollout_gpus:
|
|
raise ValueError(f'tp*dp ({world_size_per_replica}) exceeds rollout GPUs ({rollout_gpus})')
|
|
if rollout_gpus % world_size_per_replica != 0:
|
|
raise ValueError(f'rollout GPUs ({rollout_gpus}) must be divisible by '
|
|
f'tp*dp ({world_size_per_replica})')
|
|
n_replicas = rollout_gpus // world_size_per_replica
|
|
mode = RolloutMode.HYBRID if is_hybrid else RolloutMode.STANDALONE
|
|
|
|
replicas: List['RolloutReplica'] = []
|
|
bundle_infos = pool.bundle_infos
|
|
for i in range(n_replicas):
|
|
offset = i * world_size_per_replica
|
|
replica_infos = bundle_infos[offset:offset + world_size_per_replica]
|
|
nodes = {info[0] for info in replica_infos}
|
|
gpus_per_node = (world_size_per_replica if len(nodes) == 1 else world_size_per_replica // len(nodes))
|
|
replica = cls(
|
|
config, mode=mode, replica_rank=i, template_kwargs=template_kwargs, actor_name_prefix=actor_name_prefix)
|
|
replica._spawn_actors(replica_infos, gpus_per_node)
|
|
replicas.append(replica)
|
|
|
|
cls._parallel_launch_all(replicas)
|
|
|
|
logger.info('Rollout: %d replica(s) in %s mode (tp=%d, dp=%d, total_gpus=%d)', n_replicas, mode.value.upper(),
|
|
config.vllm_tensor_parallel_size, config.vllm_data_parallel_size, rollout_gpus)
|
|
return replicas
|
|
|
|
@classmethod
|
|
def _parallel_launch_all(cls, replicas: List['RolloutReplica']) -> None:
|
|
"""Phase 2: launch vLLM engines on all replicas in parallel."""
|
|
all_refs = []
|
|
for replica in replicas:
|
|
refs = replica._launch_engines_async()
|
|
all_refs.extend(refs)
|
|
if all_refs:
|
|
ray.get(all_refs)
|
|
for replica in replicas:
|
|
logger.info('RolloutReplica[replica=%d, mode=%s]: launched %d server(s) (tp=%d, model=%s)',
|
|
replica.replica_rank, replica.mode.value, len(replica._servers),
|
|
replica.config.vllm_tensor_parallel_size, replica.config.model)
|
|
|
|
def _spawn_actors(
|
|
self,
|
|
worker_infos: List[Tuple[str, str]],
|
|
gpus_per_node: int,
|
|
) -> None:
|
|
"""Phase 1: create Ray actors without starting engines.
|
|
|
|
``num_gpus=0`` + ``NOSET_CVD`` + explicit visible-device env so
|
|
the actor sees exactly the GPUs from pool bundles; NodeAffinity
|
|
pins each actor to the correct node.
|
|
"""
|
|
from ray.runtime_env import RuntimeEnv
|
|
from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
|
|
from transformers.utils import is_torch_npu_available
|
|
|
|
from .vllm_server import VllmServer
|
|
|
|
visible_key = 'ASCEND_RT_VISIBLE_DEVICES' if is_torch_npu_available() else 'CUDA_VISIBLE_DEVICES'
|
|
|
|
node_groups = self._group_by_node(worker_infos)
|
|
self._nnodes = len(node_groups)
|
|
|
|
actor_cls = ray.remote(num_gpus=0, num_cpus=1)(VllmServer)
|
|
|
|
for node_rank, (node_id, gpu_ids) in enumerate(node_groups):
|
|
cvd = ','.join(gpu_ids)
|
|
env_vars: Dict[str, str] = {
|
|
'VLLM_USE_V1': '1',
|
|
'RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES': '1',
|
|
'RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES': '1',
|
|
'RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES': '1',
|
|
'RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES': '1',
|
|
'NCCL_CUMEM_ENABLE': '0',
|
|
visible_key: cvd,
|
|
}
|
|
handle = actor_cls.options(
|
|
scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=node_id, soft=False),
|
|
runtime_env=RuntimeEnv(env_vars=env_vars),
|
|
name=f'{self.actor_name_prefix}_{self.replica_rank}_{node_rank}',
|
|
max_concurrency=10,
|
|
).remote(
|
|
node_rank=node_rank,
|
|
nnodes=self._nnodes,
|
|
gpus_per_node=gpus_per_node,
|
|
cuda_visible_devices=cvd,
|
|
)
|
|
self._servers.append(handle)
|
|
|
|
def _launch_engines_async(self) -> List[ray.ObjectRef]:
|
|
"""Phase 2: issue launch_server calls, return ObjectRefs (non-blocking)."""
|
|
launch_kw = self.config.to_launch_kwargs(self.mode.value, template_kwargs=self.template_kwargs)
|
|
nnodes = getattr(self, '_nnodes', 1)
|
|
|
|
if nnodes > 1:
|
|
master_address, master_port, dp_rpc_port = ray.get(self._servers[0].get_master_address.remote())
|
|
refs = [
|
|
server.launch_server.remote(
|
|
master_address=master_address, master_port=master_port, dp_rpc_port=dp_rpc_port, **launch_kw)
|
|
for server in self._servers
|
|
]
|
|
else:
|
|
refs = [self._servers[0].launch_server.remote(**launch_kw)]
|
|
return refs
|
|
|
|
@staticmethod
|
|
def _group_by_node(worker_infos: List[Tuple[str, str]]) -> List[Tuple[str, List[str]]]:
|
|
"""Group worker infos by node, preserving order.
|
|
|
|
Returns list of ``(node_id, [accelerator_id, ...])`` in node
|
|
encounter order.
|
|
"""
|
|
ordered_nodes: List[str] = []
|
|
node_gpus: Dict[str, List[str]] = defaultdict(list)
|
|
for node_id, acc_id in worker_infos:
|
|
if node_id not in node_gpus:
|
|
ordered_nodes.append(node_id)
|
|
node_gpus[node_id].append(acc_id)
|
|
return [(nid, node_gpus[nid]) for nid in ordered_nodes]
|
|
|
|
@property
|
|
def primary(self) -> Any:
|
|
"""The node_rank=0 ``VllmServer`` actor handle.
|
|
|
|
Callers that need ``sleep`` / ``wake_up`` / ``reset_prefix_cache``
|
|
/ ``update_weights_ipc`` / ``update_weights_direct`` talk to
|
|
this handle directly.
|
|
"""
|
|
if not self._servers:
|
|
raise RuntimeError('RolloutReplica: not launched yet')
|
|
return self._servers[0]
|
|
|
|
@property
|
|
def servers(self) -> List[Any]:
|
|
"""All server actor handles (one per node)."""
|
|
return list(self._servers)
|
|
|
|
def sleep(self, level: int = 1) -> None:
|
|
"""Put all servers in this replica to sleep."""
|
|
refs = [server.sleep.remote(level) for server in self._servers]
|
|
ray.get(refs)
|
|
|
|
def wake_up(self, tags=None) -> None:
|
|
"""Wake all servers in this replica."""
|
|
refs = [server.wake_up.remote(tags=tags) for server in self._servers]
|
|
ray.get(refs)
|
|
|
|
def generate(
|
|
self,
|
|
infer_requests: List[Any],
|
|
request_config: Any = None,
|
|
) -> ray.ObjectRef:
|
|
"""Submit generation to the primary server, returns an ObjectRef."""
|
|
return self.primary.generate.remote(infer_requests, request_config)
|
|
|
|
def shutdown(self) -> None:
|
|
for server in self._servers:
|
|
try:
|
|
ray.get(server.shutdown.remote(), timeout=30)
|
|
except Exception as e: # noqa: BLE001
|
|
logger.warning('RolloutReplica shutdown error: %s', e)
|
|
try:
|
|
ray.kill(server, no_restart=True)
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
self._servers = []
|