357 lines
15 KiB
Python
357 lines
15 KiB
Python
################################################################################
|
|
# 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 os
|
|
from collections import defaultdict
|
|
from typing import TYPE_CHECKING, Dict, List
|
|
|
|
import vllm.envs as envs
|
|
from vllm.executor.ray_distributed_executor import RayDistributedExecutor
|
|
from vllm.executor.ray_utils import RayWorkerWrapper, ray
|
|
# from vllm.logger import init_logger
|
|
from vllm.platforms import current_platform
|
|
from vllm.ray.ray_env import get_env_vars_to_copy
|
|
from vllm.utils import get_distributed_init_method, get_ip, get_open_port
|
|
from vllm_br import envs as envs_br
|
|
|
|
if ray is not None:
|
|
from ray.actor import ActorHandle
|
|
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
|
else:
|
|
ActorHandle = None
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.util.placement_group import PlacementGroup
|
|
|
|
from vllm.executor.ray_distributed_executor import RayWorkerMetaData, logger
|
|
|
|
|
|
def get_supernode_pp_tp_global_rank_map(tp_size, pp_size):
|
|
rank_map = {}
|
|
tp_driver_rank = []
|
|
for pp_rank in range(pp_size):
|
|
for tp_rank in range(tp_size):
|
|
# PP=2, TP=8
|
|
# pp_tp_workers = [[0, 1, 2, 3, 8, 9, 10, 11], [4, 5, 6, 7, 12, 13, 14, 15]]
|
|
if tp_rank < 4 and pp_rank < 1:
|
|
rank = (pp_rank * pp_size) + tp_rank
|
|
elif tp_rank >= 4 and pp_rank < 1:
|
|
rank = (pp_rank * pp_size) + tp_rank + 4
|
|
elif tp_rank < 4 and pp_rank >= 1:
|
|
rank = (pp_rank * pp_size) + tp_rank + 2
|
|
elif tp_rank >= 4 and pp_rank >= 1:
|
|
rank = (pp_rank * pp_size) + tp_rank + 6
|
|
rank_map[(pp_rank, tp_rank)] = rank
|
|
if tp_rank == 0:
|
|
tp_driver_rank.append(rank)
|
|
return rank_map, tp_driver_rank
|
|
|
|
|
|
def _init_workers_ray_br(self, placement_group: "PlacementGroup",
|
|
**ray_remote_kwargs):
|
|
num_gpus = envs.VLLM_RAY_PER_WORKER_GPUS
|
|
if envs_br.VLLM_BR_ENABLE_TP_GROUPS_IN_SUPERNODE:
|
|
rank_map, tp_driver_rank = get_supernode_pp_tp_global_rank_map(
|
|
self.parallel_config.tensor_parallel_size,
|
|
self.parallel_config.pipeline_parallel_size)
|
|
|
|
# The driver dummy worker does not actually use any resources.
|
|
# It holds the resource for the driver worker.
|
|
self.driver_dummy_worker = None
|
|
# The remaining workers are the actual ray actors.
|
|
self.workers = []
|
|
|
|
# Used in ray compiled DAG: indexed first by PP rank,
|
|
# and then TP rank. In other words, the inner list is
|
|
# the TP group of workers for a PP rank.
|
|
self.pp_tp_workers = []
|
|
|
|
if self.parallel_config.ray_workers_use_nsight:
|
|
ray_remote_kwargs = self._configure_ray_workers_use_nsight(
|
|
ray_remote_kwargs)
|
|
|
|
logger.info("use_ray_spmd_worker: %s", self.use_ray_spmd_worker)
|
|
|
|
# Create the workers.
|
|
bundle_indices: List[int]
|
|
if envs.VLLM_RAY_BUNDLE_INDICES:
|
|
# Use the bundle indices specified by the user.
|
|
bundle_indices = list(map(int,
|
|
envs.VLLM_RAY_BUNDLE_INDICES.split(",")))
|
|
assert len(bundle_indices) == self.parallel_config.world_size, \
|
|
("VLLM_RAY_BUNDLE_INDICES must have the same size"
|
|
f" as the world size, but got {bundle_indices=} "
|
|
f"and {self.parallel_config.world_size=}")
|
|
assert len(set(bundle_indices)) == len(bundle_indices), \
|
|
("VLLM_RAY_BUNDLE_INDICES cannot have duplicate values,"
|
|
f" but got {bundle_indices=}")
|
|
else:
|
|
# use the first N bundles that have GPU resources.
|
|
bundle_indices = []
|
|
for bundle_id, bundle in enumerate(placement_group.bundle_specs):
|
|
if bundle.get(current_platform.ray_device_key, 0):
|
|
bundle_indices.append(bundle_id)
|
|
bundle_indices = bundle_indices[:self.parallel_config.world_size]
|
|
|
|
worker_metadata: List[RayWorkerMetaData] = []
|
|
driver_ip = get_ip()
|
|
for rank, bundle_id in enumerate(bundle_indices):
|
|
scheduling_strategy = PlacementGroupSchedulingStrategy(
|
|
placement_group=placement_group,
|
|
placement_group_capture_child_tasks=True,
|
|
placement_group_bundle_index=bundle_id,
|
|
)
|
|
|
|
if current_platform.ray_device_key == "GPU":
|
|
# NV+AMD GPUs, and Intel XPUs
|
|
worker = ray.remote(
|
|
num_cpus=0,
|
|
num_gpus=num_gpus,
|
|
scheduling_strategy=scheduling_strategy,
|
|
**ray_remote_kwargs,
|
|
)(RayWorkerWrapper).remote(vllm_config=self.vllm_config,
|
|
rpc_rank=rank)
|
|
else:
|
|
worker = ray.remote(
|
|
num_cpus=0,
|
|
num_gpus=0,
|
|
resources={current_platform.ray_device_key: num_gpus},
|
|
scheduling_strategy=scheduling_strategy,
|
|
**ray_remote_kwargs,
|
|
)(RayWorkerWrapper).remote(vllm_config=self.vllm_config,
|
|
rpc_rank=rank)
|
|
worker_metadata.append(
|
|
RayWorkerMetaData(worker=worker, created_rank=rank))
|
|
|
|
worker_ips = ray.get([
|
|
each.worker.get_node_ip.remote() # type: ignore[attr-defined]
|
|
for each in worker_metadata
|
|
])
|
|
|
|
for each, ip in zip(worker_metadata, worker_ips):
|
|
each.ip = ip
|
|
|
|
if not self.use_ray_spmd_worker:
|
|
for i, each in enumerate(worker_metadata):
|
|
# find and remove the dummy worker from the list
|
|
worker = each.worker
|
|
worker_ip = each.ip
|
|
if self.driver_dummy_worker is None and worker_ip == driver_ip:
|
|
# If the worker is on the same node as the driver, we use it
|
|
# as the resource holder for the driver process.
|
|
self.driver_dummy_worker = worker
|
|
self.driver_worker = RayWorkerWrapper(
|
|
vllm_config=self.vllm_config, rpc_rank=0)
|
|
worker_metadata.pop(i)
|
|
break
|
|
|
|
logger.debug("workers: %s", worker_metadata)
|
|
logger.debug("driver_dummy_worker: %s", self.driver_dummy_worker)
|
|
if not self.use_ray_spmd_worker and self.driver_dummy_worker is None:
|
|
raise ValueError(
|
|
"Ray does not allocate any GPUs on the driver node."
|
|
f"Driver IP: {driver_ip}, worker IPs: {worker_ips}."
|
|
"Consider adjusting the Ray placement group or running "
|
|
"the driver on a GPU node.")
|
|
|
|
ip_counts: Dict[str, int] = {}
|
|
for ip in worker_ips:
|
|
ip_counts[ip] = ip_counts.get(ip, 0) + 1
|
|
|
|
def sort_by_driver_then_worker_ip(item: RayWorkerMetaData):
|
|
"""
|
|
Sort the workers based on 3 properties:
|
|
1. If the worker is on the same node as the driver (vllm engine),
|
|
it should be placed first.
|
|
2. Then, if the worker is on a node with fewer workers, it should
|
|
be placed first.
|
|
3. Finally, if the work is on a node with smaller IP address, it
|
|
should be placed first.
|
|
"""
|
|
ip = item.ip
|
|
return (0 if ip == driver_ip else 1, ip_counts[ip], ip)
|
|
|
|
# After sorting, the workers on the same node will be
|
|
# close to each other, and the workers on the driver
|
|
# node will be placed first.
|
|
sorted_worker_metadata = sorted(worker_metadata,
|
|
key=sort_by_driver_then_worker_ip)
|
|
start_rank = 0 if self.use_ray_spmd_worker else 1
|
|
for i, item in enumerate(sorted_worker_metadata):
|
|
item.adjusted_rank = i + start_rank
|
|
self.workers = [item.worker for item in sorted_worker_metadata]
|
|
rerank_mapping = {
|
|
item.created_rank: item.adjusted_rank
|
|
for item in sorted_worker_metadata
|
|
}
|
|
self._run_workers("adjust_rank", rerank_mapping)
|
|
|
|
# Get the set of GPU IDs used on each node.
|
|
worker_node_and_gpu_ids = []
|
|
for worker in [self.driver_dummy_worker] + self.workers:
|
|
if worker is None:
|
|
# driver_dummy_worker can be None when using ray spmd worker.
|
|
continue
|
|
worker_node_and_gpu_ids.append(
|
|
ray.get(worker.get_node_and_gpu_ids.remote()) \
|
|
) # type: ignore
|
|
|
|
node_workers = defaultdict(list) # node id -> list of worker ranks
|
|
node_gpus = defaultdict(list) # node id -> list of gpu ids
|
|
|
|
for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids):
|
|
node_workers[node_id].append(i)
|
|
# `gpu_ids` can be a list of strings or integers.
|
|
# convert them to integers for consistency.
|
|
# NOTE: gpu_ids can be larger than 9 (e.g. 16 GPUs),
|
|
# string sorting is not sufficient.
|
|
# see https://github.com/vllm-project/vllm/issues/5590
|
|
gpu_ids = [int(x) for x in gpu_ids]
|
|
node_gpus[node_id].extend(gpu_ids)
|
|
for node_id, gpu_ids in node_gpus.items():
|
|
node_gpus[node_id] = sorted(gpu_ids)
|
|
|
|
all_ips = set(worker_ips + [driver_ip])
|
|
n_ips = len(all_ips)
|
|
n_nodes = len(node_workers)
|
|
|
|
if n_nodes != n_ips:
|
|
raise RuntimeError(
|
|
f"Every node should have a unique IP address. Got {n_nodes}"
|
|
f" nodes with node ids {list(node_workers.keys())} and "
|
|
f"{n_ips} unique IP addresses {all_ips}. Please check your"
|
|
" network configuration. If you set `VLLM_HOST_IP`"
|
|
" environment variable, make sure it is unique for"
|
|
" each node.")
|
|
|
|
# Set environment variables for the driver and workers.
|
|
all_args_to_update_environment_variables = [{
|
|
current_platform.device_control_env_var:
|
|
",".join(map(str, node_gpus[node_id])),
|
|
} for (node_id, _) in worker_node_and_gpu_ids]
|
|
|
|
# Environment variables to copy from driver to workers
|
|
env_vars_to_copy = get_env_vars_to_copy(
|
|
exclude_vars=self.WORKER_SPECIFIC_ENV_VARS,
|
|
additional_vars=set(current_platform.additional_env_vars).union(
|
|
self.ADDITIONAL_ENV_VARS),
|
|
destination="workers")
|
|
|
|
# Copy existing env vars to each worker's args
|
|
for args in all_args_to_update_environment_variables:
|
|
# TODO: refactor platform-specific env vars
|
|
for name in env_vars_to_copy:
|
|
if name in os.environ:
|
|
args[name] = os.environ[name]
|
|
|
|
self._env_vars_for_all_workers = (all_args_to_update_environment_variables)
|
|
|
|
self._run_workers("update_environment_variables",
|
|
self._get_env_vars_to_be_updated())
|
|
|
|
if len(node_gpus) == 1:
|
|
# in single node case, we don't need to get the IP address.
|
|
# the loopback address is sufficient
|
|
# NOTE: a node may have several IP addresses, one for each
|
|
# network interface. `get_ip()` might return any of them,
|
|
# while they might not work for communication inside the node
|
|
# if the network setup is complicated. Using the loopback address
|
|
# solves this issue, as it always works for communication inside
|
|
# the node.
|
|
driver_ip = "127.0.0.1"
|
|
distributed_init_method = get_distributed_init_method(
|
|
driver_ip, get_open_port())
|
|
|
|
# Initialize the actual workers inside worker wrapper.
|
|
all_kwargs = []
|
|
for rank, (node_id, _) in enumerate(worker_node_and_gpu_ids):
|
|
local_rank = node_workers[node_id].index(rank)
|
|
if envs_br.VLLM_BR_ENABLE_TP_GROUPS_IN_SUPERNODE:
|
|
kwargs = dict(
|
|
vllm_config=self.vllm_config,
|
|
local_rank=local_rank,
|
|
rank=rank,
|
|
distributed_init_method=distributed_init_method,
|
|
is_driver_worker=(not self.parallel_config)
|
|
or (rank in tp_driver_rank),
|
|
)
|
|
else:
|
|
kwargs = dict(
|
|
vllm_config=self.vllm_config,
|
|
local_rank=local_rank,
|
|
rank=rank,
|
|
distributed_init_method=distributed_init_method,
|
|
is_driver_worker=(not self.parallel_config)
|
|
or (rank % self.parallel_config.tensor_parallel_size == 0),
|
|
)
|
|
all_kwargs.append(kwargs)
|
|
self._run_workers("init_worker", all_kwargs)
|
|
|
|
self._run_workers("init_device")
|
|
self._run_workers("load_model",
|
|
max_concurrent_workers=self.parallel_config.
|
|
max_parallel_loading_workers)
|
|
|
|
if self.use_ray_spmd_worker:
|
|
if envs_br.VLLM_BR_ENABLE_TP_GROUPS_IN_SUPERNODE:
|
|
for pp_rank in range(self.parallel_config.pipeline_parallel_size):
|
|
self.pp_tp_workers.append([])
|
|
for tp_rank in range(
|
|
self.parallel_config.tensor_parallel_size):
|
|
# PP=8, TP=2
|
|
# pp_tp_workers = [[0, 1, 2, 3, 8, 9, 10, 11], [4, 5, 6, 7, 12, 13, 14, 15]]
|
|
rank = rank_map[(pp_rank, tp_rank)]
|
|
self.pp_tp_workers[pp_rank].append(self.workers[rank])
|
|
else:
|
|
for pp_rank in range(self.parallel_config.pipeline_parallel_size):
|
|
self.pp_tp_workers.append([])
|
|
for tp_rank in range(
|
|
self.parallel_config.tensor_parallel_size):
|
|
# PP=2, TP=4
|
|
# pp_tp_workers = [[0, 1, 2, 3], [4, 5, 6, 7]]
|
|
rank = (pp_rank * self.parallel_config.tensor_parallel_size
|
|
) + tp_rank
|
|
assert len(self.pp_tp_workers[pp_rank]) == tp_rank
|
|
assert pp_rank < len(self.pp_tp_workers)
|
|
self.pp_tp_workers[pp_rank].append(self.workers[rank])
|
|
|
|
# This is the list of workers that are rank 0 of each TP group EXCEPT
|
|
# global rank 0. These are the workers that will broadcast to the
|
|
# rest of the workers.
|
|
self.tp_driver_workers = []
|
|
# This is the list of workers that are not drivers and not the first
|
|
# worker in a TP group. These are the workers that will be
|
|
# broadcasted to.
|
|
self.non_driver_workers = []
|
|
|
|
# Enforce rank order for correct rank to return final output.
|
|
for index, worker in enumerate(self.workers):
|
|
# The driver worker is rank 0 and not in self.workers.
|
|
rank = index + 1
|
|
if envs_br.VLLM_BR_ENABLE_TP_GROUPS_IN_SUPERNODE:
|
|
if rank in tp_driver_rank:
|
|
self.tp_driver_workers.append(worker)
|
|
else:
|
|
self.non_driver_workers.append(worker)
|
|
else:
|
|
if rank % self.parallel_config.tensor_parallel_size == 0:
|
|
self.tp_driver_workers.append(worker)
|
|
else:
|
|
self.non_driver_workers.append(worker)
|
|
|
|
|
|
RayDistributedExecutor._init_workers_ray = _init_workers_ray_br # noqa: E501
|