DEV Community

Rodion Kazennov
Rodion Kazennov

Posted on Originally published at allkeep.org

PyTorch on a Mac: why memory looks freed after .to("cpu")

#pytorch #python #machinelearning #macos

On a Mac, a tensor moved from mps to the CPU seems to disappear from memory. The process footprint, the number Activity Monitor shows in its Memory column, drops to almost nothing, and so does the MPS driver count. The copy is still there and still takes the same memory. This is the minimal example I put together for a diffusers discussion about offloading, with a 4 GiB tensor standing in for a model's weights (torch 2.14.0, macOS 26.5.2, M5 Pro with 48 GB):

import ctypes, os, time
import torch

# rusage_info_v2 from <libproc.h>; phys_footprint is the Memory column of Activity Monitor
FIELDS = ("user_time", "system_time", "pkg_idle_wkups", "interrupt_wkups", "pageins",
          "wired_size", "resident_size", "phys_footprint")
class RUsageInfoV2(ctypes.Structure):
    _fields_ = [("uuid", ctypes.c_uint8 * 16)] + [(f, ctypes.c_uint64) for f in FIELDS] \
        + [(f"rest{i}", ctypes.c_uint64) for i in range(10)]

libproc = ctypes.CDLL("/usr/lib/libproc.dylib")

def show(step):
    torch.mps.synchronize()
    info = RUsageInfoV2()
    libproc.proc_pid_rusage(os.getpid(), 2, ctypes.byref(info))
    print(f"{step:<14} footprint {info.phys_footprint / 2**30:4.1f} GiB, "
          f"mps driver {torch.mps.driver_allocated_memory() / 2**30:4.1f} GiB")

show("start")
w = torch.randn(2**30, device="mps")      # 4 GiB of float32, standing in for a model's weights
show("on mps")
w = w.to("cpu")                           # what model CPU offload does when a model's turn ends
torch.mps.empty_cache(); time.sleep(2)
show("moved to cpu")
float(w.sum())                            # read every page of the cpu copy
show("cpu copy read")
Enter fullscreen mode Exit fullscreen mode 
start          footprint  0.1 GiB, mps driver  0.0 GiB
on mps         footprint  4.2 GiB, mps driver  4.0 GiB
moved to cpu   footprint  0.1 GiB, mps driver  0.0 GiB
cpu copy read  footprint  4.2 GiB, mps driver  0.0 GiB
Enter fullscreen mode Exit fullscreen mode

I ran it three times and got the same output each time.

What the numbers say

Right after the move, both counters are back near zero, but w holds all 4 GiB of its values. Once the CPU reads every page of the copy, the footprint is 4.2 GiB again. So a footprint taken between .to("cpu") and the next time the CPU touches those weights does not show them.

The reason the move frees nothing in the first place is the hardware. On Apple Silicon the CPU and the GPU use one pool of physical memory. On a PC, moving weights from the GPU to the CPU frees VRAM for the next model. On a Mac, the weights "on the CPU" sit in the same RAM the GPU needs.

Why it matters for model CPU offload

diffusers' enable_model_cpu_offload() keeps each model on the CPU and moves it to the GPU only while it runs. I ran Qwen-Image-2.1 at 1024 × 1024 that way on the same Mac (trace and write-up). The process footprint never went above 18.5 GB, which would fit on a 48 GB Mac with room to spare. Swap grew by 3.5 GB during denoising and had grown by 9.6 GB by the VAE decode, and my memory guard stopped the run before it wrote an image.

A chart of the footprint alone would have called that run a fit. The example above shows one way the footprint can stay low while the weights are still in memory. I didn't measure which pages went to swap in the offload run.

How to check your own run

Look at swap, not only at the footprint:

sysctl vm.swapusage             # compare "used" before and after the run
sysctl kern.memorystatus_level  # system-wide free memory, in percent
Enter fullscreen mode Exit fullscreen mode

If swap grows while the job runs, it doesn't fit, whatever the footprint says. stageload (pip install stageload), a small library I wrote for this, records both in a trace every 0.5 s with MemoryMeter, and stageload guard stops a command once swap grows past a budget:

stageload guard --wait-free 40 --swap-budget 8G -- python generate.py
Enter fullscreen mode Exit fullscreen mode

What does free the memory

For a pipeline that runs its models one after another, the fix is to load a model when its stage starts and release it when the stage ends, instead of parking it on the CPU. stageload's release(module) moves every parameter and buffer to PyTorch's meta device, collects garbage and empties the MPS cache; a tensor on meta has no storage, so the memory goes back even while other code still holds a reference to the module. With the text encoder (16.3 GB) and the transformer (13.3 GB) loaded one at a time this way, Qwen-Image-2.1 peaked at 19.0 GB on the same Mac, and swap did not grow in either of two runs. The full setup is in How to run Qwen-Image-2.1 on a 48 GB Mac without swap.

What I did not measure

  • The free memory of the whole system during the example: other jobs on this Mac move it by several GiB between runs, so I couldn't get a clean number.
  • Whether del plus torch.mps.empty_cache() gives the memory back as fully as moving to meta, and how quickly.
  • Other macOS and torch versions, and Macs with more memory.

Top comments (0)