160 lines
7.1 KiB
Python
160 lines
7.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import gc
|
|
import os
|
|
|
|
import torch
|
|
from vllm.logger import init_logger
|
|
from vllm.platforms import current_platform
|
|
from vllm.tracing import instrument
|
|
from vllm.utils.mem_utils import MemorySnapshot, format_gib
|
|
from vllm.utils.torch_utils import set_random_seed
|
|
from vllm.v1.utils import report_usage_stats
|
|
from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
|
|
from vllm.v1.worker.workspace import init_workspace_manager
|
|
|
|
from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask
|
|
from vllm_omni.platforms import current_omni_platform
|
|
from vllm_omni.worker.base import OmniGPUWorkerBase
|
|
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
|
|
from vllm_omni.worker.memory_utils import request_memory_tolerant
|
|
from vllm_omni.worker.mixins import OmniWorkerMixin
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
|
|
"""GPU worker for autoregressive omni model stages.
|
|
|
|
Extends the base GPUWorker to initialize and manage autoregressive
|
|
model runners for text generation stages (e.g., thinker stages).
|
|
"""
|
|
|
|
model_runner_cls = GPUARModelRunner
|
|
|
|
@instrument(span_name="Init device")
|
|
def init_device(self):
|
|
if self.device_config.device_type in ("cuda", "musa"):
|
|
# This env var set by Ray causes exceptions with graph building.
|
|
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
|
|
parallel_config = self.parallel_config
|
|
if (
|
|
parallel_config.distributed_executor_backend not in ("ray", "external_launcher")
|
|
and parallel_config.data_parallel_backend != "ray"
|
|
and parallel_config.nnodes_within_dp == 1
|
|
):
|
|
# Use local DP rank if available, otherwise use global DP rank.
|
|
dp_local_rank = self.parallel_config.data_parallel_rank_local
|
|
if dp_local_rank is None:
|
|
dp_local_rank = self.parallel_config.data_parallel_index
|
|
|
|
tp_pp_world_size = (
|
|
self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size
|
|
)
|
|
|
|
# DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
|
|
self.local_rank += dp_local_rank * tp_pp_world_size
|
|
|
|
# Publish the logical-to-physical mapping for topology queries
|
|
# such as NIC affinity and P2P checks (upstream PR #45026).
|
|
assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids
|
|
if assigned_physical_gpu_ids is not None:
|
|
from vllm.platforms.interface import set_assigned_physical_gpu_ids
|
|
|
|
set_assigned_physical_gpu_ids(assigned_physical_gpu_ids)
|
|
assert self.local_rank < len(assigned_physical_gpu_ids), (
|
|
f"local_rank {self.local_rank} is out of bounds for "
|
|
f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}"
|
|
)
|
|
if parallel_config.distributed_executor_backend not in ("ray", "external_launcher"):
|
|
assert self.parallel_config.local_world_size <= len(assigned_physical_gpu_ids), (
|
|
f"local_world_size ({self.parallel_config.local_world_size}) "
|
|
"exceeds assigned_physical_gpu_ids count "
|
|
f"({len(assigned_physical_gpu_ids)})"
|
|
)
|
|
else:
|
|
assert self.local_rank < torch.accelerator.device_count(), (
|
|
f"DP adjusted local rank {self.local_rank} is out of "
|
|
f"bounds for {torch.accelerator.device_count()} devices."
|
|
)
|
|
visible_device_index = current_platform.logical_device_id_to_visible_device_id(self.local_rank)
|
|
self.device = current_omni_platform.get_torch_device(visible_device_index)
|
|
torch.accelerator.set_device_index(self.device)
|
|
|
|
current_platform.check_if_supports_dtype(self.model_config.dtype)
|
|
|
|
# Initialize the distributed environment BEFORE taking
|
|
# memory snapshot
|
|
# This ensures NCCL buffers are allocated before we measure
|
|
# available memory
|
|
init_worker_distributed_environment(
|
|
self.vllm_config,
|
|
self.rank,
|
|
self.distributed_init_method,
|
|
self.local_rank,
|
|
current_platform.dist_backend,
|
|
)
|
|
|
|
# Set random seed.
|
|
set_random_seed(self.model_config.seed)
|
|
|
|
# Now take memory snapshot after NCCL is initialized
|
|
gc.collect()
|
|
torch.accelerator.empty_cache()
|
|
|
|
# take current memory snapshot
|
|
self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
|
|
self.requested_memory = request_memory_tolerant(init_snapshot, self.cache_config)
|
|
logger.debug("worker init memory snapshot: %r", self.init_snapshot)
|
|
logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory))
|
|
else:
|
|
raise RuntimeError(f"Not support device type: {self.device_config.device}")
|
|
|
|
# Initialize workspace manager
|
|
num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
|
|
init_workspace_manager(self.device, num_ubatches)
|
|
|
|
model_stage = getattr(self.model_config, "model_stage", None)
|
|
stage_id = getattr(self.model_config, "stage_id", None)
|
|
|
|
self.use_v2_model_runner = bool(getattr(self.model_config, "use_v2_model_runner", False))
|
|
if self.use_v2_model_runner:
|
|
from vllm_omni.worker_v2.omni_ar_model_runner import (
|
|
OmniARModelRunner,
|
|
)
|
|
|
|
logger.info(
|
|
"Using MR v2 OmniARModelRunner for omni AR stage (stage_id=%s model_stage=%s).",
|
|
stage_id,
|
|
model_stage,
|
|
)
|
|
self.model_runner = OmniARModelRunner(self.vllm_config, self.device)
|
|
else:
|
|
self.use_v2_model_runner = False
|
|
self.model_runner = self.model_runner_cls(self.vllm_config, self.device)
|
|
|
|
if self.rank == 0:
|
|
# If usage stat is enabled, collect relevant info.
|
|
report_usage_stats(self.vllm_config)
|
|
|
|
def handle_sleep_task(self, task: OmniSleepTask | dict) -> OmniACK:
|
|
"""
|
|
Explicitly handle sleep commands.
|
|
Calls the implementation in the base class OmniGPUWorkerBase.
|
|
"""
|
|
logger.debug(f"[AR Worker {self.rank}] Resolving handle_sleep_task dispatch")
|
|
if isinstance(task, dict):
|
|
task = OmniSleepTask(**task)
|
|
return super().handle_sleep_task(task)
|
|
|
|
def handle_wake_task(self, task: OmniWakeTask | dict) -> OmniACK:
|
|
"""
|
|
Explicitly handle wake-up commands.
|
|
Calls the implementation in the base class OmniGPUWorkerBase.
|
|
"""
|
|
logger.debug(f"[AR Worker {self.rank}] Resolving handle_wake_task dispatch")
|
|
if isinstance(task, dict):
|
|
task = OmniWakeTask(**task)
|
|
return super().handle_wake_task(task)
|