← Назад в лабу

PyTorch на Mac: почему память кажется свободной после .to("cpu")

pytorchapple-siliconmemoryguide

На Mac тензор, перенесённый из mps на CPU, будто исчезает из памяти. Footprint процесса, то самое число в колонке «Память» в Мониторинге системы, падает почти до нуля, и счётчик драйвера MPS тоже. Копия при этом на месте и занимает ту же память. Это минимальный пример, который я собрал для обсуждения выгрузки в diffusers; тензор на 4 ГиБ изображает веса модели (torch 2.14.0, macOS 26.5.2, M5 Pro с 48 ГБ):

На Mac у CPU и GPU одна память: после .to("cpu") веса лежат в том же пуле, а Мониторинг системы показывает 0,1 ГБ

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")
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

Я запускал его три раза, и вывод каждый раз был тот же.

Что говорят числа

Сразу после переноса оба счётчика почти на нуле, хотя в w лежат все 4 ГиБ значений. Как только CPU прочитает каждую страницу копии, footprint снова 4,2 ГиБ. Значит, footprint, снятый между .to("cpu") и следующим обращением CPU к этим весам, их не показывает.

А почему перенос вообще ничего не освобождает, решает железо. На Apple Silicon у CPU и GPU один пул физической памяти. На ПК перенос весов с GPU на CPU освобождает видеопамять для следующей модели. На Mac веса «на CPU» лежат в той же RAM, которая нужна GPU.

Почему это важно для model CPU offload

enable_model_cpu_offload() из diffusers держит каждую модель на CPU и переносит её на GPU только на время её работы. Я так запускал Qwen-Image-2.1 в 1024 × 1024 на том же Mac (трасса и разбор). Footprint процесса ни разу не поднялся выше 18,5 ГБ, а это с запасом влезает в Mac на 48 ГБ. Своп вырос на 3,5 ГБ за шумоподавление и на 9,6 ГБ к декодированию VAE, и мой сторож памяти остановил прогон до того, как тот записал картинку.

По графику одного footprint этот прогон сочли бы влезающим. Пример выше показывает один из способов, каким footprint остаётся низким, пока веса всё ещё в памяти. Какие страницы ушли в своп в прогоне с offload, я не мерил.

Как проверить свой прогон

Смотрите на своп, а не только на footprint:

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

Если своп растёт, пока работает задача, она не влезает, что бы ни говорил footprint. stageload (pip install stageload), небольшая библиотека, которую я написал для этого, пишет и то и другое в трассу каждые 0,5 с через MemoryMeter, а stageload guard останавливает команду, как только своп вырастет сверх бюджета:

stageload guard --wait-free 40 --swap-budget 8G -- python generate.py

Что на самом деле освобождает память

Для конвейера, который запускает модели одну за другой, выход такой: загружать модель, когда начинается её стадия, и освобождать, когда стадия кончается, а не парковать её на CPU. release(module) из stageload переносит каждый параметр и буфер на устройство meta в PyTorch, собирает мусор и чистит кеш MPS; у тензора на meta нет хранилища, поэтому память возвращается, даже пока другой код держит ссылку на модуль. Когда текстовый энкодер (16,3 ГБ) и трансформер (13,3 ГБ) так загружались по одному, пик Qwen-Image-2.1 на том же Mac был 19,0 ГБ, и своп не вырос ни в одном из двух прогонов. Вся настройка в руководстве Как запустить Qwen-Image-2.1 на Mac с 48 ГБ без свопа.

Что я не мерил

PyTorch на Mac: почему память кажется свободной после .to("cpu") · Rodion Kazennov