Ошибка CUDA out of memory как уменьшить использование памяти GPU

[docker_kun]

Обучаю трансформер на RTX 3080 (10GB VRAM), батч сайз 16, длина последовательности 512. Через несколько шагов вылетает RuntimeError: CUDA out of memory. Tried to allocate 2.50 GiB. Уменьшил батч до 8 — то же самое. До 4 — стало работать, но это очень медленно и, читал, влияет на качество обучения. Как правильно решать такую проблему?

[k8s_master]

Gradient accumulation — первое что нужно попробовать. Смысл в том, что вы делаете несколько forward+backward пассов без обновления весов, накапливаете градиенты, и только потом вызываете optimizer.step(). Эффективный батч сайз = реальный батч * accumulation_steps. Код: добавляете счётчик шагов, loss.backward() каждый шаг, optimizer.step() + optimizer.zero_grad() каждые N шагов. В PyTorch Lightning это параметр accumulate_grad_batches=4.

[ivan21]

Обязательно включите mixed precision training: from torch.cuda.amp import autocast, GradScaler. Scaler = GradScaler(), и оборачиваете forward pass в with autocast(). Это переводит вычисления в float16 где возможно, память падает почти вдвое, скорость растёт. На RTX 3080 это работает отлично, tensor cores заточены под fp16. В PyTorch 2.x это ещё проще: trainer = Trainer(precision='16-mixed').

[corvet]

Несколько практических трюков которые помогают. Первое: после валидации и в начале эпохи вызывайте torch.cuda.empty_cache() — не освобождает память занятую тензорами, но убирает кэш фрагментов. Второе: убедитесь что у вас нет утечки графов вычислений — если сохраняете loss в список для логгирования, пишите loss.item() а не просто loss, иначе граф не освобождается. Третье: используйте del tensor и затем empty_cache() для временных больших тензоров.

[valru]

Для трансформеров конкретно — посмотрите на gradient checkpointing: model.gradient_checkpointing_enable() в HuggingFace или torch.utils.checkpoint.checkpoint() для кастомных моделей. Суть: не хранить все промежуточные активации в памяти, а пересчитывать их при backward pass. Замедляет обучение на 20-30%, но память сокращается радикально — иногда в 4-10 раз. Для длинных последовательностей это меняет всё.

[nashnet]

Ещё можно профилировать что именно занимает память: print(torch.cuda.memory_summary()) даёт подробную картину. Часто оказывается что проблема не в батче а в том, что накапливаются данные в val_step без torch.no_grad(). Оберните всю валидацию в with torch.no_grad(): и увидите разницу.

[smith_kate]

Если ничего не помогает и нужна скорость — смотрите в сторону DeepSpeed ZeRO-2 или ZeRO-3. Это уже серьёзный инструмент, разбивает состояния оптимизатора и веса по GPU. Для одной карты тоже даёт профит через offload на CPU RAM. Интегрируется с HuggingFace Trainer через одну строчку конфига.

[kube22]

Две вещи, которых в треде не хватает: оптимизатор и фрагментация. AdamW держит два состояния на каждый параметр — в fp32 это легко больше памяти, чем сама модель. AdamW8bit из bitsandbytes срезает это в разы заменой одной строчки, качество у меня не просело. И смотрите внимательно на текст ошибки: если reserved сильно больше allocated — это фрагментация, а не нехватка как таковая, помогает PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True перед запуском.

[thumper416]

@ivan21, дополню: на 3080 (Ampere) я бы сразу брал bf16 вместо fp16 — диапазон как у fp32, GradScaler не нужен и нет ловли NaN на поздних этапах. У меня на fp16 лосс переполнялся примерно к десятой тысяче шагов, та же модель на bf16 доехала до конца без танцев. В Lightning это просто precision='bf16-mixed', экономия памяти та же самая.