Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
OpenDAS
vllm_cscc
Commits
4a9ce178
Unverified
Commit
4a9ce178
authored
Apr 01, 2025
by
Lionel Villard
Committed by
GitHub
Mar 31, 2025
Browse files
[sleep mode] clear pytorch cache after sleep (#15248)
Signed-off-by: <villard@us.ibm.com>
parent
7e4e709b
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
7 additions
and
3 deletions
+7
-3
vllm/device_allocator/cumem.py
vllm/device_allocator/cumem.py
+7
-3
No files found.
vllm/device_allocator/cumem.py
View file @
4a9ce178
...
...
@@ -8,6 +8,7 @@
# not sure why, they are created from a different context.
# the only successful approach is to call cuda driver API in C.
import
dataclasses
import
gc
import
os
from
contextlib
import
contextmanager
from
typing
import
Any
,
Callable
,
Dict
,
Optional
,
Tuple
,
Union
...
...
@@ -204,6 +205,9 @@ class CuMemAllocator:
data
.
cpu_backup_tensor
=
cpu_backup_tensor
unmap_and_release
(
handle
)
gc
.
collect
()
torch
.
cuda
.
empty_cache
()
def
wake_up
(
self
):
"""
Wake up the allocator from sleep mode.
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment