Нехватка памяти MPS в PyTorch 2.14: как исправить обучение в 2026 году

Процесс обучения внезапно получает MPS out of memory, macOS показывает высокое давление памяти, а после уменьшения batch size программа всё равно постепенно занимает память и перестаёт отвечать.

Самое быстрое решение — не отключать ограничения MPS, а последовательно отделить реальные тензоры от кэша, неосвобождённого графа, динамических форм и CPU fallback; затем проверить минимальный скрипт в чистой среде PyTorch 2.14.

Кому нужна эта проверка

Эта инструкция рассчитана на исследователей, которые обучают или запускают инференс на Apple Silicon и получают ошибку MPS out of memory.

Она также пригодится сопровождающим проект, где один и тот же код должен давать сопоставимый результат на Linux GPU и macOS, а также техническим руководителям лабораторий, которым нужно временно получить чистую среду без покупки отдельного Mac.

Последнее обновление: 12 сентября 2026 года. Дата выпуска PyTorch 2.14 и сведения об изменениях MPS сверены с официальным объявлением о PyTorch 2.14, а описание API и переменных окружения — с документацией PyTorch.

Сначала зафиксируйте, какой именно сбой происходит

В PyTorch 2.14 были изменены MPS-кэш-аллокатор и некоторые пути работы с памятью и копированием — это подтверждено в официальном объявлении релиза. Однако из этого не следует, что обновление устраняет все случаи нехватки памяти или роста потребления. Релиз может изменить поведение распределителя, но не исправит ссылку на граф вычислений, слишком большую форму входа или неподдерживаемую операцию.

Перед изменениями сохраните базовую запись:

• версию PyTorch и Python;
• версию macOS;
• модель чипа Apple Silicon;
• фактический размер модели, batch size, длину последовательности и разрешение входа;
• полный текст ошибки;
• момент сбоя: при загрузке модели, первом проходе, backward, оптимизаторе или после длительной работы;
• наличие CPU fallback;
• значения памяти до запуска, после загрузки модели и после нескольких одинаковых итераций.

Разделяйте четыре симптома:

1. Ошибка распределения MPS. PyTorch не может выделить очередной блок памяти и завершает операцию.
2. Завершение процесса macOS. Система может остановить приложение из-за общего давления памяти. Это не всегда тот же случай, что исключение PyTorch.
3. Постепенное замедление и зависание. Причиной может быть активное обращение к виртуальной памяти или конкуренция с другими приложениями.
4. Постоянный рост resident memory. Здесь нужно искать ссылки на тензоры, графы, результаты и динамические формы, а не просто уменьшать batch size.

Для оценки состояния системы используйте описание давления памяти в Activity Monitor от Apple. Показатель системного давления нельзя подменять одним числом из PyTorch: приложение может иметь умеренный показатель MPS, но при этом вся система уже испытывать нехватку ресурсов.

Как быстро понять: модель слишком велика или проблема появилась позже

Размер параметров — только первая часть расчёта. Во время обучения одновременно расходуются параметры модели, активации, градиенты, состояния оптимизатора и временные буферы операций. При инференсе часть этих расходов исчезает, но промежуточные результаты всё равно могут быть крупными.

Поэтому сначала изменяйте только одну переменную:

1. Запустите минимальную конфигурацию с тем же random seed.
2. Уменьшите batch size, сохранив длину последовательности и разрешение.
3. Если сбой остаётся, верните batch size и уменьшите длину последовательности или разрешение.
4. После этого отдельно проверьте размер модели и режим вычислений.
5. Повторите одинаковый минимальный запуск несколько итераций и сравните память после каждой итерации.

Такой порядок важнее, чем одновременное уменьшение всех параметров. Если одновременно изменить batch size, форму входа, precision и модель, станет невозможно определить, что именно пересекало границу доступной памяти.

Единая память Apple Silicon также не означает, что весь объём компьютера можно безусловно отдать MPS. macOS, Python, файловый кэш, графический интерфейс и другие процессы используют ту же физическую память. Для решения нужно учитывать доступность ресурсов в конкретный момент, а не только установленный объём памяти.

Для каждого варианта задайте условие прохождения: минимальный скрипт должен завершить несколько одинаковых итераций, не увеличивая потребление после каждой из них и не создавая ошибку на backward. Если уменьшенная конфигурация проходит только один шаг, это ещё не подтверждение исправления.

Что показывают счётчики MPS и почему их нельзя смешивать

Чтобы понять, где растёт память, запрашивайте как минимум два значения:

import torch

def report_memory(label):
    allocated = torch.mps.current_allocated_memory()
    driver = torch.mps.driver_allocated_memory()
    print(
        f"{label}: "
        f"current_allocated_memory={allocated:,} bytes, "
        f"driver_allocated_memory={driver:,} bytes"
    )

report_memory("before model")
# загрузка модели
report_memory("after model")
# один forward или один train step
report_memory("after step")

current_allocated_memory() показывает память, занятую тензорами, которыми управляет PyTorch. Его назначение описано в официальной документации API current_allocated_memory.

driver_allocated_memory() отражает память, выделенную MPS-драйвером для размещения тензоров и кэша. Это другой уровень наблюдения; описание доступно в документации driver_allocated_memory.

Если current allocated стабилен, а driver allocated растёт, подозрение падает на кэш, фрагментацию, временные буферы или особенности распределителя. Если растут оба значения после одинаковых итераций, ищите живые ссылки на тензоры, графы и результаты. Если оба показателя относительно стабильны, но macOS сообщает о высоком давлении, проверяйте другие процессы и общую память системы.

Для дополнительного ориентира можно записать torch.mps.recommended_max_memory(). Этот показатель следует использовать как рекомендацию распределителя, а не как обещание, что приложение сможет безопасно занять весь результат. Его назначение приведено в документации recommended_max_memory.

Почему torch.mps.empty_cache() не освобождает всю память

empty_cache() освобождает неиспользуемые блоки кэша. Он не удаляет тензоры, на которые всё ещё существуют ссылки, и не разрушает граф вычислений, необходимый для backward. Поэтому после вызова память может почти не измениться — это ожидаемо, если проблема не в свободном кэше. Ограничение явно описано в документации torch.mps.empty_cache.

Проверяйте это отдельным экспериментом:

import gc
import torch

# удалить переменные, которые действительно больше не нужны
del temporary_output

gc.collect()
torch.mps.empty_cache()
report_memory("after cleanup")

Если переменная temporary_output всё ещё находится в списке, словаре, объекте логгера или замыкании, удаление локального имени не освободит сам тензор. Такой тест нужно проводить после явного удаления контейнеров и только в отладочной копии скрипта.

Важно: не следует начинать с принудительной отмены ограничений высокого уровня памяти. Переменные окружения MPS имеют конкретное назначение и риски; сверяйте их с официальной документацией переменных окружения MPS, а не с советом из случайного обсуждения.

Почему уменьшение batch size не останавливает рост

Если после уменьшения batch size память продолжает расти, размер пакета, вероятно, не является единственной причиной. Наиболее частые варианты связаны с жизненным циклом объектов Python и графа:

• loss добавляется в список без .detach() или преобразования в обычное число;
• предсказания сохраняются вместе с графом;
• скрытое состояние рекуррентной модели переносится между итерациями без отсоединения;
• результаты для логирования остаются на устройстве;
• в обучении не очищаются или неправильно накапливаются градиенты;
• в инференсе отсутствует torch.no_grad() или torch.inference_mode() там, где они допустимы;
• словарь метрик хранит GPU- или MPS-тензоры вместо скалярных значений.

Небезопасный шаблон выглядит так:

history.append(loss)
predictions.append(output)

Для диагностического варианта используйте явное отсоединение:

history.append(float(loss.detach().cpu()))
predictions.append(output.detach().cpu())

В инференсе отдельно проверьте область действия:

with torch.inference_mode():
    output = model(batch)

Это не универсальная замена обучению и не должно механически добавляться вокруг train loop. Задача проверки — доказать, удерживает ли код граф и выходные данные между итерациями.

Сделайте два прогона: один с исходным логированием, второй без сохранения результатов. Если рост исчезает только во втором варианте, причина находится в коде удержания ссылок, а не обязательно в MPS-аллокаторе.

Динамические формы, кэш и повторная компиляция

Постоянно меняющиеся формы входа усложняют анализ. Разные длины последовательностей, разрешения изображений или размеры последнего batch могут приводить к созданию новых буферов и сохранению блоков в распределителе. Поэтому сравните два режима:

• фиксированная форма входа на всех итерациях;
• те же данные в исходном динамическом режиме.

Запишите current_allocated_memory() и driver_allocated_memory() до каждой итерации. Если фиксированная форма стабилизирует значения, не объявляйте это доказательством утечки: это может быть следствием кэширования и повторного использования буферов.

Проверяйте также обработку последнего неполного batch, предварительную компиляцию, изменение масок и создание новых тензоров внутри цикла. Для минимального теста исключите аугментации, логирование, сохранение чекпойнтов и внешние callback-функции. Затем возвращайте их по одному.

Отдельный сигнал — рост памяти только после перехода к новым формам, а не после повторения одной и той же формы. Такой результат направляет диагностику к динамическим буферам и кэшу, но не доказывает наличие дефекта в конкретной версии PyTorch.

CPU fallback меняет не только скорость, но и путь хранения данных

В MPS есть операции, которые могут быть неподдерживаемыми или иметь ограничения. При включённом CPU fallback часть вычислений может выполняться на CPU, а входы и результаты — перемещаться между устройствами. В таком случае рост общей памяти нельзя автоматически приписать MPS.

Проверьте:

• все ли входы действительно находятся на mps;
• не перемещается ли часть модели на CPU вручную;
• включён ли fallback и какие операции его вызывают;
• в какой строке появляется ошибка;
• меняются ли результаты при полном запуске на CPU;
• нет ли неявного копирования в логгере или обработчике метрик.

Для воспроизводимой проверки создайте два минимальных запуска: один на MPS, второй на CPU. Используйте одинаковые входные данные, seed и порядок операций. Сравнивайте не производительность, а этап сбоя, форму результата, наличие nan и факт завершения.

Официальные ограничения и устройство MPS описаны в заметке PyTorch о бэкенде MPS. Отдельные GitHub issue можно использовать для поиска похожего симптома, но пользовательский отчёт о конкретной проблеме не подтверждает, что тот же дефект присутствует во всех моделях, версиях macOS или конфигурациях Apple Silicon.

Чистая среда решает вопрос «код или машина»

После локальных исправлений не переносите сразу всю рабочую среду. Зафиксируйте версии в отдельном окружении, установите только PyTorch 2.14 и минимальные зависимости, очистите экспериментальные кэши и запустите короткий скрипт.

Порядок проверки:

1. Создайте новое изолированное окружение Python.
2. Зафиксируйте версию PyTorch, Python и системную версию macOS.
3. Установите только зависимости, которые нужны минимальному примеру.
4. Подготовьте маленький входной файл, модель или синтетический тензор.
5. Запишите модель чипа Apple Silicon и способ подключения к Mac.
6. Запустите CPU-вариант и сохраните факт завершения.
7. Запустите MPS-вариант с одинаковым seed.
8. Снимите current allocated, driver allocated и системное давление памяти до и после одинаковых итераций.
9. Повторите тест после изменения только одного параметра.
10. Сверьте результаты с Linux GPU, не объявляя их автоматически идентичными.

В лаборатории без Mac такая изолированная проверка возможна на удалённом физическом Apple Silicon Mac. Это полезнее, чем пытаться судить о MPS по эмулятору или по машине коллеги с неизвестными пакетами. Для подключения и ограничений удалённой работы заранее проверьте справочную информацию JexMac, а варианты временного доступа сопоставьте с актуальными условиями аренды Mac.

