Як виправити помилку «PyTorch CUDA Out of Memory» під час навчання моделі

Запуск навчання в PyTorch може працювати кілька кроків, а потім зупинитися з помилкою torch.OutOfMemoryError або повідомленням на кшталт CUDA out of memory. Tried to allocate .... Безпосередня причина проста: наступне виділення пам'яті CUDA не вмістилося. Корисне питання полягає в тому, чому воно не вмістилося.

Під час навчання пам'ять GPU може містити параметри моделі, градієнти, стан оптимізатора, вхідні тензори, тимчасові робочі області та активації, збережені для зворотного поширення. PyTorch також використовує кешуючий аллокатор, тому пам'ять, що відображається як «зарезервована», не тотожна пам'яті, яку наразі займають активні тензори. Це розрізнення важливе при вирішенні, чи зменшувати навантаження, чи досліджувати фрагментацію аллокатора.

Цей посібник базується на актуальній документації PyTorch і використовує сучасні назви API AMP. Зокрема, PyTorch тепер документує torch.amp.autocast("cuda") та torch.amp.GradScaler("cuda"); старі точки входу torch.cuda.amp.* є застарілими. Див. документацію PyTorch щодо автоматичної змішаної точності.

Швидка діагностика: з яким типом OOM ви маєте справу?

Крок 1: Виміряйте пам'ять перед зміною рецептури навчання

Почніть із запису розміру пакету, розмірності вхідних даних, точності та моменту, коли відбувається збій. Потім перевірте пам'ять активних тензорів та пам'ять, зарезервовану аллокатором. PyTorch надає функції memory_allocated(), memory_reserved(), пікові варіанти та memory_summary(). Поточна документація з управління пам'яттю CUDA пояснює, що кешуючий аллокатор зберігає блоки, які можна повторно використовувати, саме тому невикористана зарезервована пам'ять може все ще відображатися як зайнята в інструментах моніторингу GPU. Управління пам'яттю CUDA в PyTorch.

import torch

torch.cuda.reset_peak_memory_stats()

# Виконайте тут один репрезентативний крок навчання.

print("allocated GB:",
      torch.cuda.memory_allocated() / 1024**3)
print("reserved GB:",
      torch.cuda.memory_reserved() / 1024**3)
print("peak allocated GB:",
      torch.cuda.max_memory_allocated() / 1024**3)
print(torch.cuda.memory_summary(abbreviated=True))

Якщо простого резюме недостатньо, PyTorch може створювати знімки аллокатора для глибшого аналізу. Його інструменти пам'яті можуть записувати історію виділень і створювати знімок, який можна переглянути за допомогою візуалізатора пам'яті PyTorch. PyTorch зазначає, що ці інструменти бачать пам'ять, керовану аллокатором PyTorch; виділення, зроблені безпосередньо іншими бібліотеками CUDA, можуть там не відображатися. Посібник PyTorch щодо розуміння використання пам'яті CUDA.

Не розглядайте torch.cuda.empty_cache() як загальне рішення для OOM

torch.cuda.empty_cache() звільняє невикористані кешовані блоки, щоб інші додатки GPU могли їх використовувати. PyTorch явно зазначає, що це не звільняє пам'ять, зайняту активними тензорами, і тому не збільшує обсяг пам'яті GPU, доступний для PyTorch для тензорів, які все ще живі. Це може бути корисним між окремими експериментами або після видалення великих об'єктів, але це не замінює зменшення активного обсягу пам'яті.

Крок 2: Спочатку зменшіть активний робочий набір

Найнадійнішим першим виправленням зазвичай є менший мікро-пакет: кількість зразків, оброблених одним прямим/зворотним проходженням. Пам'ять активацій зазвичай зростає з розміром пакету, роздільною здатністю зображення, довжиною послідовності та іншими розмірностями вхідних даних. Якщо модель навчається при розмірі пакету 32, але не працює при 64, зменшення пакету не є обхідним шляхом у негативному сенсі; це пряме зменшення пікової потреби в пам'яті.

Для зображень зниження просторової роздільної здатності або розміру обрізки може суттєво вплинути. Для трансформерів та інших моделей послідовностей зменшення довжини послідовності може бути ще важливішим, оскільки деякі проміжні тензори сильно зростають з довжиною послідовності. Точне масштабування залежить від архітектури, тому вимірюйте, а не припускайте.

