CUDA out of memory: 15 способов починить

CUDA OOM — главная боль ML-разработчика. Разбираем 15 проверенных способов: от torch.cuda.empty_cache() до flash-attention. С метриками и кодом.

Обновлено: 2026-05-19

TL;DR

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate X.XX GiB — главная боль ML-разработчика. В 80% случаев решается одной из этих мер: torch.cuda.empty_cache(), mixed precision, gradient checkpointing, уменьшение batch_size, PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. Ниже — 15 проверенных способов с реальными цифрами по экономии памяти.

Как читать ошибку

Типичное сообщение PyTorch:

torch.cuda.OutOfMemoryError: CUDA out of memory.
Tried to allocate 2.50 GiB. GPU 0 has a total capacity of 24.00 GiB
of which 1.20 GiB is free. Including non-PyTorch memory, this process
has 22.80 GiB memory in use. Of the allocated memory 20.80 GiB is
allocated by PyTorch, and 1.70 GiB is reserved by PyTorch but unallocated.

Расшифровка:

Tried to allocate 2.50 GiB — размер блока, который пытался выделить.
1.20 GiB is free — реально свободно.
22.80 GiB memory in use — общий занятый объём.
1.70 GiB reserved by PyTorch but unallocated — закэшировано (не отдано назад в OS), но не используется. Может быть отдано через empty_cache().

Два основных кейса:

Tried to allocate >> free, total ≈ used — модель/батч физически не помещается. Нужно уменьшать.
Tried to allocate ≈ small, reserved >> tried — fragmentation. Память «есть», но нет contiguous block. Часто лечится empty_cache() или expandable_segments:True.

1. torch.cuda.empty_cache() (мгновенно)

Освобождает кэшированную, но неиспользуемую память.

import torch
import gc

del large_tensor  # явно удалить ссылки
gc.collect()
torch.cuda.empty_cache()

Когда помогает: между эпохами, после inference большого батча, между обработкой разных моделей.

Когда НЕ помогает: если активные тензоры физически занимают всю VRAM. empty_cache не освобождает активные тензоры.

2. Mixed precision (fp16 / bf16)

torch.cuda.amp.autocast() запускает вычисления в half-precision. Память активаций — вдвое меньше, скорость на 50-80% выше на современных GPU.

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in data:
    optimizer.zero_grad()
    with autocast():
        output = model(batch)
        loss = loss_fn(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Экономия: ~50% памяти активаций.

Когда не работает: на старых GPU (Maxwell, Pascal). Для bf16 нужен Ampere+ (A100, RTX 30xx+).

3. Gradient checkpointing

Не хранить активации forward pass — пересчитывать их во время backward. Trade-off память на время.

from torch.utils.checkpoint import checkpoint

# Включить на model:
model.gradient_checkpointing_enable()

# Или вручную для отдельных блоков:
def custom_forward(x):
    return self.transformer_block(x)
output = checkpoint(custom_forward, input_tensor)

Экономия: 50-70% памяти активаций. Стоимость: +30% времени обучения.

Особенно эффективно для трансформеров с длинным контекстом.

4. Уменьшить batch_size

Самый прямой способ. Память активаций линейно зависит от batch_size.

Если на batch_size=32 OOM, попробуйте 16. Если важно сохранить эффективный batch_size — используйте gradient accumulation:

accum_steps = 4
optimizer.zero_grad()
for i, batch in enumerate(data):
    loss = model(batch).loss / accum_steps
    loss.backward()
    if (i + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Эффективный batch_size = batch_size × accum_steps, память — как для batch_size.

5. PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Решает fragmentation. Переменная окружения до запуска:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
python train.py

Windows:

$env:PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
python train.py

Когда помогает: «free=5 GB, Tried to allocate 2 GB → OOM». Аллокатор может расти и сжиматься, уменьшая фрагментацию.

Альтернатива: max_split_size_mb:512 — ограничивает максимальный размер блока.

6. flash-attention

Оптимизированная реализация self-attention. Сложность памяти O(N²) → O(N).

pip install flash-attn --no-build-isolation

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.3-70B-Instruct",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
)

Экономия на длинном контексте: в 4-10× для seq_len > 4096. На коротких — минимальная.

7. Квантизация для inference

Постоянная экономия памяти на весах.

# 8-bit квантизация (через bitsandbytes)
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.3-70B-Instruct",
    quantization_config=bnb_config,
)
# 70B fp16 = 140 GB → 70B int8 ≈ 70 GB.

# 4-bit квантизация
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)
# 70B 4-bit ≈ 35-40 GB.

Альтернативы: AWQ, GPTQ (внешняя предварительная квантизация → результат загружается уже квантизованным).

8. CPU offloading

