This library addresses the challenge of managing large memory footprints in PyTorch, particularly for GPU-bound workloads. It allows users to temporarily release the memory occupied by tensors, freeing up GPU resources, and then resume them later, preserving their virtual addresses. This is beneficial for scenarios like large model inference or training where memory constraints are a bottleneck, enabling the execution of larger models or longer sequences than otherwise possible.

How It Works

The core mechanism involves defining regions of code where tensors are marked as "pauseable" using the torch_memory_saver.region() context manager. Upon calling torch_memory_saver.pause(), the physical CUDA memory allocated to these tensors is released, while their virtual addresses are maintained. When torch_memory_saver.resume() is invoked, the physical memory is re-allocated, restoring the tensors to their previous state without altering their virtual addresses. This approach allows for dynamic memory management, crucial for optimizing GPU utilization.

Quick Start & Requirements

• Installation: make reinstall for development/local testing.
• Prerequisites: PyTorch, CUDA (implied by GPU tensor operations).
• Testing: pytest is used for local testing.
• Documentation: Examples and details are linked via GitHub issues: https://github.com/sgl-project/sglang/issues/2542#issuecomment-2563641647 and `https://github.com/sgl-project/sglang/issues/70