Add a level 3 sleep/wake_up that offloads tensors to disk

Co-authored-by: aavarghese <[email protected]>
Co-authored-by: manoelmarques <[email protected]>
Signed-off-by: Manoel Marques <[email protected]>

Author: manoelmarques

tests/basic_correctness/test_cumem.py

@@ -179,6 +179,30 @@ def test_end_to_end(monkeypatch: pytest.MonkeyPatch, model: str, use_v1: bool):
         # cmp output
         assert output[0].outputs[0].text == output3[0].outputs[0].text
 
+        # test sleep level 3 here.
+        llm.sleep(level=3)
+
+        free_gpu_bytes_after_sleep, total = torch.cuda.mem_get_info()
+        used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline
+        # now the memory usage is mostly cudagraph memory pool,
+        # and it should be less than the model weights (1B model, 2GiB weights)
+
+        # NOTE: In V1, the memory buffer for logits (max_num_reqs x vocab_size)
+        # is captured but cannot be releasesd from PyTorch due to a known bug,
+        # therefore high memory usage after `llm.sleep` is called is expected.
+        # FIXME(youkaichao & ywang96): Fix memory buffer issue with sleep mode
+        # in V1.
+        if use_v1:
+            assert used_bytes < 7 * GiB_bytes
+        else:
+            assert used_bytes < 2 * GiB_bytes
+
+        llm.wake_up()
+        output2 = llm.generate(prompt, sampling_params)
+
+        # cmp output
+        assert output[0].outputs[0].text == output2[0].outputs[0].text
+
 
 @create_new_process_for_each_test()
 def test_deep_sleep():

tests/entrypoints/openai/test_sleep.py

@@ -60,3 +60,16 @@ def test_sleep_mode():
         response = requests.get(remote_server.url_for("is_sleeping"))
         assert response.status_code == 200
         assert response.json().get("is_sleeping") is False
+
+        response = requests.post(remote_server.url_for("/sleep"),
+                                 data={"level": "3"})
+        assert response.status_code == 200
+        response = requests.get(remote_server.url_for("/is_sleeping"))
+        assert response.status_code == 200
+        assert response.json().get("is_sleeping") is True
+
+        response = requests.post(remote_server.url_for("/wake_up"))
+        assert response.status_code == 200
+        response = requests.get(remote_server.url_for("/is_sleeping"))
+        assert response.status_code == 200
+        assert response.json().get("is_sleeping") is False

vllm/device_allocator/cumem.py

@@ -8,20 +8,68 @@
 # both of them failed because of cuda context mismatch.
 # not sure why, they are created from a different context.
 # the only successful approach is to call cuda driver API in C.
+import contextlib
+import ctypes
 import dataclasses
 import gc
+import io
+import mmap
 import os
+import struct
+import uuid
 from contextlib import contextmanager
 from typing import Any, Callable, Optional, Union
 
 import torch
 
+import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.utils import is_pin_memory_available
 
 logger = init_logger(__name__)
 
 
+def _copy_from_cuda_to_bytes(scr_ptr: int, size_in_bytes: int) -> bytes:
+    dest_ptr = ctypes.create_string_buffer(size_in_bytes)
+    libcudart.cudaMemcpy(dest_ptr, scr_ptr, size_in_bytes)
+    return bytes(dest_ptr)
+
+
+def _copy_from_bytes_to_cuda(dest_ptr: int, data: bytes) -> None:
+    # reserve space for 0 termination
+    scr_ptr = ctypes.create_string_buffer(data, len(data))
+    libcudart.cudaMemcpy(dest_ptr, scr_ptr, len(data))
+
+
+def _write_bytes(data: bytes, binary_file: io.BufferedWriter) -> None:
+    # Pack the length as a 4-byte unsigned integer (little-endian)
+    data_len = len(data)
+    header = struct.pack("<I", data_len)
+    binary_file.write(header)
+    if data_len > 0:
+        binary_file.write(data)
+
+
+def _read_bytes(mmap_obj: mmap.mmap) -> bytes:
+    header = mmap_obj.read(4)
+    if not header:
+        raise ValueError("Missing header read")
+
+    if len(header) != 4:
+        raise ValueError("Incomplete header read")
+
+    data_len = struct.unpack("<I", header)[0]
+    if data_len == 0:
+        return b''
+
+    data = mmap_obj.read(data_len)
+
+    if len(data) != data_len:
+        raise ValueError("Incomplete data read")
+
+    return data
+
+
 def find_loaded_library(lib_name) -> Optional[str]:
     """
     According to according to https://man7.org/linux/man-pages/man5/proc_pid_maps.5.html,
@@ -161,6 +209,8 @@ def __init__(self):
         self.python_malloc_callback = self._python_malloc_callback
         self.python_free_callback = self._python_free_callback
 
+        self.cache_filepath = ""
+
     def _python_malloc_callback(self, allocation_handle: HandleType) -> None:
         """
         Internal method to store the allocation data
