| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Configuration for model weight offloading.""" |
| 4 | |
| 5 | import warnings |
| 6 | from typing import Literal |
| 7 | |
| 8 | from pydantic import Field, model_validator |
| 9 | |
| 10 | from vllm.config.utils import config |
| 11 | |
| 12 | OffloadBackend = Literal["auto", "uva", "prefetch"] |
| 13 | |
| 14 | |
| 15 | @config |
| 16 | class UVAOffloadConfig: |
| 17 | """Configuration for UVA (Unified Virtual Addressing) CPU offloading. |
| 18 | |
| 19 | Uses zero-copy access from CPU-pinned memory. Simple but requires |
| 20 | fast CPU-GPU interconnect. |
| 21 | """ |
| 22 | |
| 23 | cpu_offload_gb: float = Field(default=0, ge=0) |
| 24 | """The space in GiB to offload to CPU, per GPU. Default is 0, which means |
| 25 | no offloading. Intuitively, this argument can be seen as a virtual way to |
| 26 | increase the GPU memory size. For example, if you have one 24 GB GPU and |
| 27 | set this to 10, virtually you can think of it as a 34 GB GPU. Then you can |
| 28 | load a 13B model with BF16 weight, which requires at least 26GB GPU memory. |
| 29 | Note that this requires fast CPU-GPU interconnect, as part of the model is |
| 30 | loaded from CPU memory to GPU memory on the fly in each model forward pass. |
| 31 | This uses UVA (Unified Virtual Addressing) for zero-copy access. |
| 32 | """ |
| 33 | |
| 34 | cpu_offload_params: set[str] = Field(default_factory=set) |
| 35 | """The set of parameter name segments to target for CPU offloading. |
| 36 | Unmatched parameters are not offloaded. If this set is empty, parameters |
| 37 | are offloaded non-selectively until the memory limit defined by |
| 38 | `cpu_offload_gb` is reached. |
| 39 | Examples: |
| 40 | - For parameter name "mlp.experts.w2_weight": |
| 41 | - "experts" or "experts.w2_weight" will match. |
| 42 | - "expert" or "w2" will NOT match (must be exact segments). |
| 43 | This allows distinguishing parameters like "w2_weight" and "w2_weight_scale". |
| 44 | """ |
| 45 | |
| 46 | |
| 47 | @config |
| 48 | class PrefetchOffloadConfig: |
| 49 | """Configuration for prefetch-based CPU offloading. |
| 50 | |
| 51 | Groups layers and uses async H2D prefetch to hide transfer latency. |
| 52 | """ |
| 53 | |
| 54 | offload_group_size: int = Field(default=0, ge=0) |
| 55 | """Group every N layers together. Offload last `offload_num_in_group` |
| 56 | layers of each group. Default is 0 (disabled). |
| 57 | Example: group_size=8, num_in_group=2 offloads layers 6,7,14,15,22,23,... |
| 58 | Unlike cpu_offload_gb, this uses explicit async prefetching to hide transfer |
| 59 | latency. |
| 60 | """ |
| 61 | |
| 62 | offload_num_in_group: int = Field(default=1, ge=1) |
| 63 | """Number of layers to offload per group. |
| 64 | Must be <= offload_group_size. Default is 1.""" |
| 65 | |
| 66 | offload_prefetch_step: int = Field(default=1, ge=0) |
| 67 | """Number of layers to prefetch ahead. |
| 68 | Higher values hide more latency but use more GPU memory. Default is 1.""" |
| 69 | |
| 70 | offload_params: set[str] = Field(default_factory=set) |
| 71 | """The set of parameter name segments to target for prefetch offloading. |
| 72 | Unmatched parameters are not offloaded. If this set is empty, ALL |
| 73 | parameters of each offloaded layer are offloaded. |
| 74 | Uses segment matching: "w13_weight" matches "mlp.experts.w13_weight" |
| 75 | but not "mlp.experts.w13_weight_scale". |
| 76 | """ |
| 77 | |
| 78 | |
| 79 | @config |
| 80 | class OffloadConfig: |
| 81 | """Configuration for model weight offloading to reduce GPU memory usage.""" |
| 82 | |
| 83 | offload_backend: OffloadBackend = "auto" |
| 84 | """The backend for weight offloading. Options: |
| 85 | - "auto": Selects based on which sub-config has non-default values |
| 86 | (prefetch if offload_group_size > 0, uva if cpu_offload_gb > 0). |
| 87 | - "uva": UVA (Unified Virtual Addressing) zero-copy offloading. |
| 88 | - "prefetch": Async prefetch with group-based layer offloading. |
| 89 | """ |
| 90 | |
| 91 | uva: UVAOffloadConfig = Field(default_factory=UVAOffloadConfig) |
| 92 | """Parameters for UVA offloading backend.""" |
| 93 | |
| 94 | prefetch: PrefetchOffloadConfig = Field(default_factory=PrefetchOffloadConfig) |
| 95 | """Parameters for prefetch offloading backend.""" |
| 96 | |
| 97 | @model_validator(mode="after") |
| 98 | def validate_offload_config(self) -> "OffloadConfig": |
| 99 | """Validate offload configuration constraints.""" |
| 100 | if self.offload_backend == "prefetch" or self.prefetch.offload_group_size > 0: |
| 101 | if self.prefetch.offload_num_in_group > self.prefetch.offload_group_size: |
| 102 | raise ValueError( |
| 103 | f"offload_num_in_group ({self.prefetch.offload_num_in_group})" |
| 104 | f" must be <= offload_group_size" |
| 105 | f" ({self.prefetch.offload_group_size})" |
| 106 | ) |
| 107 | if self.prefetch.offload_prefetch_step < 1: |
| 108 | raise ValueError( |
| 109 | f"offload_prefetch_step" |
| 110 | f" ({self.prefetch.offload_prefetch_step})" |
| 111 | f" must be >= 1 when prefetch offloading is enabled" |
| 112 | f" (offload_group_size > 0)" |
| 113 | ) |
| 114 | |
| 115 | # Warn if both backends have non-default values |
| 116 | uva_active = self.uva.cpu_offload_gb > 0 |
| 117 | prefetch_active = self.prefetch.offload_group_size > 0 |
| 118 | if self.offload_backend == "uva" and prefetch_active: |
| 119 | warnings.warn( |
| 120 | "Prefetch offload fields are set but offload_backend='uva'. " |
| 121 | "Prefetch settings will be ignored.", |
| 122 | stacklevel=2, |
| 123 | ) |
| 124 | elif self.offload_backend == "prefetch" and uva_active: |
| 125 | warnings.warn( |
| 126 | "UVA offload fields are set but offload_backend='prefetch'. " |
| 127 | "UVA settings will be ignored.", |
| 128 | stacklevel=2, |
| 129 | ) |
| 130 | elif self.offload_backend == "auto" and uva_active and prefetch_active: |
| 131 | warnings.warn( |
| 132 | "Both UVA and prefetch offload fields are set with " |
| 133 | "offload_backend='auto'. Prefetch backend will be selected. " |
| 134 | "Set offload_backend explicitly to suppress this warning.", |
| 135 | stacklevel=2, |
| 136 | ) |
| 137 | return self |
| 138 | |
| 139 | def compute_hash(self) -> str: |
| 140 | """Provide a hash that uniquely identifies all the offload configs. |
| 141 | |
| 142 | All fields are included because PrefetchOffloader patches module |
| 143 | forwards and inserts custom ops (wait_prefetch, start_prefetch) |
| 144 | into the computation graph. Changing any offload setting can |
| 145 | alter which layers are hooked and how prefetch indices are |
| 146 | computed, so the compilation cache must distinguish them. |
| 147 | """ |
| 148 | from vllm.config.utils import get_hash_factors, hash_factors |
| 149 | |
| 150 | factors = get_hash_factors(self, ignored_factors=set()) |
| 151 | hash_str = hash_factors(factors) |
| 152 | return hash_str |