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

[danila_spb]

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

[lev_loop]

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

[ruslan_data]

Обязательно включите 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').

[irina2236]

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

[maria_api21]

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

[apiflow8899]

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

[roman_lab]

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