Комплект приёмки должен содержать:

• файл фиксации окружения;
• минимальный воспроизводящий скрипт;
• входной пример;
• полный текст ошибки;
• значения трёх уровней памяти: PyTorch allocated, driver allocated и системное давление;
• краткое описание результата CPU и MPS;
• Linux GPU-сравнение;
• дату и идентификатор повторного запуска.

Если второй человек не может повторить результат по этому комплекту, причина ещё не локализована. Если ошибка появляется только в одной комбинации версии PyTorch, macOS и модели, допустимо временно зафиксировать рабочую версию и вести две ветки проверки — Linux GPU и Apple Silicon — вместо заявления о полной эквивалентности.

Пошаговая проверка перед изменением окружения

• [ ] Сохранить полный текст ошибки и момент её появления.
• [ ] Записать версии PyTorch, Python и macOS.
• [ ] Указать модель чипа Apple Silicon и форму входных данных.
• [ ] Снять current_allocated_memory().
• [ ] Снять driver_allocated_memory().
• [ ] Проверить давление памяти macOS в Activity Monitor.
• [ ] Повторить один и тот же шаг с фиксированной формой входа.
• [ ] Изменить только batch size, затем отдельно длину последовательности или разрешение.
• [ ] Проверить списки, словари и логгеры, где могут сохраняться тензоры.
• [ ] Для инференса проверить границы inference_mode() или no_grad().
• [ ] Проверить CPU fallback и неявные перемещения между CPU и MPS.
• [ ] Сравнить минимальный MPS-скрипт с CPU-запуском.
• [ ] Повторить тест в чистом окружении PyTorch 2.14.
• [ ] Зафиксировать рабочую комбинацию версий, если ошибка зависит от среды.
• [ ] Не отключать ограничения памяти до понимания причины.

Как выбрать следующий шаг по результатам диагностики

Эта таблица не заменяет эксперимент: оценка уверенности относится к направлению поиска, а не к окончательному диагнозу. Особенно осторожно трактуйте случаи, где задействованы динамические формы или CPU fallback.

Когда удалённый Mac оправдан, а когда лучше остаться на Linux

Если задача состоит в длительном стабильном обучении на больших моделях, требуется конкретный физический интерфейс или нужна предсказуемая серверная эксплуатация, Linux GPU может оставаться основным контуром. Apple Silicon в такой схеме разумно использовать как дополнительную платформу для воспроизведения macOS-специфичных ошибок, проверки переносимости и локальной разработки.

Если же лаборатория не имеет Mac, а ошибка возникает только в MPS-пути, покупка устройства ради короткой диагностики создаёт дополнительные расходы и ответственность за обслуживание. Краткосрочная аренда реального Mac позволяет проверить именно нужную связку PyTorch 2.14, macOS, чипа и минимального скрипта, не выдавая результаты эмуляции за результат физического MPS.

У текущего варианта — проверки только на Linux или на случайном Mac коллеги — есть несколько недостатков: отсутствует точная macOS-среда, трудно исключить старые зависимости, нельзя гарантировать повторяемость и приходится тратить время на настройку устройства, которое не используется постоянно. В такой ситуации аренда Mac у JexMac может быть рациональнее для короткого эксперимента: можно получить отдельную среду с root-доступом, провести приёмку по описанному комплекту и затем решить, нужен ли постоянный Mac, Linux GPU или двойной контур. Начать можно с описания вариантов доступа JexMac.

Главный критерий выбора — не сам факт успешного запуска, а способность второго участника воспроизвести тот же сбой и подтвердить исправление по одинаковому скрипту, входу, версиям и журналу памяти. Пока это не выполнено, безопаснее сохранять Linux GPU и Apple Silicon как два проверочных контура, а не объявлять результат полностью эквивалентным.