Також переконайтеся, що код оцінювання не будує градієнти без потреби. Рекомендації щодо продуктивності PyTorch радять вимикати обчислення градієнтів для валідації або висновку, коли градієнти не потрібні, оскільки autograd інакше зберігає проміжні буфери. Типовий шаблон:

model.eval()
with torch.no_grad():
    for x, y in val_loader:
        x = x.cuda(non_blocking=True)
        y = y.cuda(non_blocking=True)
        pred = model(x)

Під час навчання використовуйте optimizer.zero_grad(set_to_none=True), якщо ваш алгоритм не покладається на поведінкову різницю між нульовим градієнтом та градієнтом None. Документація оптимізатора PyTorch зазначає, що встановлення градієнтів у None зазвичай має менший обсяг пам'яті та може незначно покращити продуктивність. Документація zero_grad оптимізатора PyTorch.

Крок 3: Зберігайте більший ефективний пакет за допомогою AMP та акумуляції градієнтів

Використовуйте автоматичну змішану точність, якщо модель це підтримує

Автоматична змішана точність (AMP) виконує придатні операції з нижчою точністю, зберігаючи операції, які потребують більшого діапазону або точності, у відповідних типах. PyTorch документує, що AMP може покращити продуктивність і зменшити обсяг пам'яті для багатьох навантажень CUDA, але він не є числово придатним для кожної моделі. Зокрема, PyTorch попереджає, що деякі моделі, попередньо навчені в bfloat16, можуть переповнюватися у float16.

Актуальний шаблон навчання CUDA AMP:

scaler = torch.amp.GradScaler("cuda")

