FoundationVision/LlamaGen
64
1"""A GPU worker class."""2import gc3import os4from typing import Any, Dict, List, Optional, Set, Tuple5 6import torch7import torch.distributed8 9from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, LoRAConfig,10 ModelConfig, ParallelConfig, SchedulerConfig,11 VisionLanguageConfig)12from vllm.distributed import (broadcast_tensor_dict,13 ensure_model_parallel_initialized,14 init_distributed_environment)15from vllm.distributed.device_communicators import pynccl_utils16from vllm.distributed.device_communicators.custom_all_reduce import (17 init_custom_ar)18from vllm.lora.request import LoRARequest19from vllm.model_executor import set_random_seed20from vllm.sequence import SamplerOutput, SequenceGroupMetadata21from vllm.worker.cache_engine import CacheEngine22# from vllm.worker.model_runner import ModelRunner23from vllm.worker.worker_base import WorkerBase24from serve.model_runner import ModelRunner25 26 27class Worker(WorkerBase):28 """A worker class that executes (a partition of) the model on a GPU.29 30 Each worker is associated with a single GPU. The worker is responsible for31 maintaining the KV cache and executing the model on the GPU. In case of32 distributed inference, each worker is assigned a partition of the model.33 """34 35 def __init__(36 self,37 model_config: ModelConfig,38 parallel_config: ParallelConfig,39 scheduler_config: SchedulerConfig,40 device_config: DeviceConfig,41 cache_config: CacheConfig,42 load_config: LoadConfig,43 local_rank: int,44 rank: int,45 distributed_init_method: str,46 lora_config: Optional[LoRAConfig] = None,47 vision_language_config: Optional[VisionLanguageConfig] = None,48 is_driver_worker: bool = False,49 ) -> None:50 self.model_config = model_config51 self.parallel_config = parallel_config52 self.scheduler_config = scheduler_config53 self.device_config = device_config54 self.cache_config = cache_config55 self.local_rank = local_rank56 self.rank = rank57 self.distributed_init_method = distributed_init_method58 self.lora_config = lora_config59 self.load_config = load_config60 self.is_driver_worker = is_driver_worker61 if self.is_driver_worker:62 assert self.rank == 0, "The driver worker must have rank 0."63 64 if self.model_config.trust_remote_code:65 # note: lazy import to avoid importing torch before initializing66 from vllm.utils import init_cached_hf_modules67 init_cached_hf_modules()68 self.vision_language_config = vision_language_config69 if self.vision_language_config:70 assert not self.lora_config, (71 "To be tested: vision language model with LoRA settings.")72 73 self.model_runner = ModelRunner(74 model_config,75 parallel_config,76 scheduler_config,77 device_config,78 load_config=load_config,79 lora_config=self.lora_config,80 kv_cache_dtype=self.cache_config.cache_dtype,81 is_driver_worker=is_driver_worker,82 vision_language_config=vision_language_config,83 )84 # Uninitialized cache engine. Will be initialized by85 # initialize_cache.86 self.cache_engine: CacheEngine87 self.gpu_cache: List[torch.Tensor]88 89 def init_device(self) -> None:90 if self.device_config.device.type == "cuda":91 # torch.distributed.all_reduce does not free the input tensor until92 # the synchronization point. This causes the memory usage to grow93 # as the number of all_reduce calls increases. This env var disables94 # this behavior.95 # Related issue:96 # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/19157397 os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"98 99 # This env var set by Ray causes exceptions with graph building.100 os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)101 self.device = torch.device(f"cuda:{self.local_rank}")102 torch.cuda.set_device(self.device)103 104 _check_if_gpu_supports_dtype(self.model_config.dtype)105 torch.cuda.empty_cache()106 self.init_gpu_memory = torch.cuda.mem_get_info()[0]107 else:108 raise RuntimeError(109 f"Not support device type: {self.device_config.device}")110 # Initialize the distributed environment.111 init_worker_distributed_environment(self.parallel_config, self.rank,112 self.distributed_init_method,113 self.local_rank)114 # Set random seed.115 set_random_seed(self.model_config.seed)116 117 def load_model(self, args):118 self.model_runner.load_model(args)119 120 @torch.inference_mode()121 def determine_num_available_blocks(self) -> Tuple[int, int]:122 """Profiles the peak memory usage of the model to determine how many123 KV blocks may be allocated without OOMs.124 125 The engine will first conduct a profiling of the existing memory usage.126 Then, it calculate the maximum possible number of GPU and CPU blocks127 that can be allocated with the remaining free memory.128 129 .. tip::130 You may limit the usage of GPU memory131 by adjusting the `gpu_memory_utilization` parameter.132 """133 # Profile the memory usage of the model and get the maximum number of134 # cache blocks that can be allocated with the remaining free memory.135 torch.cuda.empty_cache()136 137 # Execute a forward pass with dummy inputs to profile the memory usage138 # of the model.139 self.model_runner.profile_run()140 141 # Calculate the number of blocks that can be allocated with the142 # profiled peak memory.143 torch.cuda.synchronize()144 free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info()145 # NOTE(woosuk): Here we assume that the other processes using the same146 # GPU did not change their memory usage during the profiling.147 peak_memory = self.init_gpu_memory - free_gpu_memory148 assert peak_memory > 0, (149 "Error in memory profiling. This happens when the GPU memory was "150 "not properly cleaned up before initializing the vLLM instance.")151 152 cache_block_size = self.get_cache_block_size_bytes()153 num_gpu_blocks = int(154 (total_gpu_memory * self.cache_config.gpu_memory_utilization -155 peak_memory) // cache_block_size)156 num_cpu_blocks = int(self.cache_config.swap_space_bytes //157 cache_block_size)158 num_gpu_blocks = max(num_gpu_blocks, 0)159 num_cpu_blocks = max(num_cpu_blocks, 0)160 if self.model_runner.lora_manager:161 self.model_runner.remove_all_loras()162 gc.collect()163 torch.cuda.empty_cache()164 return num_gpu_blocks, num_cpu_blocks165 166 def initialize_cache(self, num_gpu_blocks: int,167 num_cpu_blocks: int) -> None:168 """Allocate GPU and CPU KV cache with the specified number of blocks.169 170 This also warms up the model, which may record CUDA graphs.171 """172 raise_if_cache_size_invalid(num_gpu_blocks,173 self.cache_config.block_size,174 self.model_config.max_model_len)175 176 self.cache_config.num_gpu_blocks = num_gpu_blocks177 self.cache_config.num_cpu_blocks = num_cpu_blocks178 179 self._init_cache_engine()180 self._warm_up_model()181 182 def _init_cache_engine(self):183 assert self.cache_config.num_gpu_blocks is not None184 self.cache_engine = CacheEngine(self.cache_config, self.model_config,185 self.parallel_config)186 self.gpu_cache = self.cache_engine.gpu_cache187 self.model_runner.set_block_size(self.cache_engine.block_size)188 189 def _warm_up_model(self) -> None:190 if not self.model_config.enforce_eager:191 self.model_runner.capture_model(self.gpu_cache)192 # Reset the seed to ensure that the random state is not affected by193 # the model initialization and profiling.194 set_random_seed(self.model_config.seed)195 196 def cache_swap(197 self,198 blocks_to_swap_in: Dict[int, int],199 blocks_to_swap_out: Dict[int, int],200 blocks_to_copy: Dict[int, List[int]],201 ) -> None:202 # Issue cache operations.203 # TODO(woosuk): Profile swapping overhead and optimize if needed.204 if blocks_to_swap_in:205 self.cache_engine.swap_in(blocks_to_swap_in)206 if blocks_to_swap_out:207 self.cache_engine.swap_out(blocks_to_swap_out)208 if blocks_to_copy:209 self.cache_engine.copy(blocks_to_copy)210 211 @torch.inference_mode()212 def execute_model(213 self,214 seq_group_metadata_list: Optional[List[SequenceGroupMetadata]] = None,215 blocks_to_swap_in: Optional[Dict[int, int]] = None,216 blocks_to_swap_out: Optional[Dict[int, int]] = None,217 blocks_to_copy: Optional[Dict[int, List[int]]] = None,218 num_lookahead_slots: int = 0,219 ) -> List[SamplerOutput]:220 221 if self.is_driver_worker:222 assert seq_group_metadata_list is not None223 num_seq_groups = len(seq_group_metadata_list)224 assert blocks_to_swap_in is not None225 assert blocks_to_swap_out is not None226 assert blocks_to_copy is not None227 data: Dict[str, Any] = {228 "num_seq_groups": num_seq_groups,229 "blocks_to_swap_in": blocks_to_swap_in,230 "blocks_to_swap_out": blocks_to_swap_out,231 "blocks_to_copy": blocks_to_copy,232 }233 broadcast_tensor_dict(data, src=0)234 else:235 data = broadcast_tensor_dict(src=0)236 num_seq_groups = data["num_seq_groups"]237 blocks_to_swap_in = data["blocks_to_swap_in"]238 blocks_to_swap_out = data["blocks_to_swap_out"]239 blocks_to_copy = data["blocks_to_copy"]240 241 assert blocks_to_swap_in is not None242 assert blocks_to_swap_out is not None243 assert blocks_to_copy is not None244 self.cache_swap(blocks_to_swap_in, blocks_to_swap_out, blocks_to_copy)245 246 # If there is no input, we don't need to execute the model.247 if num_seq_groups == 0:248 return []249 250 output = self.model_runner.execute_model(seq_group_metadata_list,251 self.gpu_cache)252 253 # Worker only supports single-step execution. Wrap the output in a list254 # to conform to interface.255 return [output]256 257 def add_lora(self, lora_request: LoRARequest) -> bool:258 return self.model_runner.add_lora(lora_request)259 260 def remove_lora(self, lora_id: int) -> bool:261 return self.model_runner.remove_lora(lora_id)262 263 def list_loras(self) -> Set[int]:264 return self.model_runner.list_loras()265 266 @property267 def max_model_len(self) -> int:268 return self.model_config.max_model_len269 270 @property271 def vocab_size(self) -> int:272 return self.model_runner.vocab_size273 274 def get_cache_block_size_bytes(self) -> int:275 """Get the size of the KV cache block size in bytes.276 """277 return CacheEngine.get_cache_block_size(self.cache_config,278 self.model_config,279 self.parallel_config)280 281 282def init_worker_distributed_environment(283 parallel_config: ParallelConfig,284 rank: int,285 distributed_init_method: Optional[str] = None,286 local_rank: int = -1,287) -> None:288 """Initialize the distributed environment."""289 init_distributed_environment(parallel_config.world_size, rank,290 distributed_init_method, local_rank)291 292 if pynccl_utils.is_initialized():293 pynccl_world_size = pynccl_utils.get_world_size()294 if pynccl_world_size != parallel_config.world_size:295 raise RuntimeError(296 "pynccl is already initialized but the pynccl world "297 "size does not match parallel_config.world_size "298 f"({pynccl_world_size} vs. {parallel_config.world_size}).")299 elif parallel_config.world_size > 1:300 # NOTE(woosuk): We don't initialize pynccl process group when world size301 # is 1.302 pynccl_utils.init_process_group(303 world_size=parallel_config.world_size,304 local_rank=local_rank,305 rank=rank,306 init_method=distributed_init_method,307 )308 309 ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,310 parallel_config.pipeline_parallel_size)311 312 # Initialize a custom fast all-reduce implementation.313 if not parallel_config.disable_custom_all_reduce:314 init_custom_ar()315 316 # A small all_reduce for warmup.317 torch.distributed.all_reduce(torch.zeros(1).cuda())318 if pynccl_utils.is_initialized():319 pynccl_utils.all_reduce(torch.zeros(1).cuda())320 321 322def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):323 # Check if the GPU supports the dtype.324 if torch_dtype == torch.bfloat16:325 compute_capability = torch.cuda.get_device_capability()326 if compute_capability[0] < 8:327 gpu_name = torch.cuda.get_device_name()328 raise ValueError(329 "Bfloat16 is only supported on GPUs with compute capability "330 f"of at least 8.0. Your {gpu_name} GPU has compute capability "331 f"{compute_capability[0]}.{compute_capability[1]}. "332 "You can use float16 instead by explicitly setting the"333 "`dtype` flag in CLI, for example: --dtype=half.")334 335 336def raise_if_cache_size_invalid(num_gpu_blocks, block_size,337 max_model_len) -> None:338 if num_gpu_blocks <= 0:339 raise ValueError("No available memory for the cache blocks. "340 "Try increasing `gpu_memory_utilization` when "341 "initializing the engine.")342 max_seq_len = block_size * num_gpu_blocks343 if max_model_len > max_seq_len:344 raise ValueError(345 f"The model's max seq len ({max_model_len}) "346 "is larger than the maximum number of tokens that can be "347 f"stored in KV cache ({max_seq_len}). Try increasing "348 "`gpu_memory_utilization` or decreasing `max_model_len` when "349 "initializing the engine.")