torch-memory-saver 0.0.6__tar.gz → 0.0.9rc1__tar.gz

torch_memory_saver-0.0.9rc1/PKG-INFO

@@ -0,0 +1,7 @@
+Metadata-Version: 2.4
+Name: torch_memory_saver
+Version: 0.0.9rc1
+Requires-Python: >=3.9
+License-File: LICENSE
+Dynamic: license-file
+Dynamic: requires-python

torch_memory_saver-0.0.9rc1/README.md

@@ -0,0 +1,104 @@
+# Torch Memory Saver
+
+A PyTorch library that allows tensor memory to be temporarily released and resumed later.
+
+Please refer to https://github.com/sgl-project/sglang/issues/2542#issuecomment-2563641647 for details.
+
+## Examples and Features
+
+### Basic Example
+
+```python
+# 1. For tensors that wants to be paused, create them within `region`
+with torch_memory_saver.region():
+    pauseable_tensor = torch.full((1_000_000_000,), 100, dtype=torch.uint8, device='cuda')
+
+# 2. After `pause`, CUDA memory is released for those tensors.
+# For example, check `nvidia-smi`'s memory usage to verify.
+torch_memory_saver.pause()
+
+# 3. After `resume`, CUDA memory is re-occupied for those tensors.
+torch_memory_saver.resume()
+```
+
+During the pause, physical memory is released and virtual address is preserved. When resume, virtual address is kept unchanged, while physical memory is re-allocated
+
+### Multiple Tags
+
+Please refer to https://github.com/sgl-project/sglang/issues/7009 for details.
+
+```python
+# 1. Create tensors with different tags
+with torch_memory_saver.region(tag="type1"):
+    tensor1 = torch.full((5_000_000_000,), 100, dtype=torch.uint8, device='cuda')
+
+with torch_memory_saver.region(tag="type2"):
+    tensor2 = torch.full((5_000_000_000,), 100, dtype=torch.uint8, device='cuda')
+
+# 2. Pause and resume with different tags selectively
+torch_memory_saver.pause("type1")
+torch_memory_saver.pause("type2")
+
+torch_memory_saver.resume("type2")
+torch_memory_saver.resume("type1")
+
+torch_memory_saver.pause("type1")
+torch_memory_saver.resume("type1")
+```
+
+### Release Memory in CUDA Graph
+
+Not only does torch_memory_saver make tensors compatible with CUDA graph, but we can also release the memory held by CUDA graph (i.e. the intermediate tensors).
+
+API: Change `torch.cuda.graph(...)` to `torch_memory_saver.cuda_graph(...)`
+
+### CPU Backup
+
+By default, in order to save time, the content is thrown away. This is useful for, for example, KV cache that are to be staled, or model weights that are to be updated.
+
+If you want the tensor content to be kept unchanged, use `enable_cpu_backup`.
+
+```python
+with torch_memory_saver.region(enable_cpu_backup=True):
+    tensor1 = torch.full((5_000_000_000,), 42, dtype=torch.uint8, device='cuda')
+
+torch_memory_saver.pause()
+torch_memory_saver.resume()
+
+assert tensor1[0] == 42, "content is kept unchanged"
+```
+
+### Hook Modes
+
+There are two hook modes:
+
+* **preload**: Use `LD_PRELOAD` to hook CUDA's malloc and free API to change allocation behavior.
+* **torch**: Use torch's custom allocator API to change allocation behavior.
+
+The mode can be chosen by:
+
+```python
+torch_memory_saver.hook_mode = "torch"
+```
+
+### Example of RL with CUDA Graph
+
+Please refer to `rl_example.py` for details.
+
+## Development
+
+```bash
+make reinstall
+```
+
+You can use this command for local testing:
+
+```bash
+pytest /path/to/torch_memory_saver/test
+```
+
+Or this one to test a single case (e.g. the `simple` one here):
+
+```bash
+pytest /path/to/torch_memory_saver/test/test_examples.py::test_simple -s
+```

torch_memory_saver-0.0.9rc1/csrc/api_forwarder.cpp

