cpu_executor.py 14.1 KB
Newer Older
1
2
3
import os
from functools import partial
from typing import Any, Awaitable, List, Optional, Set, Tuple, Union
4
5
6

import torch

7
import vllm.envs as envs
8
9
from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
                         SchedulerConfig)
10
from vllm.executor.executor_base import ExecutorAsyncBase, ExecutorBase
11
12
from vllm.executor.multiproc_worker_utils import (ProcessWorkerWrapper,
                                                  ResultHandler, WorkerMonitor)
13
14
from vllm.logger import init_logger
from vllm.lora.request import LoRARequest
15
from vllm.model_executor.layers.sampler import SamplerOutput
16
from vllm.prompt_adapter.request import PromptAdapterRequest
17
from vllm.sequence import ExecuteModelRequest
18
from vllm.utils import (GiB_bytes, get_distributed_init_method, get_open_port,
19
20
                        get_vllm_instance_id, make_async)
from vllm.worker.worker_base import WorkerWrapperBase
21
22
23
24
25
26

logger = init_logger(__name__)


class CPUExecutor(ExecutorBase):

27
28
    uses_ray: bool = False

29
30
31
    def _init_executor(self) -> None:
        assert self.device_config.device_type == "cpu"
        assert self.lora_config is None, "cpu backend doesn't support LoRA"
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

        #
        # Environment variables for CPU executor
        #

        # Ensure that VLLM_INSTANCE_ID is set, to be inherited by workers
        os.environ["VLLM_INSTANCE_ID"] = get_vllm_instance_id()

        # Disable torch async compiling which won't work with daemonic processes
        os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"

        # Intel OpenMP setting
        ld_prealod_str = os.getenv("LD_PRELOAD", "")
        if "libiomp5.so" in ld_prealod_str:
            # The time(milliseconds) that a thread should wait after
            # completing the execution of a parallel region, before sleeping.
            os.environ['KMP_BLOCKTIME'] = "1"
            # Prevents the CPU to run into low performance state
            os.environ['KMP_TPAUSE'] = "0"
            # Provides fine granularity parallelism
            os.environ['KMP_FORKJOIN_BARRIER_PATTERN'] = "dist,dist"
            os.environ['KMP_PLAIN_BARRIER_PATTERN'] = "dist,dist"
            os.environ['KMP_REDUCTION_BARRIER_PATTERN'] = "dist,dist"

        # To hint IPEX uses shared memory based AllReduce
        os.environ["LOCAL_WORLD_SIZE"] = str(
            self.parallel_config.tensor_parallel_size)

60
61
62
63
        self.model_config = _verify_and_get_model_config(self.model_config)
        self.cache_config = _verify_and_get_cache_config(self.cache_config)
        self.scheduler_config = _verify_and_get_scheduler_config(
            self.scheduler_config)
64
65
        self.parallel_config = _verify_and_get_parallel_config(
            self.parallel_config)
66

67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
        # Multiprocessing-based executor does not support multi-node setting.
        # Since it only works for single node, we can use the loopback address
        # 127.0.0.1 for communication.
        ip = "127.0.0.1"
        port = get_open_port()
        self.distributed_init_method = get_distributed_init_method(ip, port)

        is_async = isinstance(self, CPUExecutorAsync)

        world_size = self.parallel_config.tensor_parallel_size
        result_handler = ResultHandler()
        self.parallel_worker_tasks: Optional[Union[Any, Awaitable[Any]]] = None
        self.workers = []

        if is_async:
            self.workers = [
                ProcessWorkerWrapper(
                    result_handler,
                    partial(
                        self._create_worker,
                        rank=rank,
                        local_rank=rank,
                    )) for rank in range(0, world_size)
            ]
            self.driver_worker = self.workers[0]
            self.workers = self.workers[1:]
            self.driver_method_invoker = _async_driver_method_invoker
        else:
            self.driver_worker = self._create_worker()
            self.driver_method_invoker = _driver_method_invoker

            if world_size != 1:
                self.workers = [
                    ProcessWorkerWrapper(
                        result_handler,
                        partial(
                            self._create_worker,
                            rank=rank,
                            local_rank=rank,
                        )) for rank in range(1, world_size)
                ]

109
        self.worker_monitor = None
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
        if world_size != 1 or is_async:
            if is_async:
                async_worker_list = self.workers + [self.driver_worker]
            else:
                async_worker_list = self.workers
            self.worker_monitor = WorkerMonitor(async_worker_list,
                                                result_handler)
            result_handler.start()
            self.worker_monitor.start()

        self._run_workers("init_device")
        self._run_workers("load_model")

    def _create_worker(
        self,
        local_rank: int = 0,
        rank: int = 0,
    ):
        worker_module_name = "vllm.worker.cpu_worker"
        worker_class_name = "CPUWorker"

        wrapper = WorkerWrapperBase(
            worker_module_name=worker_module_name,
            worker_class_name=worker_class_name,
        )
135

136
        assert self.distributed_init_method is not None
137

138
        kwargs = dict(
139
140
141
142
143
            model_config=self.model_config,
            parallel_config=self.parallel_config,
            scheduler_config=self.scheduler_config,
            device_config=self.device_config,
            cache_config=self.cache_config,
144
            load_config=self.load_config,
145
146
147
            local_rank=local_rank,
            rank=rank,
            distributed_init_method=self.distributed_init_method,
148
149
            lora_config=self.lora_config,
            kv_cache_dtype=self.cache_config.cache_dtype,
150
            prompt_adapter_config=self.prompt_adapter_config,
151
            is_driver_worker=rank == 0,
152
        )
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
        wrapper.init_worker(**kwargs)

        return wrapper.worker

    def _run_workers(
        self,
        method: str,
        *args,
        async_run_remote_workers_only: bool = False,
        max_concurrent_workers: Optional[int] = None,
        **kwargs,
    ) -> Any:
        """Runs the given method on all workers.

        Args:
            async_run_remote_workers_only: If True the method will be run only
                in the remote workers, not the driver worker. It will also be
                run asynchronously and return a list of futures rather than
                blocking on the results.
        """

        if max_concurrent_workers:
            raise NotImplementedError(
                "max_concurrent_workers is not supported yet.")

        # Start the workers first.
        worker_outputs = [
            worker.execute_method(method, *args, **kwargs)
            for worker in self.workers
        ]

        if async_run_remote_workers_only:
            # Just return futures
            return worker_outputs

        driver_worker_output = self.driver_method_invoker(
            self.driver_worker, method, *args, **kwargs)

        # Get the results of the workers.
        return [driver_worker_output
                ] + [output.get() for output in worker_outputs]
194

195
    def determine_num_available_blocks(self) -> Tuple[int, int]:
196
197
198
        """Determine the number of available KV blocks by invoking the
        underlying worker.
        """
199
200
        return self.driver_method_invoker(self.driver_worker,
                                          "determine_num_available_blocks")
201
202
203
204
205
206
207
208

    def initialize_cache(self, num_gpu_blocks: int,
                         num_cpu_blocks: int) -> None:
        """Initialize the KV cache by invoking the underlying worker.
        """
        # NOTE: We log here to avoid multiple logs when number of workers is
        # greater than one. We could log in the engine, but not all executors
        # have GPUs.
209
210
211
        # NOTE: `cpu block` for CPU backend is located on CPU memory but is
        # referred as `gpu block`. Because we want to reuse the existing block
        # management procedure.
212
        logger.info("# CPU blocks: %d", num_gpu_blocks)
213
214
215
216

        self._run_workers("initialize_cache",
                          num_gpu_blocks=num_gpu_blocks,
                          num_cpu_blocks=num_cpu_blocks)
217

218
219
220
    def execute_model(
            self,
            execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]:
221
222
223
224
225
226
227
228
        if (self.parallel_config.tensor_parallel_size > 1
                and self.parallel_worker_tasks is None):
            self.parallel_worker_tasks = self._run_workers(
                "start_worker_execution_loop",
                async_run_remote_workers_only=True,
            )
        output = self.driver_method_invoker(self.driver_worker,
                                            "execute_model", execute_model_req)
229
230
        return output

231
232
233
234
235
236
237
238
239
240
241
242
243
244
    def stop_remote_worker_execution_loop(self) -> None:
        if self.parallel_worker_tasks is None:
            return
        """
        Passing None will cause the driver to stop the model execution
        loop running in each of the remote workers.
        """
        self.driver_method_invoker(self.driver_worker, "execute_model", None)
        parallel_worker_tasks = self.parallel_worker_tasks
        self.parallel_worker_tasks = None
        # Ensure that workers exit model loop cleanly
        # (this will raise otherwise)
        self._wait_for_tasks_completion(parallel_worker_tasks)

245
    def add_lora(self, lora_request: LoRARequest) -> bool:
246
        return all(self._run_workers("add_lora", lora_request))
247
248

    def remove_lora(self, lora_id: int) -> bool:
249
        return all(self._run_workers("remove_lora", lora_id))
250

251
    def pin_lora(self, lora_id: int) -> bool:
252
253
254
255
256
        assert lora_id > 0, "lora_id must be greater than 0."
        return all(self._run_workers(
            "pin_lora",
            lora_id=lora_id,
        ))
257

258
    def list_loras(self) -> Set[int]:
259
        return self.driver_method_invoker(self.driver_worker, "list_loras")
260

261
262
    def add_prompt_adapter(
            self, prompt_adapter_request: PromptAdapterRequest) -> bool:
263
264
265
266
267
        return all(
            self._run_workers(
                "add_prompt_adapter",
                prompt_adapter_request,
            ))
268
269

    def remove_prompt_adapter(self, prompt_adapter_id: int) -> bool:
270
271
272
273
274
        return all(
            self._run_workers(
                "remove_prompt_adapter",
                prompt_adapter_id,
            ))
275
276

    def list_prompt_adapters(self) -> Set[int]:
277
278
        return self.driver_method_invoker(self.driver_worker,
                                          "list_prompt_adapters")
279
280

    def pin_prompt_adapter(self, prompt_adapter_id: int) -> bool:
281
282
283
284
        return all(self._run_workers(
            "pin_prompt_adapter",
            prompt_adapter_id,
        ))
285

286
    def check_health(self) -> None:
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
        """Raises an error if engine is unhealthy."""
        if self.worker_monitor is not None and not self.worker_monitor.is_alive(
        ):
            raise RuntimeError("Worker processes are not running")

    def shutdown(self):
        if (worker_monitor := getattr(self, "worker_monitor",
                                      None)) is not None:
            worker_monitor.close()

    def _wait_for_tasks_completion(self, parallel_worker_tasks: Any) -> None:
        """Wait for futures returned from _run_workers() with
        async_run_remote_workers_only to complete."""
        for result in parallel_worker_tasks:
            result.get()
302

303
304
305
306
307
308
    def start_profile(self) -> None:
        self.driver_method_invoker(self.driver_worker, "start_profile")

    def stop_profile(self) -> None:
        self.driver_method_invoker(self.driver_worker, "stop_profile")

309

310
311
312
class CPUExecutorAsync(CPUExecutor, ExecutorAsyncBase):

    async def execute_model_async(
313
314
            self,
            execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]:
315
        output = await make_async(self.execute_model
316
                                  )(execute_model_req=execute_model_req, )
317
318
319
        return output

    async def check_health_async(self) -> None:
320
        self.check_health()
321
322


323
324
325
326
327
328
329
330
331
332
333
334
def _verify_and_get_model_config(config: ModelConfig) -> ModelConfig:
    if config.dtype == torch.float16:
        logger.warning("float16 is not supported on CPU, casting to bfloat16.")
        config.dtype = torch.bfloat16
    if not config.enforce_eager:
        logger.warning(
            "CUDA graph is not supported on CPU, fallback to the eager "
            "mode.")
        config.enforce_eager = True
    return config


335
336
337
338
339
340
341
342
343
def _verify_and_get_scheduler_config(
        config: SchedulerConfig) -> SchedulerConfig:
    if config.chunked_prefill_enabled:
        logger.warning("Chunked prefill is not supported on CPU, disable it.")
        config.chunked_prefill_enabled = False

    return config


344
345
346
347
348
def _verify_and_get_cache_config(config: CacheConfig) -> CacheConfig:
    if config.enable_prefix_caching:
        logger.warning("Prefix caching is not supported on CPU, disable it.")
        config.enable_prefix_caching = False

349
    kv_cache_space = envs.VLLM_CPU_KVCACHE_SPACE
350
351
352

    if kv_cache_space >= 0:
        if kv_cache_space == 0:
353
            config.cpu_kvcache_space_bytes = 4 * GiB_bytes  # type: ignore
354
355
356
            logger.warning("Environment variable VLLM_CPU_KVCACHE_SPACE (GB) "
                           "for CPU backend is not set, using 4 by default.")
        else:
357
            config.cpu_kvcache_space_bytes = kv_cache_space * GiB_bytes  # type: ignore
358
359
360
361
362
363
    else:
        raise RuntimeError(
            "Invalid environment variable VLLM_CPU_KVCACHE_SPACE"
            f" {kv_cache_space}, expect a positive integer value.")

    return config
364
365


366
367
368
369
370
371
372
373
374
375
def _verify_and_get_parallel_config(config: ParallelConfig) -> ParallelConfig:
    if (config.distributed_executor_backend is not None
            and config.distributed_executor_backend != "mp"):
        logger.warning(
            "%s is not supported on CPU, fallback to mp distributed executor "
            "backend.", config.distributed_executor_backend)
        config.distributed_executor_backend = "mp"
    return config


376
377
378
379
380
381
def _driver_method_invoker(driver, method: str, *args, **kwargs):
    return getattr(driver, method)(*args, **kwargs)


def _async_driver_method_invoker(driver, method: str, *args, **kwargs):
    return driver.execute_method(method, *args, **kwargs).get()