From bcc2306cefa4179c548d3e638e7a22a88d281733 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Wed, 15 Apr 2026 11:43:29 +0800 Subject: [PATCH] [Bugfix] Respect VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY in prefetch offloader (#37699) Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Co-authored-by: Isotr0py --- vllm/model_executor/offloader/__init__.py | 2 ++ vllm/model_executor/offloader/base.py | 14 ++++++++++++++ vllm/model_executor/offloader/prefetch.py | 9 ++++----- vllm/model_executor/offloader/uva.py | 9 +++------ 4 files changed, 23 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/offloader/__init__.py b/vllm/model_executor/offloader/__init__.py index a6522ff7c0..f1b49c69ef 100644 --- a/vllm/model_executor/offloader/__init__.py +++ b/vllm/model_executor/offloader/__init__.py @@ -8,6 +8,7 @@ from vllm.model_executor.offloader.base import ( create_offloader, get_offloader, set_offloader, + should_pin_memory, ) from vllm.model_executor.offloader.prefetch import PrefetchOffloader from vllm.model_executor.offloader.uva import UVAOffloader @@ -20,4 +21,5 @@ __all__ = [ "create_offloader", "get_offloader", "set_offloader", + "should_pin_memory", ] diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index 7cb0ddfd18..b8c1b6cfa4 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -10,7 +10,9 @@ from typing import TYPE_CHECKING import torch.nn as nn +import vllm.envs as envs from vllm.logger import init_logger +from vllm.utils.platform_utils import is_pin_memory_available if TYPE_CHECKING: from vllm.config import OffloadConfig @@ -18,6 +20,18 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def should_pin_memory() -> bool: + """Check if pinned memory should be used for weight offloading. + + Combines the platform capability check with the user override env var. + On unified-memory systems (e.g. GH200) pinned memory eats into GPU + memory, so users can disable it via VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY. + """ + return ( + is_pin_memory_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY + ) + + """ class relation: diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index 5bdde8c3a1..cc04367d54 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -20,8 +20,7 @@ import torch.nn as nn # Import prefetch_ops to register custom ops at module load time import vllm.model_executor.offloader.prefetch_ops # noqa: F401 from vllm.logger import init_logger -from vllm.model_executor.offloader.base import BaseOffloader -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.model_executor.offloader.base import BaseOffloader, should_pin_memory logger = init_logger(__name__) @@ -528,7 +527,7 @@ class _ModuleOffloader: gpu_buffer = offloader._gpu_buffer assert cpu_storage is not None, "CPU storage not initialized" assert gpu_buffer is not None, "GPU buffer not assigned" - assert not is_pin_memory_available() or cpu_storage.is_pinned(), ( + assert not should_pin_memory() or cpu_storage.is_pinned(), ( f"CPU storage for {name} is not pinned! " "non_blocking=True H2D copy from non-pinned memory " "causes stream synchronization that breaks " @@ -629,7 +628,7 @@ class _CpuParamOffloader(_BaseParamOffloader): original GPU tensor is garbage collected. """ param = self._param - pin_memory = is_pin_memory_available() + pin_memory = should_pin_memory() # Create pinned CPU storage and copy current GPU data self._cpu_storage = torch.empty_strided( @@ -666,7 +665,7 @@ class _CpuParamOffloader(_BaseParamOffloader): param = self._param if param.data.device.type == "cpu": - if is_pin_memory_available() and not param.data.is_pinned(): + if should_pin_memory() and not param.data.is_pinned(): pinned = torch.empty_strided( size=param.data.size(), stride=param.data.stride(), diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py index c524e43cdd..51eb1a14fc 100644 --- a/vllm/model_executor/offloader/uva.py +++ b/vllm/model_executor/offloader/uva.py @@ -10,9 +10,9 @@ from torch.func import functional_call import vllm.envs as envs from vllm.logger import init_logger -from vllm.model_executor.offloader.base import BaseOffloader +from vllm.model_executor.offloader.base import BaseOffloader, should_pin_memory from vllm.utils.mem_utils import format_gib -from vllm.utils.platform_utils import is_pin_memory_available, is_uva_available +from vllm.utils.platform_utils import is_uva_available from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor logger = init_logger(__name__) @@ -43,10 +43,7 @@ class UVAOffloader(BaseOffloader): self.cpu_offload_bytes = 0 self.cpu_offload_params = cpu_offload_params or set() - self.pin_memory = ( - is_pin_memory_available() - and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY - ) + self.pin_memory = should_pin_memory() self.uva_offloading = ( is_uva_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_UVA )