Часть модели держится в RAM, переносится на GPU по мере необходимости.

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "...",
    device_map="auto",
    offload_folder="/tmp/offload",
    offload_state_dict=True,
)

Стоимость: значительное замедление (3-10×), особенно на disk-offload. RAM-offload приемлемо.

9. Освободить ссылки до OOM

PyTorch не освобождает тензоры, пока на них есть ссылки в Python. Типичная ошибка:

# BAD: history растёт с каждой эпохой
loss_history = []
for batch in data:
    loss = model(batch).loss
    loss_history.append(loss)  # держит граф!
    loss.backward()

# GOOD: detach() освобождает граф
loss_history = []
for batch in data:
    loss = model(batch).loss
    loss_history.append(loss.detach().item())  # только число
    loss.backward()

10. Уменьшить max sequence length

Для трансформеров память attention — O(seq_len²) без flash-attention. Сокращение max_seq_len с 2048 до 1024 экономит ~75% памяти attention.

В HuggingFace:

tokenizer.encode(text, max_length=1024, truncation=True)

11. Параллелизм между GPU (data / tensor / pipeline)

Если у вас несколько GPU, разные виды параллелизма распределяют модель и батчи.

DataParallel — деление батча между GPU. Простой, но overhead.
DistributedDataParallel (DDP) — стандарт, низкий overhead.
Tensor parallel (через vLLM, DeepSpeed) — модель распилена по слоям между GPU.
Pipeline parallel — каждая GPU держит свой набор слоёв.

Для inference 70B на 2× A100 80GB: tensor parallel через vLLM (--tensor-parallel-size 2).

12. DeepSpeed ZeRO / FSDP

Для обучения больших моделей. ZeRO распределяет optimizer state, gradients и параметры между GPU.

# С 🤗 Accelerate
from accelerate import Accelerator
accelerator = Accelerator(mixed_precision="bf16")
# accelerate config — выбрать DeepSpeed Stage 3 / FSDP FULL_SHARD

Экономия: на 8× A100 ZeRO-3 позволяет тренировать 70B+ модели.

13. Отключить градиенты для inference

При inference не нужно хранить активации для backward:

with torch.no_grad():
    output = model(input)

Или torch.inference_mode() — ещё агрессивнее (отключает version counter тоже):

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

Экономия: 30-50% памяти на inference. Это must-have для production.

14. Контролировать non-PyTorch память

Иногда nvidia-smi показывает 5 GB занято, но torch.cuda.memory_allocated() всего 100 MB. Виновники:

Другие процессы (Chrome с hardware acceleration, Discord, X server).
Запасной CUDA context от прерванной программы.

Лечение: nvidia-smi pmon найти процессы. Закрыть. Перезапустить.

На сервере без GUI этой проблемы обычно нет.

15. Сменить GPU

Если после всех мер всё равно OOM — модель просто не помещается на этом GPU.

Альтернативы:

Аренда A100 80GB на пару часов (через РФ-провайдеров или managed-API).
Перейти на меньшую модель той же серии (LLaMA 70B → LLaMA 13B).
Использовать managed-API вместо self-host (если задача — inference).

Чеклист диагностики

При OOM пройдите по списку (от простого к сложному):

torch.cuda.empty_cache() помог?
nvidia-smi показывает другие процессы — закрыть.
Перезапустить процесс PyTorch.
Включить PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True.
Для inference — torch.inference_mode().
Для обучения — mixed precision (autocast).
Уменьшить batch_size.
Уменьшить max_seq_len.
Gradient checkpointing.
Flash-attention.
Квантизация (8-bit или 4-bit) для inference.
Если ничего — сменить GPU.

Что почитать дальше

torch.cuda.OutOfMemoryError в ComfyUI — детальный разбор именно этой ошибки.
CUDA out of memory в ComfyUI — 12 способов в контексте ComfyUI.
Stable Diffusion: не хватает VRAM — для SD-сценариев.
vLLM out of memory — настройка vLLM для стабильности.
GPU для машинного обучения — гайд по выбору видеокарты.

Частые вопросы

Что делает torch.cuda.empty_cache()?

Освобождает кэшированную (но неиспользуемую) память PyTorch. Не освобождает активные тензоры. Запускайте между эпохами или после крупных операций.

Помогает ли gradient checkpointing?

Сильно. Снижает память активаций в 2-3 раза за счёт recompute (~30% медленнее). Включается одной строкой: model.gradient_checkpointing_enable().

Что такое mixed precision?

fp16/bf16 вместо fp32. Память вдвое меньше, скорость на 50-80% выше на современных GPU. Запускается через torch.cuda.amp.autocast() context manager.

Поможет ли flash-attention?

Да, для transformer моделей. Снижает память attention с O(N²) до O(N). На длинных контекстах разница в 4-10×. Установка: pip install flash-attn.