Torch Memory Saver

A PyTorch library that allows tensor memory to be temporarily released and resumed later.

Please refer to sgl-project/sglang#2542 (comment) for details.

Examples and Features

Basic Example

# 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 sgl-project/sglang#7009 for details.

# 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(...)

If a retained CUDA graph contains NCCL operations on TMS-managed buffers, NCCL graph registrations can become stale when TMS replaces the buffers' physical backing during pause() / resume(), and graph replay may hang (#88). Either keep buffers passed to NCCL operations outside TMS-managed regions, or disable NCCL graph user-buffer registration before starting the workload:

export NCCL_GRAPH_REGISTER=0

This leaves NCCL and CUDA graphs enabled, but may give up the associated zero-copy optimization from NCCL CUDA-graph user-buffer registration.

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.

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"

The default host shadow is pinned (cpu_backup_backend="pinned"). If process RSS remains resident after a non-retaining resume with the pinned backend, CUDA users can select cpu_backup_backend="mmap" (or TMS_INIT_CPU_BACKUP_BACKEND=mmap) so releasing the host shadow uses munmap instead of cudaFreeHost. Both backends release the host shadow on resume unless retention is enabled.

with torch_memory_saver.region(enable_cpu_backup=True, cpu_backup_backend="mmap"):
    ...

ROCm stays pinned (hipHostMalloc); cpu_backup_backend="mmap" is rejected. On the legacy ROCm path, host shadows are retained across resume. XPU host shadows are pageable malloc/free; cpu_backup_backend is not supported there.

TMS_INIT_CPU_BACKUP_BACKEND=mmap|pinned sets the process default for preload / env-driven integrations; an explicit cpu_backup_backend= argument overrides it.

On CUDA, pinned or mmap CPU backups can be retained across resume cycles to avoid reallocating them before every pause:

torch_memory_saver.retain_cpu_backup = True

Set TMS_RETAIN_CPU_BACKUP=1 before process startup to enable the same policy, including for preload mode. Retention is CUDA-only and can consume host RAM equal to the backed-up allocations. While enabled, get_cpu_backup continues to expose the retained backup after resume. Without retention on CUDA, get_cpu_backup is only valid while allocations are paused.

The retention policy is consulted on resume. Setting it to False does not eagerly free backups for active allocations; they are released by a later non-retaining resume, when the allocation is freed, or when the process exits.

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:

torch_memory_saver.hook_mode = "torch"

Example of RL with CUDA Graph

Please refer to rl_example.py for details.

Platform Support

Platform	Backend	Hook modes	Install
NVIDIA (CUDA)	CUDA VMM (cuMemMap/cuMemCreate)	preload, torch	pip install torch_memory_saver (prebuilt wheel)
AMD (ROCm)	HIP VMM (hipMemMap/hipMemCreate)	preload, torch	pip install torch_memory_saver (prebuilt wheel)
Intel (XPU)	Level Zero VMM (zeVirtualMemMap/zePhysicalMemCreate)	torch only	pip install from source (builds against local oneAPI)

Intel XPU

Same pause/resume behavior on Intel GPUs via Level Zero VMM, wired into PyTorch through XPUPluggableAllocator + torch.xpu.MemPool. Only hook_mode="torch" is supported (preload is CUDA/HIP-specific); pauseable CUDA-graph capture is not.

Install from source (no prebuilt wheel). TMS_PLATFORM=xpu forces the backend; --no-build-isolation lets the build see torch+xpu and ABI-match libsycl to it (otherwise a mismatched oneAPI silently produces a broken .so):

# Prerequisites: torch+xpu, Intel oneAPI (icpx) + Level Zero headers.
source /opt/intel/oneapi/setvars.sh          # put icpx on PATH (or set ICPX=...)
TMS_PLATFORM=xpu pip install --no-build-isolation .   # or `make build-xpu`

If the sourced oneAPI's libsycl major differs from your torch runtime's, pin a matching compiler with ICPX=/opt/intel/oneapi/compiler/<ver>/bin/icpx. Then use region/pause/resume as on CUDA with hook_mode="torch", device="xpu".

Note: torch.xpu.memory_allocated() / mem_get_info() do not reflect pages released by zeVirtualMemUnmap; use sysman (ZES_ENABLE_SYSMAN=1) to verify.

Development

make reinstall

You can use this command for local testing:

pytest /path/to/torch_memory_saver/test

Or this one to test a single case (e.g. the simple one here):

pytest /path/to/torch_memory_saver/test/test_examples.py::test_simple -s