Ошибка переполнения памяти CUDA, но есть свободное место ⇐ Python

[Anonymous]

У меня возникла эта проблема некоторое время, когда я пытался удвоить размер встраивания с 64 до 128 или увеличить размер пакета с 1, я продолжал получать эту ошибку. Это языковая модель с 1,2 миллионами параметров. Он работает с 300 тысячами параметров с одними и теми же данными, но если я использую больший набор данных с 200 тысячами выборок, происходит сбой. Текущий набор данных содержит 4 тыс. выборок. Образцы представляют собой последовательности из 50 белков, наложенных друг на друга. Буду признателен, если кто-нибудь поможет.
Код ошибки:

Код: Выделить всё

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 150.00 MiB (GPU 0; 31.74 GiB total capacity; 21.32 GiB already allocated; 9.86 GiB free; 21.58 GiB allowed; 21.37 GiB reserved in total by PyTorch)

Вот аргументы:

Код: Выделить всё

nb_blocks=6, embed_dim=128, nb_heads=4, nb_epochs=20, warmup_steps=300, learning_rate=0.0001, check_val_every=1000, batch_size=1

Я использую кластерную настройку с использованием графических процессоров NVIDIA V100 PCIe 32 ГБ.
Настройки Lightning:

Код: Выделить всё

slurm_args = {
"accelerator": "gpu",
"devices": int(os.environ["SLURM_GPUS_ON_NODE"]),
"num_nodes": int(os.environ["SLURM_NNODES"]),
"strategy": "ddp",
"precision": 16,
}

...

accelerator = "cuda" if torch.cuda.is_available() else "cpu"
trainer_args = {
"max_epochs": args.nb_epochs,
"log_every_n_steps": LOGGING_STEPS,
"val_check_interval": VAL_CHECK_STEPS,
"logger": wandb_logger,
"callbacks": callbacks,
"accelerator": accelerator,
**slurm_args,
}

Что я пробовал:

• Установка точности до 16 (вместо 32): вдвое уменьшилось переполнение памяти с 300 МБ. до 150 МБ (хотя бы некоторый прогресс, но проблему не решил)

• Код: Выделить всё

  torch.cuda.set_per_process_memory_fraction(0.7)

  : место освобождено, но оно не было выделено для использования при переполнении памяти, и я не могу понять, как выделить эту свободную память для решения вышеуказанной проблемы.
• Ручное переопределение распределения и перемещение весов, данных и модели в ЦП. Ничего не делал, так как почти вся инициализация уже инициализируется процессором.
• Установите акселератор в slurm_args на CPU. Решена проблема с памятью, но каждая итерация занимала от 1,2 до 25 секунд. Другими словами, это настолько медленно, что бесполезно
• Изменил стратегию на FSDPStrategy(cpu_offload=True) с ddp и ничего не изменилось
  Пытался установить PYTORCH_CUDA_ALLOC_CONF, ничего не дало
• Установил Assessment_grad_batches в аргументах с разными значениями, без изменений

Подробнее здесь: https://stackoverflow.com/questions/792 ... free-space