127 lines
4.8 KiB
Python
127 lines
4.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import multiprocessing as mp
|
|
import os
|
|
import time
|
|
|
|
import torch
|
|
from vllm.config import LoadConfig, VllmConfig, set_current_vllm_config
|
|
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 (
|
|
OmniDiffusionConfig,
|
|
set_current_omni_diffusion_config,
|
|
)
|
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
|
init_distributed_environment,
|
|
initialize_model_parallel,
|
|
)
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.worker.gpu_worker import GPUWorker, WorkerProc
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class NPUWorker(GPUWorker):
|
|
"""
|
|
A worker that executes the model on a single NPU.
|
|
Inherits from GPUWorker and overrides device-specific initialization.
|
|
"""
|
|
|
|
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_current_omni_diffusion_config(self.od_config),
|
|
set_current_vllm_config(vllm_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)
|
|
|
|
|
|
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.")
|