torch-memory-saver 0.0.1__tar.gz

torch_memory_saver-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 fzyzcjy
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torch_memory_saver-0.0.1/PKG-INFO

@@ -0,0 +1,5 @@
+Metadata-Version: 2.1
+Name: torch_memory_saver
+Version: 0.0.1
+Requires-Python: >=3.9
+License-File: LICENSE

torch_memory_saver-0.0.1/README.md

@@ -0,0 +1,11 @@
+# torch_memory_saver
+
+Allow torch tensor memory to be released and resumed later 
+
+Please refer to https://github.com/sgl-project/sglang/issues/2542#issuecomment-2563641647 for details.
+
+TODO:
+
+- [x] Implementation
+- [ ] More tests and infra
+- [ ] Publish to pypi

torch_memory_saver-0.0.1/csrc/torch_memory_saver.cpp

@@ -0,0 +1,288 @@
+#include <sys/types.h>
+#include <cuda_runtime_api.h>
+#include <cuda.h>
+#include <iostream>
+#include <stdio.h>
+#include <dlfcn.h>
+#include <unordered_map>
+#include <mutex>
+
+// #define TMS_DEBUG_LOG
+
+// ----------------------------------------------- copied code --------------------------------------------------
+
+// Cannot use pytorch (libc10.so) since LD_PRELOAD happens earlier than `import torch`
+// #include <ATen/cuda/Exceptions.h>
+
+// torch Macros.h
+#if defined(__GNUC__) || defined(__ICL) || defined(__clang__)
+#define C10_LIKELY(expr) (__builtin_expect(static_cast<bool>(expr), 1))
+#define C10_UNLIKELY(expr) (__builtin_expect(static_cast<bool>(expr), 0))
+#else
+#define C10_LIKELY(expr) (expr)
+#define C10_UNLIKELY(expr) (expr)
+#endif
+
+// ----------------------------------------------- utils --------------------------------------------------
+
+#define SIMPLE_CHECK(COND, MSG) \
+  do { \
+    if (!(COND)) { \
+        std::cerr << "[torch_memory_saver.cpp] " << MSG << std::endl; \
+        exit(1); \
+    } \
+  } while (false)
+
+// very naive
+#define CURESULT_CHECK(EXPR) \
+  do { \
+    CUresult __result = (EXPR); \
+    if (__result != CUDA_SUCCESS) { \
+        std::cerr << "[torch_memory_saver.cpp] CUresult error " \
+            << " result=" << __result << " file=" << __FILE__ << " func=" << __func__ << " line=" << __LINE__ \
+            << std::endl; \
+        exit(1); \
+    } \
+  } while (false)
+
+namespace APIForwarder {
+    static void *check_dlsym(void *value) {
+        if (nullptr == value) {
+            std::cerr << "[torch_memory_saver.cpp] dlsym failed dlerror=" << dlerror() << std::endl;
+            exit(1);
+        }
+        return value;
+    }
+
+    typedef cudaError_t (*CudaMallocFunc)(void **, size_t);
+
+    typedef cudaError_t (*CudaFreeFunc)(void *);
+
+    static CudaMallocFunc real_cudaMalloc = NULL;
+    static CudaFreeFunc real_cudaFree = NULL;
+
+    static cudaError_t call_real_cuda_malloc(void **ptr, size_t size) {
+        if (C10_UNLIKELY(nullptr == real_cudaMalloc)) {
+            real_cudaMalloc = (CudaMallocFunc) check_dlsym(dlsym(RTLD_NEXT, "cudaMalloc"));
+        }
+
+        cudaError_t ret = real_cudaMalloc(ptr, size);
+
+#ifdef TMS_DEBUG_LOG
+        std::cout << "[torch_memory_saver.cpp] cudaMalloc [MODE NORMAL]"
+                  << " ptr=" << ptr << " *ptr=" << *ptr << " size=" << size << " ret=" << ret
+                  << std::endl;
+#endif
+
+        return ret;
+    }
+
+    static cudaError_t call_real_cuda_free(void *ptr) {
+        if (C10_UNLIKELY(nullptr == real_cudaFree)) {
+            real_cudaFree = (CudaFreeFunc) check_dlsym(dlsym(RTLD_NEXT, "cudaFree"));
+        }
+
+        cudaError_t ret = real_cudaFree(ptr);
+
+#ifdef TMS_DEBUG_LOG
+        std::cout << "[torch_memory_saver.cpp] cudaFree [MODE NORMAL]"
+                  << " ptr=" << ptr << " ret=" << ret
+                  << std::endl;
+#endif
+
+        return ret;
+    }
+}
+
+namespace CUDAUtils {
+    static void cu_mem_create(CUmemGenericAllocationHandle *allocHandle, size_t size, CUdevice device) {
+        CUmemAllocationProp prop = {};
+        prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
+        prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+        prop.location.id = device;
+        CURESULT_CHECK(cuMemCreate(allocHandle, size, &prop, 0));
+    }
+
+    static void cu_mem_set_access(void *ptr, size_t size, CUdevice device) {
+        CUmemAccessDesc accessDesc = {};
+        accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+        accessDesc.location.id = device;
+        accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+        CURESULT_CHECK(cuMemSetAccess((CUdeviceptr) ptr, size, &accessDesc, 1));
+    }
+}
+
+// ----------------------------------------------- primary class --------------------------------------------------
+
+struct _AllocationMetadata {
+    size_t size;
+    CUdevice device;
+    CUmemGenericAllocationHandle allocHandle;
+};
+
+class TorchMemorySaver {
+public:
+    TorchMemorySaver() {}
+
+    cudaError_t malloc(void **ptr, size_t size) {
+        CUdevice device;
+        CURESULT_CHECK(cuCtxGetDevice(&device));
+
+        CUmemGenericAllocationHandle allocHandle;
+        CUDAUtils::cu_mem_create(&allocHandle, size, device);
+
+        CURESULT_CHECK(cuMemAddressReserve((CUdeviceptr *) ptr, size, 0, 0, 0));
+        CURESULT_CHECK(cuMemMap((CUdeviceptr) * ptr, size, 0, allocHandle, 0));
+        CUDAUtils::cu_mem_set_access(*ptr, size, device);
+
+        {
+            const std::lock_guard <std::mutex> lock(allocator_metadata_mutex_);
+            allocation_metadata_.emplace(*ptr, _AllocationMetadata{size, device, allocHandle});
+        }
+
+#ifdef TMS_DEBUG_LOG
+        std::cout << "[torch_memory_saver.cpp] TorchMemorySaver.cuda_malloc "
+                  << " ptr=" << ptr << " *ptr=" << *ptr << " size=" << size
+                  << " allocHandle=" << allocHandle
+                  << std::endl;
+#endif
+
+        return cudaSuccess;
+    }
+
+    cudaError_t free(void *ptr) {
+        _AllocationMetadata metadata;
+        {
+            const std::lock_guard <std::mutex> lock(allocator_metadata_mutex_);
+            SIMPLE_CHECK(allocation_metadata_.count(ptr), "Trying to free a pointer not allocated here");
+            metadata = allocation_metadata_[ptr];
+            allocation_metadata_.erase(ptr);
+        }
+
+        CURESULT_CHECK(cuMemUnmap((CUdeviceptr) ptr, metadata.size));
+        CURESULT_CHECK(cuMemRelease(metadata.allocHandle));
+        CURESULT_CHECK(cuMemAddressFree((CUdeviceptr) ptr, metadata.size));
+
+#ifdef TMS_DEBUG_LOG
+        std::cout << "[torch_memory_saver.cpp] TorchMemorySaver.cuda_free "
+                  << " ptr=" << ptr << " metadata.size=" << metadata.size
+                  << " metadata.allocHandle=" << metadata.allocHandle
+                  << std::endl;
+#endif
+
+        return cudaSuccess;
+    }
+
+    void pause() {
+        const std::lock_guard <std::mutex> lock(allocator_metadata_mutex_);
+
+        for (auto it = allocation_metadata_.begin(); it != allocation_metadata_.end(); ++it) {
+            void *ptr = it->first;
+            _AllocationMetadata metadata = it->second;
+
+            CURESULT_CHECK(cuMemUnmap((CUdeviceptr) ptr, metadata.size));
+            CURESULT_CHECK(cuMemRelease(metadata.allocHandle));
+
+#ifdef TMS_DEBUG_LOG
+            std::cout << "[torch_memory_saver.cpp] TorchMemorySaver.pause"
+                      << " ptr=" << ptr << " metadata.size=" << metadata.size << " metadata.allocHandle="
+                      << metadata.allocHandle
+                      << std::endl;
+#endif
+        }
+    }
+
+    void resume() {
+        const std::lock_guard <std::mutex> lock(allocator_metadata_mutex_);
+
+        for (auto it = allocation_metadata_.begin(); it != allocation_metadata_.end(); ++it) {
+            void *ptr = it->first;
+            _AllocationMetadata &metadata = it->second;
+
+            CUmemGenericAllocationHandle newAllocHandle;
+            CUDAUtils::cu_mem_create(&newAllocHandle, metadata.size, metadata.device);
+
+            CURESULT_CHECK(cuMemMap((CUdeviceptr) ptr, metadata.size, 0, newAllocHandle, 0));
+
+            CUDAUtils::cu_mem_set_access(ptr, metadata.size, metadata.device);
+
+#ifdef TMS_DEBUG_LOG
+            std::cout << "[torch_memory_saver.cpp] TorchMemorySaver.resume"
+                      << " ptr=" << ptr << " metadata.size=" << metadata.size << " (old)metadata.allocHandle="
+                      << metadata.allocHandle
+                      << " (new)newAllocHandle=" << newAllocHandle
+                      << std::endl;
+#endif
+
+            metadata.allocHandle = newAllocHandle;
+        }
+    }
+
+    static TorchMemorySaver &instance() {
+        static TorchMemorySaver instance;
+        return instance;
+    }
+
+private:
+    // Similar to torch's CUDACachingAllocator and CUDAPluggableAllocator
+    std::mutex allocator_metadata_mutex_;
+    std::unordered_map<void *, _AllocationMetadata> allocation_metadata_;
+};
+
+namespace RegionManager {
+    static thread_local bool is_interesting_region_ = false;
+
+    void enter() {
+#ifdef TMS_DEBUG_LOG
+        std::cout << "[torch_memory_saver.cpp] tms_region_enter" << std::endl;
+#endif
+        is_interesting_region_ = true;
+    }
+
+    void leave() {
+#ifdef TMS_DEBUG_LOG
+        std::cout << "[torch_memory_saver.cpp] tms_region_leave" << std::endl;
+#endif
+        is_interesting_region_ = false;
+    }
+
+    bool is_interesting_region() {
+        return is_interesting_region_;
+    }
+}
+
+// ------------------------------------------------- entrypoints ------------------------------------------------
+
+cudaError_t cudaMalloc(void **ptr, size_t size) {
+    if (RegionManager::is_interesting_region()) {
+        return TorchMemorySaver::instance().malloc(ptr, size);
+    } else {
+        return APIForwarder::call_real_cuda_malloc(ptr, size);
+    }
+}
+
+cudaError_t cudaFree(void *ptr) {
+    if (RegionManager::is_interesting_region()) {
+        return TorchMemorySaver::instance().free(ptr);
+    } else {
+        return APIForwarder::call_real_cuda_free(ptr);
+    }
+}
+
+extern "C" {
+void tms_region_enter() {
+    RegionManager::enter();
+}
+
+void tms_region_leave() {
+    RegionManager::leave();
+}
+
+void tms_pause() {
+    TorchMemorySaver::instance().pause();
+}
+
+void tms_resume() {
+    TorchMemorySaver::instance().resume();
+}
+}