@@ -0,0 +1,52 @@
+#include <iostream>
+#include "api_forwarder.h"
+#include "utils.h"
+#include "macro.h"
+
+namespace APIForwarder {
+    using CudaMallocFunc = cudaError_t (*)(void**, size_t);
+    using CudaFreeFunc = cudaError_t (*)(void*);
+
+    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;
+    }
+
+    static CudaMallocFunc real_cuda_malloc_ = NULL;
+    static CudaFreeFunc real_cuda_free_ = NULL;
+
+    cudaError_t call_real_cuda_malloc(void **ptr, size_t size) {
+        if (C10_UNLIKELY(nullptr == real_cuda_malloc_)) {
+            real_cuda_malloc_ = (CudaMallocFunc) check_dlsym(dlsym(RTLD_NEXT, "cudaMalloc"));
+        }
+
+        cudaError_t ret = real_cuda_malloc_(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;
+    }
+
+    cudaError_t call_real_cuda_free(void *ptr) {
+        if (C10_UNLIKELY(nullptr == real_cuda_free_)) {
+            real_cuda_free_ = (CudaFreeFunc) check_dlsym(dlsym(RTLD_NEXT, "cudaFree"));
+        }
+
+        cudaError_t ret = real_cuda_free_(ptr);
+
+#ifdef TMS_DEBUG_LOG
+        std::cout << "[torch_memory_saver.cpp] cudaFree [MODE NORMAL]"
+                  << " ptr=" << ptr << " ret=" << ret
+                  << std::endl;
+#endif
+
+        return ret;
+    }
+}

torch_memory_saver-0.0.9rc1/csrc/core.cpp