@@ -185,8 +235,25 @@ def _python_free_callback(self, ptr: int) -> HandleType:
             data.handle[1], data.tag, ptr)
         return data.handle
 
+    def _delete_cache_file(self):
+        """
+        Remove sleep cache file if it exists
+        """
+        if self.cache_filepath != "":
+            filepath = self.cache_filepath
+            self.cache_filepath = ""
+            try:
+                with contextlib.suppress(FileNotFoundError):
+                    os.remove(filepath)
+                    logger.info("cache file %s deleted", filepath)
+            except Exception as e:
+                logger.warning("failed to delete sleep cache file %s.",
+                               filepath,
+                               exc_info=e)
+
     def sleep(
             self,
+            level: Optional[int] = 1,
             offload_tags: Optional[Union[tuple[str, ...],
                                          str]] = None) -> None:
         """
@@ -209,6 +276,29 @@ def sleep(
         total_bytes = 0
         backup_bytes = 0
 
+        # remove previous file if exists
+        self._delete_cache_file()
+
+        # level 3 write weights to file
+        if level == 3:
+            unique_id = uuid.uuid4().hex
+            self.cache_filepath = os.path.join(envs.VLLM_CACHE_ROOT,
+                                               f"sleep_cache_{unique_id}.bin")
+            logger.info("sleep level %d writing to cache file %s", level,
+                        self.cache_filepath)
+            with open(self.cache_filepath, 'wb') as binary_file:
+                for ptr, data in self.pointer_to_data.items():
+                    handle = data.handle
+                    if data.tag in offload_tags:
+                        size_in_bytes = handle[1]
+                        data = _copy_from_cuda_to_bytes(ptr, size_in_bytes)
+                        _write_bytes(data, binary_file)
+                    else:
+                        _write_bytes(b'', binary_file)
+                    unmap_and_release(handle)
+            return
+
+        # handle other levels
         for ptr, data in self.pointer_to_data.items():
             handle = data.handle
             total_bytes += handle[1]
@@ -244,6 +334,23 @@ def wake_up(self, tags: Optional[list[str]] = None) -> None:
             back to GPU memory. If None, all memory allocation will be loaded
             back to GPU memory.
         """
+        if self.cache_filepath != "":
+            logger.info("wake_up reading from cache file %s",
+                        self.cache_filepath)
+            with open(self.cache_filepath, 'rb') as bin_file, \
+                    mmap.mmap(bin_file.fileno(),
+                              length=0, access=mmap.ACCESS_READ) as mmap_obj:
+                for ptr, data in self.pointer_to_data.items():
+                    handle = data.handle
+                    create_and_map(handle)
+                    data = _read_bytes(mmap_obj)
+                    if len(data) > 0:
+                        _copy_from_bytes_to_cuda(ptr, data)
+
+            # remove file
+            self._delete_cache_file()
+            return
+
         for ptr, data in self.pointer_to_data.items():
             if tags is None or data.tag in tags:
                 handle = data.handle

vllm/entrypoints/llm.py

@@ -1433,6 +1433,11 @@ def sleep(self, level: int = 1):
                 sleep is good for sleeping and waking up the engine to run a
                 different model or update the model, where previous model
                 weights are not needed. It reduces CPU memory pressure.
+                Level 3 sleep will offload the model weights to disk and
+                discard the kv cache. The model weights are not backed up in
+                CPU memory. The content of kv cache is forgotten. Level 3
+                sleep helps use minimum CPU memory and loads efficiently
+                from disk when woken up.
         """
         self.reset_prefix_cache()
         self.llm_engine.sleep(level=level)

vllm/v1/worker/gpu_worker.py

@@ -109,7 +109,9 @@ def sleep(self, level: int = 1) -> None:
             }
 
         allocator = CuMemAllocator.get_instance()
-        allocator.sleep(offload_tags=("weights", ) if level == 1 else tuple())
+        allocator.sleep(level,
+                        offload_tags=("weights", )
+                        if level == 1 or level == 3 else tuple())
         free_bytes_after_sleep, total = torch.cuda.mem_get_info()
         freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
         used_bytes = total - free_bytes_after_sleep