torch_memory_saver-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torch_memory_saver-0.0.1/setup.py

@@ -0,0 +1,19 @@
+from setuptools import setup
+from torch.utils import cpp_extension
+
+ext_module = cpp_extension.CppExtension(
+    'torch_memory_saver_cpp',
+    ['csrc/torch_memory_saver.cpp'],
+    extra_compile_args=['-I/usr/local/cuda/include'],
+    extra_link_args=['-lcuda'],
+)
+
+setup(
+    name='torch_memory_saver',
+    version='0.0.1',
+    # https://pytorch.org/tutorials/advanced/cpp_extension.html#writing-a-c-extension
+    ext_modules=[ext_module],
+    cmdclass={'build_ext': cpp_extension.BuildExtension},
+    python_requires=">=3.9",
+    packages=['torch_memory_saver'],
+)

torch_memory_saver-0.0.1/torch_memory_saver/__init__.py

@@ -0,0 +1,88 @@
+import ctypes
+import logging
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class TorchMemorySaver:
+    def __init__(self):
+        self._mem_pool = torch.cuda.MemPool()
+        self._id = _global_info.next_id()
+        assert self._id == 1, 'Only support one single instance yet (multi-instance will be implemented later)'
+
+    @contextmanager
+    def region(self):
+        with torch.cuda.use_mem_pool(self._mem_pool):
+            _global_info.cdll.tms_region_enter()
+            try:
+                yield
+            finally:
+                _global_info.cdll.tms_region_leave()
+
+    def pause(self):
+        _global_info.cdll.tms_pause()
+
+    def resume(self):
+        _global_info.cdll.tms_resume()
+
+
+class _GlobalInfo:
+    def __init__(self):
+        self._cdll: Optional[ctypes.CDLL] = None
+        self._last_id = 0
+
+    @property
+    def cdll(self):
+        if self._cdll is None:
+            self._cdll = _compute_cdll()
+            logger.debug(f'Use cdll={self._cdll}')
+        return self._cdll
+
+    def next_id(self):
+        self._last_id += 1
+        return self._last_id
+
+
+_global_info = _GlobalInfo()
+
+
+def _compute_cdll():
+    env_ld_preload = os.environ.get('LD_PRELOAD', '')
+    assert 'torch_memory_saver' in env_ld_preload, f'Please specify correct LD_PRELOAD (currently: {env_ld_preload})'
+    return ctypes.CDLL(env_ld_preload)
+
+
+def get_binary_path():
+    dir_package = Path(__file__).parent
+    candidates = [
+        p
+        for d in [dir_package, dir_package.parent]
+        for p in d.glob('torch_memory_saver_cpp.*.so')
+    ]
+    assert len(candidates) == 1, f'{candidates=}'
+    return candidates[0]
+
+
+@contextmanager
+def configure_subprocess():
+    with change_env('LD_PRELOAD', str(get_binary_path())):
+        yield
+
+
+@contextmanager
+def change_env(key: str, value: str):
+    old_value = os.environ.get(key, '')
+    os.environ[key] = value
+    logger.debug(f'change_env set key={key} value={value}')
+    try:
+        yield
+    finally:
+        assert os.environ[key] == value
+        os.environ[key] = old_value
+        logger.debug(f'change_env restore key={key} value={old_value}')

torch_memory_saver-0.0.1/torch_memory_saver.egg-info/PKG-INFO

@@ -0,0 +1,5 @@
+Metadata-Version: 2.1
+Name: torch_memory_saver
+Version: 0.0.1
+Requires-Python: >=3.9
+License-File: LICENSE

torch_memory_saver-0.0.1/torch_memory_saver.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+LICENSE
+README.md
+setup.py
+csrc/torch_memory_saver.cpp
+torch_memory_saver/__init__.py
+torch_memory_saver.egg-info/PKG-INFO
+torch_memory_saver.egg-info/SOURCES.txt
+torch_memory_saver.egg-info/dependency_links.txt
+torch_memory_saver.egg-info/top_level.txt

torch_memory_saver-0.0.1/torch_memory_saver.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

torch_memory_saver-0.0.1/torch_memory_saver.egg-info/top_level.txt

@@ -0,0 +1,2 @@
+torch_memory_saver
+torch_memory_saver_cpp