| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Must be imported firstly |
| 5 | import vllm.v1.worker.cpu.shm # noqa # isort: skip |
| 6 | |
| 7 | import math |
| 8 | import os |
| 9 | import sys |
| 10 | |
| 11 | import psutil |
| 12 | import torch |
| 13 | |
| 14 | from vllm import envs |
| 15 | from vllm.config import CompilationMode, VllmConfig |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.platforms import CpuArchEnum, current_platform |
| 18 | from vllm.utils.cpu_resource_utils import ( |
| 19 | get_allowed_cpu_list, |
| 20 | get_memory_node_info, |
| 21 | get_visible_memory_node, |
| 22 | ) |
| 23 | from vllm.utils.mem_utils import format_gib |
| 24 | from vllm.utils.torch_utils import set_random_seed |
| 25 | from vllm.v1.worker.cpu_model_runner import CPUModelRunner |
| 26 | from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment |
| 27 | from vllm.v1.worker.worker_base import CompilationTimes |
| 28 | |
| 29 | logger = init_logger(__name__) |
| 30 | |
| 31 | |
| 32 | class CPUWorker(Worker): |
| 33 | def __init__( |
| 34 | self, |
| 35 | vllm_config: VllmConfig, |
| 36 | local_rank: int, |
| 37 | rank: int, |
| 38 | distributed_init_method: str, |
| 39 | is_driver_worker: bool = False, |
| 40 | ): |
| 41 | # TODO: use numactl for process setup |
| 42 | # TODO: optimize for `interleaved` policy |
| 43 | # Bind memory node |
| 44 | allowed_memory_nodes = get_visible_memory_node() |
| 45 | allowed_cpu_list = get_allowed_cpu_list() |
| 46 | cpu_core = allowed_cpu_list[0] |
| 47 | |
| 48 | # TODO: some CI hosts are not correctly set, change to assertion |
| 49 | # after fix |
| 50 | if cpu_core.numa_node not in allowed_memory_nodes: |
| 51 | logger.warning( |
| 52 | "Node %s is not in available memory nodes %s.", |
| 53 | cpu_core.numa_node, |
| 54 | allowed_memory_nodes, |
| 55 | ) |
| 56 | |
| 57 | # On s390x, numa_node may be a synthetic book ID that doesn't |
| 58 | # correspond to a real memory node. Fall back to first visible node. |
| 59 | if cpu_core.numa_node in allowed_memory_nodes: |
| 60 | memory_node = cpu_core.numa_node |
| 61 | else: |
| 62 | logger.warning( |
| 63 | "CPU group key %s is not a valid memory node. " |
| 64 | "Falling back to memory node %s.", |
| 65 | cpu_core.numa_node, |
| 66 | allowed_memory_nodes[0], |
| 67 | ) |
| 68 | memory_node = allowed_memory_nodes[0] |
| 69 | |
| 70 | torch.ops._C.init_cpu_memory_env([memory_node]) |
| 71 | |
| 72 | memory_status = get_memory_node_info(memory_node) |
| 73 | memory_fraction = vllm_config.cache_config.gpu_memory_utilization |
| 74 | self.requested_cpu_memory = math.ceil( |
| 75 | memory_status.total_memory * memory_fraction |
| 76 | ) |
| 77 | available_memory = memory_status.available_memory |
| 78 | |
| 79 | if ( |
| 80 | vllm_config.cache_config.kv_cache_memory_bytes is None |
| 81 | and self.requested_cpu_memory > available_memory |
| 82 | ): |
| 83 | raise ValueError( |
| 84 | f"Available memory on node {cpu_core.numa_node} " |
| 85 | f"({format_gib(available_memory)}/" |
| 86 | f"{format_gib(memory_status.total_memory)} GiB) on startup " |
| 87 | f"is less than desired CPU memory utilization " |
| 88 | f"({vllm_config.cache_config.gpu_memory_utilization}, " |
| 89 | f"{format_gib(self.requested_cpu_memory)} GiB). " |
| 90 | "On the CPU backend, the `--gpu-memory-utilization` flag " |
| 91 | "controls the fraction of CPU memory reserved (despite its " |
| 92 | "name). To resolve: decrease `--gpu-memory-utilization` " |
| 93 | "(e.g. `--gpu-memory-utilization 0.5`) " |
| 94 | "or reduce CPU memory used by other processes." |
| 95 | ) |
| 96 | |
| 97 | super().__init__( |
| 98 | vllm_config, |
| 99 | local_rank, |
| 100 | rank, |
| 101 | distributed_init_method, |
| 102 | is_driver_worker=is_driver_worker, |
| 103 | ) |
| 104 | |
| 105 | self.parallel_config.disable_custom_all_reduce = True |
| 106 | |
| 107 | def init_device(self): |
| 108 | self.device = torch.device("cpu") |
| 109 | |
| 110 | # Check whether critical libraries are loaded |
| 111 | def check_preloaded_libs(name: str) -> bool: |
| 112 | ld_preload_list = os.environ.get("LD_PRELOAD", "") |
| 113 | if name not in ld_preload_list: |
| 114 | logger.warning( |
| 115 | "%s is not found in LD_PRELOAD. " |
| 116 | "For best performance, please follow the section " |
| 117 | "`set LD_PRELOAD` in " |
| 118 | "https://docs.vllm.ai/en/latest/getting_started/installation/cpu/ " |
| 119 | "to setup required pre-loaded libraries.", |
| 120 | name, |
| 121 | ) |
| 122 | return False |
| 123 | return True |
| 124 | |
| 125 | if sys.platform.startswith("linux"): |
| 126 | check_preloaded_libs("libtcmalloc") |
| 127 | if current_platform.get_cpu_architecture() == CpuArchEnum.X86: |
| 128 | iomp_loaded = check_preloaded_libs("libiomp") |
| 129 | if not iomp_loaded and self.vllm_config.speculative_config is not None: |
| 130 | logger.warning( |
| 131 | "Speculative decoding on CPU without Intel OpenMP in " |
| 132 | "LD_PRELOAD will cause significant performance loss. " |
| 133 | "Please follow the section `set LD_PRELOAD` in " |
| 134 | "https://docs.vllm.ai/en/latest/getting_started/" |
| 135 | "installation/cpu/ " |
| 136 | "to setup libiomp5.", |
| 137 | ) |
| 138 | |
| 139 | def skip_set_num_threads(x: int): |
| 140 | logger.warning( |
| 141 | "CPU backend doesn't allow to use " |
| 142 | "`torch.set_num_threads` after the thread binding, skip it." |
| 143 | ) |
| 144 | |
| 145 | torch.set_num_threads = skip_set_num_threads |
| 146 | |
| 147 | init_method = self.distributed_init_method |
| 148 | os.environ["VLLM_DIST_IDENT"] = ( |
| 149 | os.path.basename(init_method.removeprefix("file://")) |
| 150 | if init_method.startswith("file://") |
| 151 | else init_method.split(":")[-1] |
| 152 | ) |
| 153 | # Initialize the distributed environment. |
| 154 | init_worker_distributed_environment( |
| 155 | self.vllm_config, |
| 156 | self.rank, |
| 157 | self.distributed_init_method, |
| 158 | self.local_rank, |
| 159 | current_platform.dist_backend, |
| 160 | ) |
| 161 | if self.use_v2_model_runner: |
| 162 | logger.info_once("Using V2 Model Runner") |
| 163 | |
| 164 | # Set random seed. |
| 165 | set_random_seed(self.model_config.seed) |
| 166 | |
| 167 | # Construct the model runner |
| 168 | if self.use_v2_model_runner: |
| 169 | from vllm.v1.worker.cpu.model_runner import ( |
| 170 | CPUModelRunner as CPUModelRunnerV2, |
| 171 | ) |
| 172 | |
| 173 | self.model_runner: CPUModelRunner = CPUModelRunnerV2( # type: ignore |
| 174 | self.vllm_config, self.device |
| 175 | ) |
| 176 | else: |
| 177 | self.model_runner = CPUModelRunner(self.vllm_config, torch.device("cpu")) |
| 178 | |
| 179 | def sleep(self, level: int = 1) -> None: |
| 180 | logger.warning("sleep mode is not supported on CPU, ignore it.") |
| 181 | pass |
| 182 | |
| 183 | def wake_up(self, tags: list[str] | None = None) -> None: |
| 184 | logger.warning("sleep mode is not supported on CPU, ignore it.") |
| 185 | pass |
| 186 | |
| 187 | def discard(self, tags: tuple[str, ...]) -> None: |
| 188 | raise RuntimeError("Selective discard is not supported on CPU.") |
| 189 | |
| 190 | def _should_warm_up_model(self) -> bool: |
| 191 | # VLLM_CPU_CI_ENV always skips warmup to save CI time. |
| 192 | if envs.VLLM_CPU_CI_ENV: |
| 193 | return False |
| 194 | # With eager (CompilationMode.NONE) execution and an explicit KV |
| 195 | # cache size, warmup serves no purpose: there's no compiled graph to |
| 196 | # prewarm, and the auto KV-cache-size calc below -- which relies on |
| 197 | # a warmup forward pass having already bumped RSS to a realistic |
| 198 | # steady-state value -- is bypassed whenever the size is explicit. |
| 199 | return not ( |
| 200 | self.compilation_config.mode == CompilationMode.NONE |
| 201 | and self.cache_config.kv_cache_memory_bytes is not None |
| 202 | ) |
| 203 | |
| 204 | def determine_available_memory(self) -> int: |
| 205 | if self._should_warm_up_model(): |
| 206 | self.model_runner.warming_up_model() |
| 207 | |
| 208 | allowed_cpu_list = get_allowed_cpu_list() |
| 209 | cpu_core = allowed_cpu_list[0] |
| 210 | |
| 211 | memory_status = get_memory_node_info(cpu_core.numa_node) |
| 212 | available_memory = memory_status.available_memory |
| 213 | explicit_kv_cache_size = self.cache_config.kv_cache_memory_bytes |
| 214 | |
| 215 | kv_cache_size = None |
| 216 | msg = None |
| 217 | if explicit_kv_cache_size is not None: |
| 218 | if explicit_kv_cache_size > available_memory: |
| 219 | raise ValueError( |
| 220 | f"Available memory on node {cpu_core.numa_node} " |
| 221 | f"({format_gib(available_memory)}/" |
| 222 | f"{format_gib(memory_status.total_memory)} GiB) on kv cache" |
| 223 | f" allocation is less than requested memory for kv " |
| 224 | f"({format_gib(explicit_kv_cache_size)} GiB). " |
| 225 | "Decrease --kv-cache-memory-bytes, VLLM_CPU_KVCACHE_SPACE, " |
| 226 | "or reduce CPU memory used by other processes." |
| 227 | ) |
| 228 | kv_cache_size = explicit_kv_cache_size |
| 229 | msg = ( |
| 230 | f"Explicitly set ({format_gib(kv_cache_size)}/" |
| 231 | f"{format_gib(memory_status.total_memory)}) GiB for KV cache " |
| 232 | f"on node {cpu_core.numa_node}." |
| 233 | ) |
| 234 | else: |
| 235 | consumed_memory = psutil.Process(os.getpid()).memory_info().rss |
| 236 | requested_memory_for_kv = int(self.requested_cpu_memory - consumed_memory) |
| 237 | if ( |
| 238 | requested_memory_for_kv <= 0 |
| 239 | or requested_memory_for_kv > available_memory |
| 240 | ): |
| 241 | raise ValueError( |
| 242 | f"Available memory on node {cpu_core.numa_node} " |
| 243 | f"({format_gib(available_memory)}/" |
| 244 | f"{format_gib(memory_status.total_memory)} GiB) on kv cache" |
| 245 | f" allocation is less than requested memory for kv " |
| 246 | f"({format_gib(requested_memory_for_kv)}/" |
| 247 | f"{format_gib(self.requested_cpu_memory)} GiB). " |
| 248 | "Reduce CPU memory used by other processes." |
| 249 | ) |
| 250 | kv_cache_size = requested_memory_for_kv |
| 251 | msg = ( |
| 252 | f"Auto set ({format_gib(kv_cache_size)}/" |
| 253 | f"{format_gib(memory_status.total_memory)}) GiB for KV cache " |
| 254 | f"on node {cpu_core.numa_node}, with " |
| 255 | f"{format_gib(self.requested_cpu_memory)} GiB requested memory" |
| 256 | f" for the worker. {format_gib(consumed_memory)} GiB" |
| 257 | f" memory was consumed by non-kv usages." |
| 258 | ) |
| 259 | |
| 260 | logger.info(msg) |
| 261 | |
| 262 | return kv_cache_size |
| 263 | |
| 264 | def compile_or_warm_up_model(self) -> CompilationTimes: |
| 265 | # Note: the model has been compiled in determine_available_memory(), |
| 266 | # Only compile here for models without kv cache |
| 267 | if len(self.model_runner.kv_caches) == 0 and self._should_warm_up_model(): |
| 268 | self.model_runner.warming_up_model() |
| 269 | # Reset the seed to ensure that the random state is not affected by |
| 270 | # the model initialization and profiling. |
| 271 | set_random_seed(self.model_config.seed) |
| 272 | return CompilationTimes( |
| 273 | language_model=self.compilation_config.compilation_time, |
| 274 | encoder=self.compilation_config.encoder_compilation_time, |
| 275 | ) |