MPS backend out of memory на Mac: что значат числа и как задать watermark

Когда задаче PyTorch на Mac не хватает памяти, ошибка заканчивается советом:

RuntimeError: MPS backend out of memory (MPS allocated: 18.50 GiB, other allocations: 448.00 KiB, max allowed: 18.72 GiB). Tried to allocate 512.00 MiB on shared pool. Use PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0 to disable upper limit for memory allocations (may cause system failure).

Я хотел понять, что будет, если этому совету последовать, и заполнял память MPS шагами по 512 МиБ на своём M5 Pro с 48 ГБ (macOS 26.5.2, PyTorch 2.14.0), по одному прогону на каждую настройку, под сторожем, который останавливает процесс, когда своп вырастает на 2 ГБ. Вышло две вещи. На этом Mac предел по умолчанию выше объёма RAM, поэтому задача, которой не хватает памяти, сползает в своп и до этой ошибки так и не доходит. А чтобы понизить предел, нужны две переменные, а не одна.

Откуда берётся «max allowed»

max allowed означает предел: PYTORCH_MPS_HIGH_WATERMARK_RATIO, умноженный на то, что возвращает torch.mps.recommended_max_memory(). На моём Mac это 37,44 ГиБ, и 0,5 × 37,44 = 18,72 ГиБ, то самое число из ошибки выше. MPS allocated в моём логе совпадал с torch.mps.current_allocated_memory(), то есть с памятью под тензоры. Как посмотреть своё базовое число:

import torch
print(torch.mps.recommended_max_memory() / 2**30, "GiB")

Почему Mac с 48 ГБ сначала уходит в своп

По документации PyTorch верхний порог по умолчанию равен 1,7. На моём Mac это ставит предел примерно на 63,6 ГиБ, больше памяти, чем есть в машине. Вот что сделала каждая настройка:

Со значением по умолчанию, с 0,0 и с 1,0 своп начинал расти, как только процесс переходил 33,5–36,2 ГБ, а PyTorch так и не выдал ошибку: его предел лежал дальше, чем кончается RAM. Так что 0.0, совет из самой ошибки, здесь ничего заметного не изменил. Предел и так был недостижим, и ничто в PyTorch прогон не остановило.

Как задать: две переменные

Если задать только верхний порог, всё падает на первом обращении к MPS:

$ PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.8 python fill.py
RuntimeError: invalid low watermark ratio 1.4

Нижний порог по умолчанию на Apple Silicon равен 1,4 и не может быть выше верхнего (кроме случая, когда верхний равен 0,0 и предел выключен). Задайте оба:

PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.8 PYTORCH_MPS_LOW_WATERMARK_RATIO=0.8 python run.py

Из Python их нужно задать до того, как первый тензор попадёт на mps. PyTorch читает их при запуске MPS, а не при import torch, и изменений после этого не замечает:

import os
os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = "0.8"
os.environ["PYTORCH_MPS_LOW_WATERMARK_RATIO"] = "0.8"
import torch

Документация называет нижний порог мягким пределом и пишет, что «setting 0.0 disables adaptive commit and garbage collection». Я ставил его равным верхнему и отдельно не мерил.

Какое значение выбрать

Предел должен оставаться ниже памяти, которая у вашего Mac действительно свободна при обычных открытых программах. На моём 1,0 (37,44 ГиБ) уже было слишком много, а 0,8 (29,95 ГиБ) дал чистую ошибку без свопа. После чистой ошибки можно уменьшить batch size или разрешение и попробовать снова, а прогон, сползающий в своп, сначала тормозит весь Mac.

Ловушка при проверке

Мой первый прогон заполнял тензоры через torch.ones. Со значением по умолчанию он дошёл до поставленного мной потолка в 46 ГиБ без ошибки и вовсе без свопа: macOS сжимает страницы с одинаковыми значениями почти до нуля. Веса модели так не сжимаются, поэтому тесту памяти на Mac нужны случайные данные.

Если модель правда не влезает

Чистая ошибка только говорит, что модель не влезает. Для пайплайна из нескольких моделей загрузка одних моделей текущей стадии даёт больше любого порога, поэтому я и написал stageload. Скрипт заполнения и трассы всех прогонов выше лежат в его репозитории, в bench/results/2026-10-13-mps-watermark.