tzhouam
ea16546ae9
- Renamed `DiffusionModelRunner` to `GPUDiffusionModelRunner` for better clarity in class purpose. - Updated `DiffusionGPUWorker` to utilize the renamed `GPUDiffusionModelRunner`. - Simplified output handling in `GPUDiffusionModelRunner` by removing unnecessary checks, enhancing code readability.
67 lines
3.1 KiB
Python
67 lines
3.1 KiB
Python
from vllm.v1.worker.gpu_worker import Worker as GPUWorker
|
|
import torch
|
|
import os
|
|
import gc
|
|
from vllm.utils import GiB_bytes, MemorySnapshot
|
|
from vllm.platforms import current_platform
|
|
from vllm.worker.worker import _check_if_gpu_supports_dtype
|
|
from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
|
|
from vllm.model_executor import set_random_seed
|
|
from vllm.v1.utils import report_usage_stats
|
|
|
|
from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
|
|
|
|
|
|
class GPUDiffusionWorker(GPUWorker):
|
|
def init_device(self):
|
|
if self.device_config.device.type == "cuda":
|
|
# torch.distributed.all_reduce does not free the input tensor until
|
|
# the synchronization point. This causes the memory usage to grow
|
|
# as the number of all_reduce calls increases. This env var disables
|
|
# this behavior.
|
|
# Related issue:
|
|
# https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
|
|
os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
|
|
|
|
# This env var set by Ray causes exceptions with graph building.
|
|
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
|
|
self.device = torch.device(f"cuda:{self.local_rank}")
|
|
current_platform.set_device(self.device)
|
|
|
|
_check_if_gpu_supports_dtype(self.model_config.dtype)
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
# take current memory snapshot
|
|
self.init_snapshot = MemorySnapshot()
|
|
self.requested_memory = (self.init_snapshot.total_memory *
|
|
self.cache_config.gpu_memory_utilization)
|
|
if self.init_snapshot.free_memory < self.requested_memory:
|
|
GiB = lambda b: round(b / GiB_bytes, 2)
|
|
raise ValueError(
|
|
f"Free memory on device "
|
|
f"({GiB(self.init_snapshot.free_memory)}/"
|
|
f"{GiB(self.init_snapshot.total_memory)} GiB) on startup "
|
|
f"is less than desired GPU memory utilization "
|
|
f"({self.cache_config.gpu_memory_utilization}, "
|
|
f"{GiB(self.requested_memory)} GiB). Decrease GPU memory "
|
|
f"utilization or reduce GPU memory used by other processes."
|
|
)
|
|
else:
|
|
raise RuntimeError(
|
|
f"Not support device type: {self.device_config.device}")
|
|
# Initialize the distributed environment.
|
|
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)
|
|
|
|
# Construct the model runner
|
|
self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner(
|
|
self.vllm_config, self.device)
|
|
|
|
if self.rank == 0:
|
|
# If usage stat is enabled, collect relevant info.
|
|
report_usage_stats(self.vllm_config) |