252 lines
9.6 KiB
Python
252 lines
9.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import multiprocessing as mp
|
|
import os
|
|
import time
|
|
from collections.abc import Iterable
|
|
from contextlib import AbstractContextManager, nullcontext
|
|
|
|
import torch
|
|
from vllm.config import LoadConfig, VllmConfig
|
|
from vllm.logger import init_logger
|
|
from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
|
|
|
|
from vllm_omni.diffusion.cache.selector import get_cache_backend
|
|
from vllm_omni.diffusion.data import (
|
|
DiffusionOutput,
|
|
OmniDiffusionConfig,
|
|
)
|
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
|
destroy_distributed_env,
|
|
init_distributed_environment,
|
|
initialize_model_parallel,
|
|
)
|
|
from vllm_omni.diffusion.forward_context import set_forward_context
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.worker.gpu_diffusion_worker import WorkerProc
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class NPUWorker:
|
|
"""
|
|
A worker that executes the model on a single NPU.
|
|
Inherits from GPUWorker and overrides device-specific initialization.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
local_rank: int,
|
|
rank: int,
|
|
od_config: OmniDiffusionConfig,
|
|
):
|
|
self.local_rank = local_rank
|
|
self.rank = rank
|
|
self.od_config = od_config
|
|
self.pipeline = None
|
|
self.device = None
|
|
self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
|
|
self.init_device_and_model()
|
|
|
|
def init_device_and_model(self) -> None:
|
|
"""Initialize the NPU device and load the model."""
|
|
world_size = self.od_config.num_gpus
|
|
rank = self.rank
|
|
# Set environment variables for distributed initialization
|
|
os.environ["MASTER_ADDR"] = "localhost"
|
|
os.environ["MASTER_PORT"] = str(self.od_config.master_port)
|
|
os.environ["LOCAL_RANK"] = str(self.local_rank)
|
|
os.environ["RANK"] = str(rank)
|
|
os.environ["WORLD_SIZE"] = str(world_size)
|
|
|
|
device = torch.device(f"npu:{rank}")
|
|
torch.npu.set_device(device)
|
|
|
|
# hack
|
|
vllm_config = VllmConfig()
|
|
vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size
|
|
vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
|
|
self.vllm_config = vllm_config
|
|
with set_forward_context(vllm_config=vllm_config, omni_diffusion_config=self.od_config):
|
|
init_distributed_environment(world_size=world_size, rank=rank)
|
|
logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
|
|
parallel_config = self.od_config.parallel_config
|
|
initialize_model_parallel(
|
|
data_parallel_size=parallel_config.data_parallel_size,
|
|
cfg_parallel_size=parallel_config.cfg_parallel_size,
|
|
sequence_parallel_size=parallel_config.sequence_parallel_size,
|
|
ulysses_degree=parallel_config.ulysses_degree,
|
|
ring_degree=parallel_config.ring_degree,
|
|
tensor_parallel_size=parallel_config.tensor_parallel_size,
|
|
pipeline_parallel_size=parallel_config.pipeline_parallel_size,
|
|
)
|
|
|
|
load_config = LoadConfig()
|
|
model_loader = DiffusersPipelineLoader(load_config)
|
|
time_before_load = time.perf_counter()
|
|
with DeviceMemoryProfiler() as m:
|
|
self.pipeline = model_loader.load_model(
|
|
od_config=self.od_config,
|
|
load_device=f"npu:{rank}",
|
|
)
|
|
time_after_load = time.perf_counter()
|
|
|
|
logger.info(
|
|
"Model loading took %.4f GiB and %.6f seconds",
|
|
m.consumed_memory / GiB_bytes,
|
|
time_after_load - time_before_load,
|
|
)
|
|
logger.info(f"Worker {self.rank}: Model loaded successfully.")
|
|
|
|
# Setup cache backend based on type (both backends use enable()/reset() interface)
|
|
self.cache_backend = get_cache_backend(self.od_config.cache_backend, self.od_config.cache_config)
|
|
|
|
if self.cache_backend is not None:
|
|
self.cache_backend.enable(self.pipeline)
|
|
|
|
def generate(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
|
|
"""
|
|
Generate output for the given requests.
|
|
|
|
Args:
|
|
requests: List of diffusion requests
|
|
|
|
Returns:
|
|
DiffusionOutput with generated results
|
|
"""
|
|
return self.execute_model(requests, self.od_config)
|
|
|
|
@torch.inference_mode()
|
|
def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
|
|
"""
|
|
Execute a forward pass.
|
|
"""
|
|
assert self.pipeline is not None
|
|
if not reqs or len(reqs) == 0:
|
|
raise ValueError("Cannot execute model with empty request list")
|
|
# TODO: dealing with first req for now
|
|
req = reqs[0]
|
|
|
|
if req.generator is None and req.seed is not None:
|
|
req.generator = torch.Generator(device=self.device).manual_seed(req.seed)
|
|
|
|
# Refresh cache context if needed
|
|
if self.cache_backend is not None and self.cache_backend.is_enabled():
|
|
self.cache_backend.refresh(self.pipeline, req.num_inference_steps)
|
|
with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
|
|
output = self.pipeline.forward(req)
|
|
return output
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
return self.pipeline.load_weights(weights)
|
|
|
|
def sleep(self, level: int = 1) -> bool:
|
|
"""
|
|
Put the worker to sleep. The worker should not process any requests.
|
|
The caller should guarantee that no requests are being processed
|
|
during the sleep period, before `wake_up` is called.
|
|
|
|
Args:
|
|
level: The sleep level. Level 1 sleep will offload the model
|
|
weights and discard the kv cache.
|
|
Currently only support level 1.
|
|
"""
|
|
from vllm.device_allocator.cumem import CuMemAllocator
|
|
|
|
free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
|
|
|
|
# Save the buffers before level 2 sleep
|
|
if level == 2:
|
|
model = self.pipeline
|
|
self._sleep_saved_buffers = {name: buffer.cpu().clone() for name, buffer in model.named_buffers()}
|
|
|
|
allocator = CuMemAllocator.get_instance()
|
|
allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
|
|
free_bytes_after_sleep, total = torch.cuda.mem_get_info()
|
|
freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
|
|
used_bytes = total - free_bytes_after_sleep
|
|
assert freed_bytes >= 0, "Memory usage increased after sleeping."
|
|
logger.info(
|
|
"Sleep mode freed %.2f GiB memory, %.2f GiB memory is still in use.",
|
|
freed_bytes / GiB_bytes,
|
|
used_bytes / GiB_bytes,
|
|
)
|
|
return True
|
|
|
|
def wake_up(self, tags: list[str] | None = None) -> bool:
|
|
"""
|
|
Wake up the worker from sleep mode. See the sleep function
|
|
method for more details.
|
|
|
|
Args:
|
|
tags: An optional list of tags to reallocate the worker memory
|
|
for specific memory allocations. Values must be in
|
|
`("weights")`. If None, all memory is reallocated.
|
|
wake_up should be called with all tags (or None) before the
|
|
worker is used again.
|
|
"""
|
|
from vllm.device_allocator.cumem import CuMemAllocator
|
|
|
|
allocator = CuMemAllocator.get_instance()
|
|
allocator.wake_up(tags)
|
|
|
|
# Restore the buffers after level 2 sleep
|
|
if len(self._sleep_saved_buffers):
|
|
model = self.pipeline
|
|
for name, buffer in model.named_buffers():
|
|
if name in self._sleep_saved_buffers:
|
|
buffer.data.copy_(self._sleep_saved_buffers[name].data)
|
|
self._sleep_saved_buffers = {}
|
|
return True
|
|
|
|
def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
|
|
if self.od_config.enable_sleep_mode:
|
|
from vllm.device_allocator.cumem import CuMemAllocator
|
|
|
|
allocator = CuMemAllocator.get_instance()
|
|
if tag == "weights":
|
|
assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process."
|
|
return allocator.use_memory_pool(tag=tag)
|
|
else:
|
|
return nullcontext()
|
|
|
|
def shutdown(self) -> None:
|
|
destroy_distributed_env()
|
|
|
|
|
|
class NPUWorkerProc(WorkerProc):
|
|
"""Wrapper that runs one NPUWorker in a separate process."""
|
|
|
|
def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> NPUWorker:
|
|
"""Create an NPUWorker instead of GPUWorker."""
|
|
return NPUWorker(
|
|
local_rank=gpu_id,
|
|
rank=gpu_id,
|
|
od_config=od_config,
|
|
)
|
|
|
|
@staticmethod
|
|
def worker_main(
|
|
rank: int,
|
|
od_config: OmniDiffusionConfig,
|
|
pipe_writer: mp.connection.Connection,
|
|
broadcast_handle,
|
|
) -> None:
|
|
"""Worker initialization and execution loops."""
|
|
|
|
worker_proc = NPUWorkerProc(
|
|
od_config,
|
|
gpu_id=rank,
|
|
broadcast_handle=broadcast_handle,
|
|
)
|
|
logger.info(f"Worker {rank}: Scheduler loop started.")
|
|
pipe_writer.send(
|
|
{
|
|
"status": "ready",
|
|
"result_handle": worker_proc.result_mq_handle if rank == 0 else None,
|
|
}
|
|
)
|
|
worker_proc.worker_busy_loop()
|
|
logger.info(f"Worker {rank}: Shutdown complete.")
|