Google Colab runtime and PyTorch CUDA memory: should cleanup be automatic after cell execution?

Answer

No, the Google Colab runtime does not perform any automatic GPU memory reclamation after a cell completes. Memory management is delegated entirely to the Python process and the PyTorch CUDA caching allocator.

Technical Explanation

When a cell finishes executing, any tensors that are no longer referenced by the Python namespace are marked for deletion by Python's garbage collector. However, PyTorch does not immediately release this memory back to the GPU driver. Instead, it holds the memory in a caching allocator to avoid the high overhead of frequent CUDA malloc/free calls during training.

While torch.cuda.empty_cache() forces the allocator to release unused cached blocks back to the driver (reducing the memory reported by nvidia-smi), it does not solve memory fragmentation. If a session involves many allocations of varying sizes, the available free memory may become non-contiguous. In such cases, a large allocation request may trigger an Out-of-Memory (OOM) error even if the total free memory appears sufficient, a state that typically requires a runtime restart to resolve.

Recommended Cleanup Workflow

1. Explicit Deletion: Use del tensor_name for large models or buffers to remove references immediately.
2. Garbage Collection: Call import gc; gc.collect() to ensure Python has cleared the references.
3. Cache Clearing: Execute torch.cuda.empty_cache() to return the memory to the system.

Verification Process

To verify if memory is being reclaimed without restarting, compare the following metrics before and after your workload:

import torch
import subprocess

def get_gpu_usage():
    res = subprocess.check_output(['nvidia-smi', '--query-gpu=memory.used', '--format=csv,nounits,noheader'])
    return res.decode().strip()

print(f"Driver Memory Used: {get_gpu_usage()} MB")
print(f"PyTorch Allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
print(f"PyTorch Reserved: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

If "PyTorch Reserved" remains high after cleanup, the allocator is still holding memory for future use. If "Driver Memory Used" remains high after empty_cache(), there are likely still live tensors in the namespace. You can also inspect torch.cuda.memory_summary() for a detailed view of segment fragmentation.

Fragmentation Mitigation

Restarting the runtime is the only guaranteed full reset because it kills the process and the CUDA context, but it also clears installed packages and mounted state. Before resorting to that, note that newer PyTorch releases support PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, which can reduce fragmentation in long sessions. This option is version-sensitive and may not exist in older Colab images, so check your installed PyTorch version first. Also keep in mind that nvidia-smi reports process-level reservations, including CUDA context overhead and library workspaces, so it can over-report relative to actual tensor storage.