Исправить RuntimeError: CUDA out of Memory

Иллюстрация ошибки

PyTorch широко используется в R&D и промышленности. При работе с большими моделями или объёмными наборами данных нередко появляется ошибка “RuntimeError: CUDA out of Memory”. Она указывает, что память GPU исчерпана и дальнейшие вычисления не могут завершиться.

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

Что такое проблема “CUDA out of Memory”?

Схематичное изображение GPU и использования памяти при обучении

CUDA — это технология NVIDIA (Compute Unified Device Architecture), которая позволяет выполнять вычисления на GPU. Многие фреймворки машинного обучения (PyTorch, TensorFlow, Stable Diffusion и др.) используют GPU для ускорения обучения и инференса. Когда потребление памяти на GPU превышает доступный объём, операции прерываются и возникает ошибка “CUDA out of Memory”.

Краткое определение: CUDA out of Memory — ситуация, в которой все доступные блоки видеопамяти заняты и система не может выделить памяти для следующей операции.

Что вызывает ошибку в PyTorch

Диаграмма основных причин переполнения памяти GPU при обучении

Ниже перечислены типичные причины, почему Python/PyTorch может сообщать о нехватке памяти GPU:

Слишком большой batch size. Обработка большего количества образцов одновременно требует пропорционально больше памяти.
Сложная или большая архитектура модели — много слоёв и параметров увеличивает требование к памяти для тензоров, активаций и градиентов.
Утечки памяти в коде: накопление ссылок на тензоры с requires_grad=True, хранение историй градиентов или списков тензоров без явного освобождения.
Отсутствие очистки после итераций: если не вызывать optimizer.zero_grad() или не удалять промежуточные переменные, память будет удерживаться между батчами.
PyTorch хранит граф вычислений для всех тензоров, требующих градиентов; для больших батчей и глубокой сети это быстро накапливается.
Операции с неподходящей точностью или размещением тензоров на нескольких устройствах без явного контроля памяти.

Теперь пройдёмся по методам решения — от простых до продвинутых.

Как исправить RuntimeError: CUDA out of Memory

Ниже — проверенные подходы. Пробуйте их по порядку: сначала простые и быстрые, затем более сложные.

Устранение 1: Изменение batch size

Самое простое: уменьшите batch size. Это напрямую снижает объём памяти, требуемый для активаций и буферов.

Если вы не можете уменьшать batch size из-за стабильности обучения, используйте накопление градиентов (gradient accumulation): делайте несколько шагов backward подряд и только потом вызывайте optimizer.step(). Это эквивалентно увеличению эффективного батча без увеличения памяти.

Пример в PyTorch (комментарии переведены):

optimizer.zero_grad()  # Явно обнуляем буферы градиентов

for i in range(num_mini_batches):
    inputs, labels = next(training_data)
    inputs, labels = inputs.to(device), labels.to(device)
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()  # Обратное распространение градиента

    if (i + 1) % accumulation_steps == 0:  # Ждём несколько backward
        optimizer.step()  # Выполняем шаг оптимизатора
        optimizer.zero_grad()  # Сбрасываем градиенты

Советы:

Начните с уменьшения в 2×, 4×, затем подберите стабильный размер.
Контролируйте скорость сходимости: при очень маленьких батчах может потребоваться скорректировать learning rate.

Устранение 2: Смешанная точность (Mixed Precision Training)

Смешанная точность использует float16 (FP16) там, где это безопасно, и float32 (FP32) там, где требуется стабильность. Это уменьшает потребление памяти и ускоряет вычисления на современных GPU.

PyTorch предоставляет пакет torch.cuda.amp для автоматизации этого процесса. Пример использования:

model = ...
optimizer = ...
scaler = torch.cuda.amp.GradScaler()

for inputs, labels in data:
    inputs, labels = inputs.to(device), labels.to(device)
    optimizer.zero_grad()

    # Выполняем forward с автокастингом
    with torch.cuda.amp.autocast():
        outputs = model(inputs)
        loss = loss_fn(outputs, labels)

    # Масштабируем loss и выполняем backward с AMP
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Примечания и подводные камни:

Часто экономия памяти достигает значительных значений (в ряде случаев до ~50% по оригинальным наблюдениям), но точное снижение зависит от модели.
BatchNorm и некоторые операции чувствительны к снижению точности — можно оставить их в float32.
Не забудьте включить оптимизации cudnn, если требуется:

torch.backends.cudnn.benchmark = True
torch.backends.cudnn.enabled = True

Устранение 3: Меньшая архитектура или оптимизированные блоки

Если модель слишком велика, подумайте о замене на более компактную архитектуру или оптимизированные блоки. Некоторые архитектуры специально разработаны для экономии памяти и скорости, например MobileNet, EfficientNet и их вариации.

Пример простейшей модели в PyTorch для понимания структуры:

import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 16 * 16, 10)  # Подберите размер под задачу

Варианты оптимизаций:

Используйте depthwise-separable convolutions.
Применяйте меньший каналный размер или число слоёв.
Пробуйте pruned (обрезанные) или quantized модели для инференса.

Устранение 4: Явная очистка неиспользуемой памяти

PyTorch хранит кэш CUDA для ускорения аллокаций. При долгих скриптах или повторных запусках циклов полезно очищать кэш.

optimizer.zero_grad()
loss.backward()
optimizer.step()
torch.cuda.empty_cache()

Важно: torch.cuda.empty_cache() не освобождает память из-под процесса Python для ОС, но позволяет PyTorch перераспределять кэшированную память внутри процесса.

Устранение 5: Activation checkpointing (сохранение контрольных точек активаций)

