A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 13:30 UTC
Restore buffers when wake up from level 2 sleep (#16564) (#16889)
Signed-off-by: Han <zh950713@gmail.com>
Han Zhang committed
Apr 21, 2025 at 20:18 UTC
d41faaf9df6d1a741d5fdd4a282b16783cace888
2 files changed
+40
vllm/v1/worker/gpu_worker.py
+20
@@ -54,6 +54,9 @@ class Worker(WorkerBase):
54
from vllm.utils import init_cached_hf_modules
55
init_cached_hf_modules()
56
57
+ # Buffers saved before sleep
58
+ self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
59
+
60
# Torch profiler. Enabled and configured through env vars:
61
# VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
62
if envs.VLLM_TORCH_PROFILER_DIR:
@@ -73,6 +76,15 @@ class Worker(WorkerBase):
76
77
def sleep(self, level: int = 1) -> None:
78
free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
79
+
80
+ # Save the buffers before level 2 sleep
81
+ if level == 2:
82
+ model = self.model_runner.model
83
+ self._sleep_saved_buffers = {
84
+ name: buffer.cpu().clone()
85
+ for name, buffer in model.named_buffers()
86
+ }
87
+
88
allocator = CuMemAllocator.get_instance()
89
allocator.sleep(offload_tags=("weights", ) if level == 1 else tuple())
90
free_bytes_after_sleep, total = torch.cuda.mem_get_info()
@@ -88,6 +100,14 @@ class Worker(WorkerBase):
100
allocator = CuMemAllocator.get_instance()
101
allocator.wake_up(tags)
102
103
+ # Restore the buffers after level 2 sleep
104
+ if len(self._sleep_saved_buffers):
105
+ model = self.model_runner.model
106
+ for name, buffer in model.named_buffers():
107
+ if name in self._sleep_saved_buffers:
108
+ buffer.data.copy_(self._sleep_saved_buffers[name].data)
109
+ self._sleep_saved_buffers = {}
110
+
111
def init_device(self):
112
if self.device_config.device.type == "cuda":
113
# torch.distributed.all_reduce does not free the input tensor until
vllm/worker/worker.py
+20
@@ -95,6 +95,9 @@ class Worker(LocalOrDistributedWorkerBase):
95
self.gpu_cache: Optional[List[List[torch.Tensor]]] = None
96
self._seq_group_metadata_cache: Dict[str, SequenceGroupMetadata] = {}
97
98
+ # Buffers saved before sleep
99
+ self._sleep_saved_buffers: Dict[str, torch.Tensor] = {}
100
+
101
# Torch profiler. Enabled and configured through env vars:
102
# VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
103
if envs.VLLM_TORCH_PROFILER_DIR:
@@ -124,6 +127,15 @@ class Worker(LocalOrDistributedWorkerBase):
127
128
def sleep(self, level: int = 1) -> None:
129
free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
130
+
131
+ # Save the buffers before level 2 sleep
132
+ if level == 2:
133
+ model = self.model_runner.model
134
+ self._sleep_saved_buffers = {
135
+ name: buffer.cpu().clone()
136
+ for name, buffer in model.named_buffers()
137
+ }
138
+
139
allocator = CuMemAllocator.get_instance()
140
allocator.sleep(offload_tags=("weights", ) if level == 1 else tuple())
141
free_bytes_after_sleep, total = torch.cuda.mem_get_info()
@@ -139,6 +151,14 @@ class Worker(LocalOrDistributedWorkerBase):
151
allocator = CuMemAllocator.get_instance()
152
allocator.wake_up(tags=tags)
153
154
+ # Restore the buffers after level 2 sleep
155
+ if len(self._sleep_saved_buffers):
156
+ model = self.model_runner.model
157
+ for name, buffer in model.named_buffers():
158
+ if name in self._sleep_saved_buffers:
159
+ buffer.data.copy_(self._sleep_saved_buffers[name].data)
160
+ self._sleep_saved_buffers = {}
161
+
162
def init_device(self) -> None:
163
if self.device_config.device.type == "cuda":
164
# torch.distributed.all_reduce does not free the input tensor until