@@ -0,0 +1,332 @@
+#include "core.h"
+#include "utils.h"
+#include "macro.h"
+#include "api_forwarder.h"
+
+TorchMemorySaver::TorchMemorySaver() {}
+
+TorchMemorySaver &TorchMemorySaver::instance() {
+    static TorchMemorySaver instance;
+    return instance;
+}
+
+cudaError_t TorchMemorySaver::malloc(void **ptr, CUdevice device, size_t size, const std::string& tag, const bool enable_cpu_backup) {
+#if defined(USE_ROCM)
+    // hipDevice_t device;
+    CURESULT_CHECK(hipCtxGetDevice(&device));
+
+    // // Get granularity and calculate aligned size
+    // size_t granularity = CUDAUtils::cu_mem_get_granularity(device);
+    // size_t aligned_size = (size + granularity - 1) & ~(granularity - 1);
+
+    // //// Reserve aligned memory address, rocm will check granularity
+    // CURESULT_CHECK(hipMemAddressReserve((hipDeviceptr_t *)ptr, aligned_size, granularity, 0, 0));
+
+    hipMemAllocationProp prop = {};
+    prop.type = hipMemAllocationTypePinned;
+    prop.location.type = hipMemLocationTypeDevice;
+    prop.location.id = device;
+    prop.allocFlags.compressionType = 0x0;
+
+    size_t granularity;
+    CURESULT_CHECK(hipMemGetAllocationGranularity(&granularity, &prop,
+                                            hipMemAllocationGranularityMinimum));
+    size_t aligned_size = ((size + granularity - 1) / granularity) * granularity;
+    aligned_size = (aligned_size + MEMCREATE_CHUNK_SIZE - 1) / MEMCREATE_CHUNK_SIZE * MEMCREATE_CHUNK_SIZE;
+
+    assert(MEMCREATE_CHUNK_SIZE % granularity == 0);
+    assert(aligned_size % MEMCREATE_CHUNK_SIZE == 0);
+    assert(aligned_size % granularity == 0);
+
+
+    // Create allocation metadata
+    AllocationMetadata metadata;
+    metadata.size = size;
+    metadata.aligned_size = aligned_size;
+    metadata.device = device;
+    //// Not sure (Check these parameters)
+    metadata.tag = tag;
+    metadata.enable_cpu_backup = enable_cpu_backup;
+    metadata.cpu_backup = nullptr;
+    ////
+
+    // Get global device ID using our utility function
+    int global_device_id = DeviceUtils::get_global_device_id(device);
+
+    // rewrite numa node 
+    uint64_t node_id = 0;
+    if (global_device_id > 3) {
+        node_id = 1;
+    }
+
+#ifdef TMS_DEBUG_LOG
+    std::cout << "[torch_memory_saver.cpp] TorchMemorySaver.cuda_malloc "
+              << " ptr=" << ptr << " *ptr=" << *ptr << " size=" << size
+              << " granularity=" << granularity
+              << " aligned_size=" << aligned_size
+              << " node_id=" << node_id
+              << " device=" << device
+              << " global_device_id=" << global_device_id
+              << std::endl;
+#endif
+
+    hipDeviceptr_t d_mem;
+    // Reserve aligned memory address, rocm will check granularity
+    CURESULT_CHECK(hipMemAddressReserve(&d_mem, aligned_size, granularity, 0, node_id));
+    *ptr = (void*)d_mem;
+
+    // Create and map chunks
+    // CUDAUtils::cu_mem_create_and_map(device, size, (hipDeviceptr_t)*ptr, 
+    CUDAUtils::cu_mem_create_and_map(device, aligned_size, (hipDeviceptr_t)*ptr, 
+                                    metadata.allocHandles, metadata.chunk_sizes);
+    size_t num_chunks = metadata.allocHandles.size();
+    {
+        const std::lock_guard<std::mutex> lock(allocator_metadata_mutex_);
+        allocation_metadata_.emplace(*ptr, std::move(metadata));
+    }
+#ifdef TMS_DEBUG_LOG
+    std::cout << "[torch_memory_saver.cpp] TorchMemorySaver.cuda_malloc "
+              << " ptr=" << ptr << " *ptr=" << *ptr << " size=" << size
+              << " metadata.aligned_size=" << metadata.aligned_size
+              << " num_chunks=" << num_chunks
+              << std::endl;
+#endif
+
+#elif defined(USE_CUDA)
+    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, tag, AllocationState::ACTIVE, enable_cpu_backup, nullptr, allocHandle}
+        );
+    }
+
+#ifdef TMS_DEBUG_LOG
+    std::cout << "[torch_memory_saver.cpp] TorchMemorySaver.cuda_malloc "
+              << " ptr=" << ptr << " *ptr=" << *ptr << " size=" << size
+              << " allocHandle=" << allocHandle << " tag=" << tag
+              << std::endl;
+#endif
+
+#else
+    #error "USE_PLATFORM is not set"
+#endif
+    return cudaSuccess;
+}
+
+cudaError_t TorchMemorySaver::free(void *ptr) {
+#if defined(USE_ROCM)
+    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 = std::move(allocation_metadata_[ptr]);
+        allocation_metadata_.erase(ptr);
+    }
+
+    // Unmap and release chunks
+    CUDAUtils::cu_mem_unmap_and_release(metadata.device, metadata.size, 
+                                        (hipDeviceptr_t)ptr, metadata.allocHandles, metadata.chunk_sizes);
+
+    // Free the reserved address using stored aligned_size
+    CURESULT_CHECK(hipMemAddressFree((hipDeviceptr_t)ptr, metadata.aligned_size));
+
+#ifdef TMS_DEBUG_LOG
+    std::cout << "[torch_memory_saver.cpp] TorchMemorySaver.cuda_free "
+              << " ptr=" << ptr << " metadata.size=" << metadata.size
+              << " metadata.aligned_size=" << metadata.aligned_size
+              << " num_chunks=" << metadata.allocHandles.size()
+              << std::endl;
+#endif
+#elif defined(USE_CUDA)
+    AllocationMetadata metadata;
+    {
+        const std::lock_guard <std::mutex> lock(allocator_metadata_mutex_);
+        if (allocation_metadata_.count(ptr) == 0) {
+            return APIForwarder::call_real_cuda_free(ptr);
+        }
+
+        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));
+
+    if (nullptr != metadata.cpu_backup) {
+        CUDA_ERROR_CHECK(cudaFreeHost(metadata.cpu_backup));
+        metadata.cpu_backup = nullptr;
+    }
+
+#ifdef TMS_DEBUG_LOG
+    std::cout << "[torch_memory_saver.cpp] TorchMemorySaver.cuda_free "
+              << " ptr=" << ptr << " metadata.size=" << metadata.size
+              << " metadata.allocHandle=" << metadata.allocHandle << " tag=" << metadata.tag
+              << std::endl;
+#endif
+
+#else
+    #error "USE_PLATFORM is not set"
+#endif
+    return cudaSuccess;
+}
+
+void TorchMemorySaver::pause(const std::string& tag) {
+    const std::lock_guard <std::mutex> lock(allocator_metadata_mutex_);
+
+#if defined(USE_ROCM)
+    for (auto it = allocation_metadata_.begin(); it != allocation_metadata_.end(); ++it) {
+        void *ptr = it->first;
+        AllocationMetadata &metadata = it->second;
+
+        if (!tag.empty() && metadata.tag != tag) {
+            continue;
+        }
+        // Copy CUDA's code supporting cpu_backup to here
+        if (metadata.enable_cpu_backup) {
+            if (nullptr == metadata.cpu_backup) {
+                CUDA_ERROR_CHECK(hipMallocHost(&metadata.cpu_backup, metadata.aligned_size));
+            }
+            SIMPLE_CHECK(metadata.cpu_backup != nullptr, "cpu_backup should not be nullptr");
+            // TODO may use cudaMemcpyAsync if needed
+            CUDA_ERROR_CHECK(cudaMemcpy(metadata.cpu_backup, ptr, metadata.aligned_size, hipMemcpyDeviceToHost));
+        }
+        //
+
+        // Unmap and release chunks (but keep metadata for resume)
+        // CUDAUtils::cu_mem_unmap_and_release(metadata.device, metadata.size,
+        CUDAUtils::cu_mem_unmap_and_release(metadata.device, metadata.aligned_size,
+                                            (hipDeviceptr_t)ptr, metadata.allocHandles, metadata.chunk_sizes);
+
+        #ifdef TMS_DEBUG_LOG
+            std::cout << "[torch_memory_saver.cpp] TorchMemorySaver.pause"
+                    << " ptr=" << ptr << " metadata.size=" << metadata.size 
+                    << " metadata.aligned_size=" << metadata.aligned_size
+                    << " num_chunks=" << metadata.allocHandles.size()
+                    << std::endl;
+        #endif
+    }
+#elif defined(USE_CUDA)
+    for (auto it = allocation_metadata_.begin(); it != allocation_metadata_.end(); ++it) {
+        void *ptr = it->first;
+        AllocationMetadata& metadata = it->second;
+
+        if (!tag.empty() && metadata.tag != tag) {
+            continue;
+        }
+
+        if (metadata.state != AllocationState::ACTIVE) {
+            std::cerr << "[torch_memory_saver.cpp] Cannot pause allocation that is not active."
+                      << " tag=" << metadata.tag << " ptr=" << std::to_string((uintptr_t)ptr)
+                      << " file=" << __FILE__ << " func=" << __func__ << " line=" << __LINE__
+                      << std::endl;
+            exit(1);
+        }
+
+        if (metadata.enable_cpu_backup) {
+            if (nullptr == metadata.cpu_backup) {
+                CUDA_ERROR_CHECK(cudaMallocHost(&metadata.cpu_backup, metadata.size));
+            }
+            SIMPLE_CHECK(metadata.cpu_backup != nullptr, "cpu_backup should not be nullptr");
+            // TODO may use cudaMemcpyAsync if needed
+            CUDA_ERROR_CHECK(cudaMemcpy(metadata.cpu_backup, ptr, metadata.size, cudaMemcpyDeviceToHost));
+        }
+
+        CURESULT_CHECK(cuMemUnmap((CUdeviceptr) ptr, metadata.size));
+        CURESULT_CHECK(cuMemRelease(metadata.allocHandle));
+
+        metadata.state = AllocationState::PAUSED;
+
+#ifdef TMS_DEBUG_LOG
+        std::cout << "[torch_memory_saver.cpp] TorchMemorySaver.pause"
+                  << " ptr=" << ptr << " metadata.size=" << metadata.size << " metadata.allocHandle="
+                  << metadata.allocHandle << " tag=" << metadata.tag << " filter_tag=" << tag
+                  << " metadata.enable_cpu_backup=" << metadata.enable_cpu_backup
+                  << std::endl;
+#endif
+    }
+#else
+    #error "USE_PLATFORM is not set"
+#endif
+}
+
+void TorchMemorySaver::resume(const std::string& tag) {
+    const std::lock_guard <std::mutex> lock(allocator_metadata_mutex_);
+
+#if defined(USE_ROCM)
+    for (auto it = allocation_metadata_.begin(); it != allocation_metadata_.end(); ++it) {
+        void *ptr = it->first;
+        AllocationMetadata &metadata = it->second;
+
+        if (!tag.empty() && metadata.tag != tag) {
+            continue;
+        }
+
+        // Create new handles and map chunks
+        // CUDAUtils::cu_mem_create_and_map(metadata.device, metadata.size,
+        CUDAUtils::cu_mem_create_and_map(metadata.device, metadata.aligned_size,
+                                        (hipDeviceptr_t)ptr, metadata.allocHandles, metadata.chunk_sizes);
+
+#ifdef TMS_DEBUG_LOG
+        std::cout << "[torch_memory_saver.cpp] TorchMemorySaver.resume"
+                << " ptr=" << ptr << " metadata.size=" << metadata.size
+                << " metadata.aligned_size=" << metadata.aligned_size
+                << " num_chunks=" << metadata.allocHandles.size()
+                << std::endl;
+#endif
+    }
+
+#elif defined(USE_CUDA)
+    for (auto it = allocation_metadata_.begin(); it != allocation_metadata_.end(); ++it) {
+        void *ptr = it->first;
+        AllocationMetadata &metadata = it->second;
+
+        if (!tag.empty() && metadata.tag != tag) {
+            continue;
+        }
+
+        if (metadata.state != AllocationState::PAUSED) {
+            std::cerr << "[torch_memory_saver.cpp] Cannot resume allocation that is not paused. "
+                      << " tag=" << metadata.tag << " ptr=" << std::to_string((uintptr_t)ptr)
+                      << " file=" << __FILE__ << " func=" << __func__ << " line=" << __LINE__
+                      << std::endl;
+            exit(1);
+        }
+
+        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);
+
+        if (metadata.enable_cpu_backup) {
+            SIMPLE_CHECK(metadata.cpu_backup != nullptr, "cpu_backup should not be nullptr");
+            // TODO may use cudaMemcpyAsync if needed
+            CUDA_ERROR_CHECK(cudaMemcpy(ptr, metadata.cpu_backup, metadata.size, cudaMemcpyHostToDevice));
+            // maybe we can free host memory if needed (currently keep it there to reduce re-alloc time)
+        }
+
+#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 << " tag=" << metadata.tag << " filter_tag=" << tag
+                  << " metadata.enable_cpu_backup=" << metadata.enable_cpu_backup
+                  << std::endl;
+#endif
+
+        metadata.state = AllocationState::ACTIVE;
+        metadata.allocHandle = newAllocHandle;
+    }
+#else
+    #error "USE_PLATFORM is not set"
+#endif
+}

