FoundationVision/LlamaGen

1from typing import Dict, List, Set, Tuple, Optional, Set2import argparse3 4from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, LoRAConfig,5                         ModelConfig, ParallelConfig, SchedulerConfig,6                         SpeculativeConfig, VisionLanguageConfig)7from vllm.executor.executor_base import ExecutorAsyncBase, ExecutorBase8from vllm.logger import init_logger9from vllm.lora.request import LoRARequest10from vllm.sequence import SamplerOutput, SequenceGroupMetadata11from vllm.utils import (get_distributed_init_method, get_ip, get_open_port,12                        make_async)13 14logger = init_logger(__name__)15 16 17class GPUExecutor(ExecutorBase):18    def __init__(19        self,20        args: argparse.ArgumentParser,21        model_config: ModelConfig,22        cache_config: CacheConfig,23        parallel_config: ParallelConfig,24        scheduler_config: SchedulerConfig,25        device_config: DeviceConfig,26        load_config: LoadConfig,27        lora_config: Optional[LoRAConfig],28        vision_language_config: Optional[VisionLanguageConfig],29        speculative_config: Optional[SpeculativeConfig],30    ) -> None:31        self.args = args32        self.model_config = model_config33        self.cache_config = cache_config34        self.lora_config = lora_config35        self.load_config = load_config36        self.parallel_config = parallel_config37        self.scheduler_config = scheduler_config38        self.device_config = device_config39        self.vision_language_config = vision_language_config40        self.speculative_config = speculative_config41 42        self._init_executor()43 44    def _init_executor(self) -> None:45        """Initialize the worker and load the model.46 47        If speculative decoding is enabled, we instead create the speculative48        worker.49        """50        if self.speculative_config is None:51            self._init_non_spec_worker()52        else:53            self._init_spec_worker()54 55    def _init_non_spec_worker(self):56        # Lazy import the Worker to avoid importing torch.cuda/xformers57        # before CUDA_VISIBLE_DEVICES is set in the Worker58        # from vllm.worker.worker import Worker59        from serve.worker import Worker60 61        assert self.parallel_config.world_size == 1, (62            "GPUExecutor only supports single GPU.")63 64        distributed_init_method = get_distributed_init_method(65            get_ip(), get_open_port())66        self.driver_worker = Worker(67            model_config=self.model_config,68            parallel_config=self.parallel_config,69            scheduler_config=self.scheduler_config,70            device_config=self.device_config,71            cache_config=self.cache_config,72            load_config=self.load_config,73            local_rank=0,74            rank=0,75            distributed_init_method=distributed_init_method,76            lora_config=self.lora_config,77            vision_language_config=self.vision_language_config,78            is_driver_worker=True,79        )80        self.driver_worker.init_device()81        self.driver_worker.load_model(self.args)82 83    def _init_spec_worker(self):84        """Initialize a SpecDecodeWorker, using a draft model for proposals.85        """86        assert self.speculative_config is not None87 88        from vllm.spec_decode.multi_step_worker import MultiStepWorker89        from vllm.spec_decode.spec_decode_worker import SpecDecodeWorker90        from vllm.worker.worker import Worker91 92        distributed_init_method = get_distributed_init_method(93            get_ip(), get_open_port())94 95        target_worker = Worker(96            model_config=self.model_config,97            parallel_config=self.parallel_config,98            scheduler_config=self.scheduler_config,99            device_config=self.device_config,100            cache_config=self.cache_config,101            load_config=self.load_config,102            local_rank=0,103            rank=0,104            distributed_init_method=distributed_init_method,105            lora_config=self.lora_config,106            vision_language_config=self.vision_language_config,107            is_driver_worker=True,108        )109 110        draft_worker = MultiStepWorker(111            model_config=self.speculative_config.draft_model_config,112            parallel_config=self.speculative_config.draft_parallel_config,113            scheduler_config=self.scheduler_config,114            device_config=self.device_config,115            cache_config=self.cache_config,116            load_config=self.load_config,117            local_rank=0,118            rank=0,119            distributed_init_method=distributed_init_method,120            lora_config=self.lora_config,121            vision_language_config=self.vision_language_config,122            is_driver_worker=True,123        )124 125        spec_decode_worker = SpecDecodeWorker.from_workers(126            proposer_worker=draft_worker, scorer_worker=target_worker)127 128        assert self.parallel_config.world_size == 1, (129            "GPUExecutor only supports single GPU.")130 131        self.driver_worker = spec_decode_worker132 133        # Load model handled in spec decode worker.134        self.driver_worker.init_device()135 136    def determine_num_available_blocks(self) -> Tuple[int, int]:137        """Determine the number of available KV blocks by invoking the138        underlying worker.139        """140        return self.driver_worker.determine_num_available_blocks()141 142    def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks) -> None:143        """Initialize the KV cache by invoking the underlying worker.144        """145        # NOTE: This is logged in the executor because there can be >1 worker146        # with other executors. We could log in the engine level, but work147        # remains to abstract away the device for non-GPU configurations.148        logger.info(f"# GPU blocks: {num_gpu_blocks}, "149                    f"# CPU blocks: {num_cpu_blocks}")150 151        self.driver_worker.initialize_cache(num_gpu_blocks, num_cpu_blocks)152 153    def execute_model(154        self,155        seq_group_metadata_list: List[SequenceGroupMetadata],156        blocks_to_swap_in: Dict[int, int],157        blocks_to_swap_out: Dict[int, int],158        blocks_to_copy: Dict[int, List[int]],159        num_lookahead_slots: int,160    ) -> List[SamplerOutput]:161        output = self.driver_worker.execute_model(162            seq_group_metadata_list=seq_group_metadata_list,163            blocks_to_swap_in=blocks_to_swap_in,164            blocks_to_swap_out=blocks_to_swap_out,165            blocks_to_copy=blocks_to_copy,166            num_lookahead_slots=num_lookahead_slots,167        )168        return output169 170    def add_lora(self, lora_request: LoRARequest) -> bool:171        assert lora_request.lora_int_id > 0, "lora_id must be greater than 0."172        return self.driver_worker.add_lora(lora_request)173 174    def remove_lora(self, lora_id: int) -> bool:175        assert lora_id > 0, "lora_id must be greater than 0."176        return self.driver_worker.remove_lora(lora_id)177 178    def list_loras(self) -> Set[int]:179        return self.driver_worker.list_loras()180 181    def check_health(self) -> None:182        # GPUExecutor will always be healthy as long as183        # it's running.184        return185 186 187class GPUExecutorAsync(GPUExecutor, ExecutorAsyncBase):188 189    async def execute_model_async(190        self,191        seq_group_metadata_list: List[SequenceGroupMetadata],192        blocks_to_swap_in: Dict[int, int],193        blocks_to_swap_out: Dict[int, int],194        blocks_to_copy: Dict[int, List[int]],195    ) -> SamplerOutput:196        output = await make_async(self.driver_worker.execute_model)(197            seq_group_metadata_list=seq_group_metadata_list,198            blocks_to_swap_in=blocks_to_swap_in,199            blocks_to_swap_out=blocks_to_swap_out,200            blocks_to_copy=blocks_to_copy)201        return output