| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """UVA-based CPU offloading using Unified Virtual Addressing.""" |
| 4 | |
| 5 | from collections.abc import Generator |
| 6 | from contextlib import nullcontext |
| 7 | |
| 8 | import torch |
| 9 | import torch.nn as nn |
| 10 | from torch.func import functional_call |
| 11 | |
| 12 | import vllm.envs as envs |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.offloader.base import BaseOffloader, should_pin_memory |
| 15 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 16 | from vllm.utils.mem_utils import format_gib |
| 17 | from vllm.utils.platform_utils import is_uva_available |
| 18 | from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor |
| 19 | |
| 20 | logger = init_logger(__name__) |
| 21 | |
| 22 | |
| 23 | class UVAOffloader(BaseOffloader): |
| 24 | """Offloader using Unified Virtual Addressing (UVA) for zero-copy access. |
| 25 | |
| 26 | This offloader moves parameters to pinned CPU memory and creates CUDA views |
| 27 | using UVA. The GPU can then directly access the CPU memory without explicit |
| 28 | transfers, at the cost of PCIe bandwidth (slower than GPU memory). |
| 29 | |
| 30 | When UVA is disabled via env var, falls back to a functional_call-based |
| 31 | approach that moves parameters on-demand. |
| 32 | |
| 33 | Args: |
| 34 | cpu_offload_max_bytes: Maximum bytes to offload to CPU. |
| 35 | cpu_offload_params: Set of parameter name segments to selectively |
| 36 | offload. If empty, all parameters are eligible up to the byte limit. |
| 37 | |
| 38 | """ |
| 39 | |
| 40 | supports_tower_offload = True |
| 41 | |
| 42 | def __init__( |
| 43 | self, |
| 44 | cpu_offload_max_bytes: int, |
| 45 | cpu_offload_params: set[str] | None = None, |
| 46 | ): |
| 47 | self.cpu_offload_max_bytes = cpu_offload_max_bytes |
| 48 | self.cpu_offload_bytes = 0 |
| 49 | self.cpu_offload_params = cpu_offload_params or set() |
| 50 | |
| 51 | self.pin_memory = should_pin_memory() |
| 52 | self.uva_offloading = ( |
| 53 | is_uva_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_UVA |
| 54 | ) |
| 55 | |
| 56 | def wrap_modules( |
| 57 | self, |
| 58 | modules_generator: Generator[nn.Module, None, None], |
| 59 | prefix: str = "", |
| 60 | ) -> list[nn.Module]: |
| 61 | """Wrap modules with UVA offloading.""" |
| 62 | if prefix: |
| 63 | prefix = f"{prefix}." |
| 64 | modules = [ |
| 65 | self._maybe_offload_to_cpu(module, prefix) for module in modules_generator |
| 66 | ] |
| 67 | if self.cpu_offload_bytes > 0: |
| 68 | logger.info( |
| 69 | "Total CPU offloaded parameters: %s", |
| 70 | format_gib(self.cpu_offload_bytes), |
| 71 | ) |
| 72 | return modules |
| 73 | |
| 74 | def _maybe_offload_to_cpu(self, module: nn.Module, prefix: str = "") -> nn.Module: |
| 75 | """Offload module parameters to CPU using UVA if budget allows.""" |
| 76 | if (params := next(module.parameters(), None)) is None: |
| 77 | return module |
| 78 | |
| 79 | device = params.device |
| 80 | |
| 81 | if device == torch.device("cpu"): |
| 82 | return module |
| 83 | |
| 84 | if self.cpu_offload_bytes >= self.cpu_offload_max_bytes: |
| 85 | return module |
| 86 | |
| 87 | # offload parameters to CPU |
| 88 | # use pin_memory if possible, which helps cudagraph capture speed |
| 89 | offloaded_parameters = False |
| 90 | for name, p in module.named_parameters(): |
| 91 | if self.cpu_offload_bytes >= self.cpu_offload_max_bytes: |
| 92 | # we use per-parameter offloading |
| 93 | # one module might have some parameters offloaded and some not |
| 94 | break |
| 95 | |
| 96 | # Skip parameters an earlier wrap_modules call already offloaded. |
| 97 | # The UVA path leaves p.device as the accelerator (a view of CPU |
| 98 | # memory), so the marker is the only way to recognize those. |
| 99 | if p.device.type == "cpu" or getattr(p, "_vllm_is_uva_offloaded", False): |
| 100 | continue |
| 101 | |
| 102 | if self.cpu_offload_params: |
| 103 | # Check if parameter belongs to the offloading set |
| 104 | # Add dots here to ensure we match full segments only |
| 105 | # e.g., "experts.w2_weight" matches "mlp.experts.w2_weight" |
| 106 | # but not "mlp.experts.w2_weight_scale" |
| 107 | should_offload = any( |
| 108 | f".{param}." in f".{prefix}{name}." |
| 109 | for param in self.cpu_offload_params |
| 110 | ) |
| 111 | if not should_offload: |
| 112 | continue |
| 113 | |
| 114 | cpu_data = p.data.to(device="cpu") |
| 115 | if self.pin_memory: |
| 116 | cpu_data = cpu_data.pin_memory() |
| 117 | |
| 118 | if not self.uva_offloading: |
| 119 | p.data = cpu_data |
| 120 | else: |
| 121 | p.data = get_accelerator_view_from_cpu_tensor(cpu_data) |
| 122 | p._vllm_is_uva_offloaded = True |
| 123 | |
| 124 | self.cpu_offload_bytes += p.data.numel() * p.data.element_size() |
| 125 | offloaded_parameters = True |
| 126 | |
| 127 | if offloaded_parameters and not self.uva_offloading: |
| 128 | original_forward = module.forward |
| 129 | |
| 130 | def forward(*args, **kwargs): |
| 131 | module.forward = original_forward |
| 132 | with nullcontext() if self.pin_memory else gpu_sync_allowed(): |
| 133 | device_state = { |
| 134 | # here we blindly call `to(device)` |
| 135 | # if the parameter is already on the device, |
| 136 | # it will be a no-op |
| 137 | k: v.to(device, non_blocking=True) |
| 138 | for k, v in module.state_dict().items() |
| 139 | } |
| 140 | |
| 141 | # set `tie_weights=False` as tied weights in original model |
| 142 | # become untied when calling .to(device) individually |
| 143 | output = functional_call( |
| 144 | module, |
| 145 | device_state, |
| 146 | args=args, |
| 147 | kwargs=kwargs, |
| 148 | tie_weights=False, |
| 149 | ) |
| 150 | module.forward = forward |
| 151 | return output |
| 152 | |
| 153 | module.forward = forward |
| 154 | |
| 155 | return module |