torch_memory_saver-0.0.9rc1/csrc/entrypoint.cpp

@@ -0,0 +1,122 @@
+#include "utils.h"
+#include "core.h"
+#include "api_forwarder.h"
+#include <optional>
+#include "macro.h"
+
+// ----------------------------------------------- threadlocal configs --------------------------------------------------
+
+class ThreadLocalConfig {
+public:
+    std::string current_tag_ = "default";
+
+    bool is_interesting_region() {
+        if (!is_interesting_region_.has_value()) {
+            is_interesting_region_ = get_bool_env_var("TMS_INIT_ENABLE");
+        }
+        return is_interesting_region_.value();
+    }
+
+    void set_interesting_region(bool value) {
+        is_interesting_region_ = value;
+    }
+
+    bool enable_cpu_backup() {
+        if (!enable_cpu_backup_.has_value()) {
+            enable_cpu_backup_ = get_bool_env_var("TMS_INIT_ENABLE_CPU_BACKUP");
+        }
+        return enable_cpu_backup_.value();
+    }
+
+    void set_enable_cpu_backup(bool value) {
+        enable_cpu_backup_ = value;
+    }
+
+private:
+    std::optional<bool> is_interesting_region_;
+    std::optional<bool> enable_cpu_backup_;
+};
+static thread_local ThreadLocalConfig thread_local_config;
+
+// ------------------------------------------------- entrypoints :: hook ------------------------------------------------
+
+#ifdef TMS_HOOK_MODE_PRELOAD
+cudaError_t cudaMalloc(void **ptr, size_t size) {
+    if (thread_local_config.is_interesting_region()) {
+        return TorchMemorySaver::instance().malloc(
+            ptr, CUDAUtils::cu_ctx_get_device(), size, thread_local_config.current_tag_, thread_local_config.enable_cpu_backup());
+    } else {
+        return APIForwarder::call_real_cuda_malloc(ptr, size);
+    }
+}
+
+cudaError_t cudaFree(void *ptr) {
+    if (thread_local_config.is_interesting_region()) {
+        return TorchMemorySaver::instance().free(ptr);
+    } else {
+        return APIForwarder::call_real_cuda_free(ptr);
+    }
+}
+#endif
+
+#ifdef TMS_HOOK_MODE_TORCH
+extern "C" {
+void *tms_torch_malloc(ssize_t size, int device, cudaStream_t stream) {
+#ifdef TMS_DEBUG_LOG
+    std::cout << "[torch_memory_saver.cpp] tms_torch_malloc "
+              << " size=" << size << " device=" << device << " stream=" << stream
+              << std::endl;
+#endif
+    SIMPLE_CHECK(thread_local_config.is_interesting_region(), "only support interesting region");
+    void *ptr;
+    TorchMemorySaver::instance().malloc(
+        &ptr, CUDAUtils::cu_device_get(device), size, thread_local_config.current_tag_, thread_local_config.enable_cpu_backup());
+    return ptr;
+}
+
+void tms_torch_free(void *ptr, ssize_t ssize, int device, cudaStream_t stream) {
+#ifdef TMS_DEBUG_LOG
+    std::cout << "[torch_memory_saver.cpp] tms_torch_free "
+              << " ptr=" << ptr << " ssize=" << ssize << " device=" << device << " stream=" << stream
+              << std::endl;
+#endif
+    SIMPLE_CHECK(thread_local_config.is_interesting_region(), "only support interesting region");
+    TorchMemorySaver::instance().free(ptr);
+}
+}
+#endif
+
+// ------------------------------------------------- entrypoints :: others ------------------------------------------------
+
+extern "C" {
+void tms_set_interesting_region(bool is_interesting_region) {
+    thread_local_config.set_interesting_region(is_interesting_region);
+}
+
+bool tms_get_interesting_region() {
+    return thread_local_config.is_interesting_region();
+}
+
+void tms_set_current_tag(const char* tag) {
+    SIMPLE_CHECK(tag != nullptr, "tag should not be null");
+    thread_local_config.current_tag_ = tag;
+}
+
+bool tms_get_enable_cpu_backup() {
+    return thread_local_config.enable_cpu_backup();
+}
+
+void tms_set_enable_cpu_backup(bool enable_cpu_backup) {
+    thread_local_config.set_enable_cpu_backup(enable_cpu_backup);
+}
+
+void tms_pause(const char* tag) {
+    std::string tag_str = (tag != nullptr) ? std::string(tag) : "";
+    TorchMemorySaver::instance().pause(tag_str);
+}
+
+void tms_resume(const char* tag) {
+    std::string tag_str = (tag != nullptr) ? std::string(tag) : "";
+    TorchMemorySaver::instance().resume(tag_str);
+}
+}