Activation checkpointing сохраняет только часть активаций во время forward-прохода и вычисляет их заново при backward. Это снижает пиковое потребление памяти в обмен на дополнительное время вычислений.

В PyTorch есть утилиты для checkpointing (torch.utils.checkpoint). Пример использования:

from torch.utils.checkpoint import checkpoint

# Внутри forward:
out = checkpoint(module_or_function, input_tensor)

Примечание: checkpointing особенно полезен для очень глубоких сетей, где большинство активаций занимает память.

Устранение 6: Модульное распределение и шардирование (model sharding, tensor/model parallelism)

Если одна GPU недостаточна, распределите модель по нескольким устройствам:

DataParallel — простой, но не самый эффективный способ.
DistributedDataParallel — предпочтительный способ для обучения на нескольких GPU/нодах.
ZeRO (DeepSpeed) и FSDP (Fully Sharded Data Parallel) дают продвинутую экономию памяти за счёт шардирования оптимизаторов, градиентов и параметров.

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

Устранение 7: Отключение ненужных логов и хранение промежуточных тензоров

Иногда данные для отладки или логирования (например, сохранение выходов каждого слоя) удерживают память. Убедитесь, что вы не сохраняете ссылки на тензоры, требующие градиента, если это не целенаправленно.

Пример плохой практики:

outputs_list = []
for batch in data:
    out = model(batch)
    outputs_list.append(out)  # Нельзя хранить тензор с requires_grad=True

Лучше сохранять detatched копии или переносить данные на CPU:

outputs_list.append(out.detach().cpu())

Устранение 8: Диагностика и поиск утечек памяти

Как отладить проблему:

Выведите torch.cuda.memory_summary() и torch.cuda.memory_allocated()/reserved() перед и после ключевых операций.
Используйте профайлер PyTorch (torch.profiler) для визуализации потребления памяти.
Проверяйте наличие больших списков, словарей или кэшей в коде.

Пример вызова памяти:

print(torch.cuda.memory_summary(device=None, abbreviated=False))
print('Allocated:', torch.cuda.memory_allocated())
print('Reserved:', torch.cuda.memory_reserved())

Когда предложенные методы не помогут

Контрпримеры и ограничения:

Если модель действительно больше, чем физический объём памяти всех доступных GPU и нет возможности шардирования или распределения, единственный выход — уменьшать модель или перебирать решения с offloading на CPU.
На этапах инференса часто помогает квантование до int8; но это шаг для продакшена и требует валидации качества.
При аппаратных ограничениях (старые GPU с маленькой памятью) многие оптимизации будут лишь частично эффективны.

Эвристика для выбора подхода (как решать быстро)

Уменьшите batch size в 2× — если помогло, подберите максимальный возможный.
Включите gradient accumulation, если batch size критичен для обучения.
Включите AMP — обычно даёт хорошее соотношение памяти/скорости.
Если по-прежнему недостаточно — пробуйте activation checkpointing и/или уменьшайте модель.
При распределённом обучении рассматривайте ZeRO/DeepSpeed/FSDP.

Шаблонный план действий (SOP)

Перезапустите Python-контейнер/ядро, чтобы очистить состояние GPU.
Запустите минимальный скрипт с небольшим batch size и проверьте успешный проход.
Пошагово увеличивайте batch size, контролируя torch.cuda.memory_summary().
Если ошибка появляется, включите AMP и повторите тест.
Если всё ещё не хватает — включите activation checkpointing или уменьшите модель.
Если нужно масштабирование — настройте DDP/ZeRO или перенос на кластер.

Роли и чек-листы

Для исследователя/разработчика:

Попробовать уменьшить batch size.
Включить accumulation и AMP.
Проверить код на удержание тензоров (detach/cpu).

Для инженера ML Ops:

Предоставить лог доступа к GPU и метрики памяти.
Рассмотреть добавление узлов с большим объёмом памяти или настройку кластера.
Настроить мониторинг CUDA и алерты.

Для продакшн-инженера при инференсе:

Использовать квантование и оптимизированные рантаймы (TensorRT, ONNX Runtime).
Профилировать латентность и потребление памяти при пиковых нагрузках.

Критерии приёмки

Скрипт обучения проходит хотя бы одну эпоху без ошибки OOM для выбранного batch size.
Пиковое потребление памяти стабилизировано и не растёт между эпохами (нет утечки).
Точность/показатели модели не ухудшились неприемлемо после оптимизаций (AMP/quantize).

Короткое руководство по отладке (быстрые команды)

Перезапуск окружения: закройте процессы, использующие GPU (nvidia-smi) и перезапустите kernel.
Проверка занятых процессов: nvidia-smi
Вывод статистики памяти в коде: torch.cuda.memory_summary()

Итог

Причин ошибки “CUDA out of Memory” немного, но способов её исправить много. Начните с уменьшения batch size и очистки кэша, затем переходите к смешанной точности и activation checkpointing. Для масштабных моделей используйте распределённые подходы (DDP, ZeRO). Всегда диагностируйте проблему с помощью инструментов PyTorch и системных утилит, чтобы понять, где именно происходит расход памяти.

Если вам нужен разбор конкретного кода или вывод nvidia-smi — вставьте фрагменты, и я помогу с пошаговой диагностикой.

Дальнейшее чтение:

How to Fix Driver Overran Stack Buffer Error?
How to Check HDD and SSD Health in Windows?
What is Killer Network Service? Everything You Need to Know

Краткое резюме

Уменьшение batch size и накопление градиентов — самые простые методы.
AMP даёт хорошую экономию памяти и ускорение, но требует тестирования для batch norm.
Activation checkpointing и шардирование модели помогают для очень больших сетей.
Всегда используйте torch.cuda.memory_summary() и профайлеры для точной диагностики.