for inputs, targets in train_loader:
    inputs = inputs.cuda(non_blocking=True)
    targets = targets.cuda(non_blocking=True)
    optimizer.zero_grad(set_to_none=True)

    with torch.amp.autocast("cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

PyTorch рекомендує виконувати пряме проходження та функцію втрат у контексті autocast, а потім виходити з контексту autocast перед зворотним проходженням. Якщо float16 спричиняє нестабільність, дослідіть, чи підтримується bfloat16 і чи є він доречним для вашого апаратного забезпечення та моделі, замість того, щоб припускати, що всі режими змішаної точності поводяться однаково.

Використовуйте акумуляцію градієнтів, коли потрібен більший ефективний пакет

Акумуляція градієнтів обробляє кілька менших мікро-пакетів перед оновленням оптимізатора. Якщо мікро-пакет дорівнює 4, і ви акумулятуєте 8 кроків, ефективний пакет для одного оновлення оптимізатора становить 32 зразки на воркера, за умови, що кожен мікро-пакет має чотири зразки, а налаштування паралелізму даних не змінює цю арифметику.

accum_steps = 8
optimizer.zero_grad(set_to_none=True)

for step, (inputs, targets) in enumerate(train_loader):
    inputs = inputs.cuda(non_blocking=True)
    targets = targets.cuda(non_blocking=True)

    with torch.amp.autocast("cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets) / accum_steps

    scaler.scale(loss).backward()

    if (step + 1) % accum_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)

Для виробничого коду також обробляйте фінальне неповне вікно акумуляції, коли кількість пакетів не ділиться на accum_steps. Якщо ви використовуєте розподілене навчання, поведінка синхронізації градієнтів може змінити компроміс пам'ять/продуктивність, тому дотримуйтесь рекомендацій щодо акумуляції розподіленого API, а не копіюйте цикл для одного GPU без змін.

Крок 4: Обміняйте обчислення на пам'ять, потім дослідіть утримання та фрагментацію

Контрольні точки активацій

Контрольні точки активацій зменшують пам'ять, не зберігаючи вибрані активації прямого проходження живими до зворотного проходження. Замість цього PyTorch перераховує їх під час зворотного проходження. Це обмінює додаткові обчислення на нижчий обсяг пам'яті для активацій. Поточна документація PyTorch щодо контрольних точок рекомендує явно передавати use_reentrant=False. Документація PyTorch щодо контрольних точок активацій.

from torch.utils.checkpoint import checkpoint

def forward(self, x):
    x = checkpoint(self.block1, x, use_reentrant=False)
    x = checkpoint(self.block2, x, use_reentrant=False)
    return self.head(x)

Використовуйте контрольні точки для шарів з великими збереженими активаціями та прийнятною вартістю перерахунку. Не припускайте, що контрольні точки для кожної операції є оптимальними; це може суттєво сповільнити навчання.

Шукайте тензори, які утримують графи обчислень живими

Якщо пам'ять зростає з кожною ітерацією, а не досягає піку приблизно на тому ж рівні, перевірте посилання Python. Типовий шаблон — зберігання тензорів, пов'язаних з графом, у списку:

# Ризиковано, якщо зберігається багато кроків:
loss_history.append(loss)

# Зберігайте число Python замість цього:
loss_history.append(loss.item())

Та сама проблема може виникати, коли ви кешуєте вихідні дані моделі, карти уваги, приховані стани або тензори валідації без від'єднання або переміщення їх з GPU. Видаляйте посилання, які вам більше не потрібні, і використовуйте detach() лише тоді, коли ви навмисно хочете отримати тензор, від'єднаний від autograd.

Налаштовуйте аллокатор лише після того, як статистика вказує на фрагментацію

Поточна документація PyTorch віддає перевагу змінній середовища PYTORCH_ALLOC_CONF. Старіша PYTORCH_CUDA_ALLOC_CONF залишається псевдонімом для зворотної сумісності. Ця деталь щодо іменування змінилася в поточній документації, тому нові конфігурації мають використовувати переважну назву. Змінні середовища CUDA в PyTorch.

Два параметри аллокатора особливо актуальні:

• expandable_segments:True є експериментальним і призначений для зменшення непридатних шматків пам'яті, коли розміри виділень змінюються, наприклад, у навантаженнях, де розміри пакетів або тензорів варіюються.
• max_split_size_mb може зменшити фрагментацію за допомогою нативного аллокатора, але PyTorch явно описує його як останній засіб для навантажень, які зазнають збоїв через OOM, демонструючи велику кількість неактивних розділених блоків. Він також може погіршити продуктивність і ігнорується бекендом cudaMallocAsync.

# Приклад для навантаження зі змінними розмірами виділень:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Не копіюйте прапорці аллокатора з іншої машини без перевірки memory_summary() або знімка. Справжня проблема з місткістю — де активні тензори вже заповнюють GPU — не буде вирішена налаштуванням фрагментації.

Коли один GPU все ще не може вмістити модель

Якщо один зразок при найменшому практичному розмірі вхідних даних все ще спричиняє OOM, проблема може бути в моделі та стані оптимізатора, а не в пакеті. У цьому випадку розгляньте меншу архітектуру, параметри з нижчою точністю там, де це числово доречно, стратегії CPU/offload або розподілене навчання з шардуванням.

Fully Sharded Data Parallel (FSDP) у PyTorch може шардувати параметри моделі між воркерами паралелізму даних, а його стратегія FULL_SHARD також шардує градієнти та стани оптимізатора. Це може зменшити пам'ять на GPU порівняно з повністю реплікованим паралелізмом даних, ціною комунікації та складнішої поведінки навчання. Документація PyTorch FSDP.

Практичний порядок дій

Чек-лист: як дізнатися, що OOM дійсно виправлено

• Запустіть кілька репрезентативних ітерацій навчання, а не лише одне успішне пряме проходження.
• Скиньте та запишіть max_memory_allocated(), щоб знати новий пік.
• Переконайтеся, що пам'ять GPU досягає стабільного діапазону, замість того щоб зростати з кожною ітерацією.
• Перевірте втрати та градієнти після увімкнення змішаної точності.
• Переконайтеся, що акумуляція градієнтів зберігає запланований вами графік оновлення оптимізатора.
• Запустіть проходження валідації в контексті torch.no_grad(), коли градієнти не потрібні.
• Якщо ви змінили налаштування аллокатора, порівняйте статистику пам'яті та пропускну здатність до і після.
• Не вважайте проблему вирішеною лише тому, що nvidia-smi показує менше зарезервованої пам'яті після empty_cache(); саме навчальне навантаження має завершуватися при своєму звичайному піку.

OOM CUDA найкраще розглядати як проблему бюджету пам'яті, а не як одну помилку PyTorch. Виміряйте пік, спочатку зменшіть живий робочий набір, потім використовуйте змішану точність, акумуляцію та контрольні точки як свідомі компроміси. Переходьте до налаштування аллокатора лише тоді, коли статистика аллокатора вказує на фрагментацію, і переходьте до шардування або іншої моделі, коли сама модель більше не вміщується комфортно на одному GPU.