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, 16:00 UTC
[sleep mode] clear pytorch cache after sleep (#15248)
Signed-off-by: <villard@us.ibm.com>
Lionel Villard committed
Apr 1, 2025 at 01:58 UTC
4a9ce1784c4497d71b9a4497619f60d97271bd00
1 file changed
+7
-3
vllm/device_allocator/cumem.py
+7
-3
@@ -8,6 +8,7 @@
8
# not sure why, they are created from a different context.
9
# the only successful approach is to call cuda driver API in C.
10
import dataclasses
11
+import gc
12
import os
13
from contextlib import contextmanager
14
from typing import Any, Callable, Dict, Optional, Tuple, Union
@@ -175,7 +176,7 @@ class CuMemAllocator:
176
str]] = None) -> None:
177
"""
178
Put the allocator in sleep mode.
178
- All data in the memory allocation with the specified tag will be
179
+ All data in the memory allocation with the specified tag will be
180
offloaded to CPU memory, and others will be discarded.
181
182
:param offload_tags: The tags of the memory allocation that will be
@@ -204,10 +205,13 @@ class CuMemAllocator:
205
data.cpu_backup_tensor = cpu_backup_tensor
206
unmap_and_release(handle)
207
208
+ gc.collect()
209
+ torch.cuda.empty_cache()
210
+
211
def wake_up(self):
212
"""
213
Wake up the allocator from sleep mode.
210
- All data that is previously offloaded will be loaded back to GPU
214
+ All data that is previously offloaded will be loaded back to GPU
215
memory, and the rest of the data will have empty memory."""
216
for ptr, data in self.pointer_to_data.items():
217
handle = data.handle
@@ -225,7 +229,7 @@ class CuMemAllocator:
229
def use_memory_pool(self, tag: Optional[str] = None):
230
"""
231
A context manager to use the memory pool.
228
- All memory allocation created inside the context will be allocated
232
+ All memory allocation created inside the context will be allocated
233
in the memory pool, and has the specified tag.
234
235
:param tag: The tag of the memory allocation. If None, the default tag