Как да поправите грешката „PyTorch CUDA Out of Memory“ по време на обучение на модел

Едно обучение с PyTorch може да работи няколко стъпки, след което да спре с torch.OutOfMemoryError или съобщение като CUDA out of memory. Tried to allocate .... Непосредствената причина е проста: следващото CUDA разпределение на памет не е могло да се побере. Полезният въпрос е защо не е могло да се побере.

По време на обучението паметта на GPU може да съдържа параметри на модела, градиенти, състояние на оптимизатора, входни тензори, временни работни пространства и активации, запазени за обратното разпространение. PyTorch използва и кеш аллокатор, така че паметта, показана като „reserved“ (резервирана), не е идентична с паметта, която в момента е заета от живи тензори. Това разграничение е важно, когато решавате дали да намалите натоварването, или да разследвате фрагментацията на аллокатора.

Това ръководство следва текущата документация на 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()

# Executirajte edna predstavitelna stъpka na obuchenie tuk.

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. Чест модел е съхраняването на тензори, свързани с графа, в списък:

# Risky if retained for many steps:
loss_history.append(loss)

# Store a Python number instead:
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.

# Example for a workload with varying allocation sizes:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Не копирайте флагове на аллокатора от друга машина, без да проверите memory_summary() или снимка. Истински проблем с капацитета—където живите тензори вече запълват GPU—няма да бъде решен чрез настройка на фрагментацията.

Когато един GPU все още не може да побере модела

Ако един пример при най-малкия практичен размер на входа все още предизвиква OOM, проблемът може да е в модела и състоянието на оптимизатора, а не в партидата. В този момент помислете за по-малка архитектура, параметри с по-ниска точност, където е числено уместно, стратегии за CPU/офлоуд или разпределено обучение с шардиране.

Fully Sharded Data Parallel (FSDP) на PyTorch може да шардира параметрите на модела между работници за паралелизъм на данните, а неговата стратегия FULL_SHARD също шардира градиентите и състоянията на оптимизатора. Това може да намали паметта на GPU в сравнение с напълно реплициран паралелизъм на данните, за сметка на комуникацията и по-сложното поведение при обучение. Документация за FSDP на PyTorch.

Практически ред на операциите

Контролен списък: как да разберете, че OOM е наистина поправен

• Изпълнете няколко представителни итерации на обучение, а не само едно успешно правото преминаване.
• Нулирайте и запишете max_memory_allocated(), за да знаете новия пик.
• Потвърдете, че паметта на GPU достига стабилен диапазон, вместо да се увеличава при всяка итерация.
• Валидирайте загубата и градиентите след активиране на смесената точност.
• Потвърдете, че акумулирането на градиенти запазва графика за актуализация на оптимизатора, който сте предвидили.
• Изпълнете преминаване за валидиране под torch.no_grad(), когато градиентите не са необходими.
• Ако сте променили настройките на аллокатора, сравнете статистиката за паметта и пропускливостта преди и след това.
• Не считайте проблема за решен само защото nvidia-smi показва по-малко резервирана памет след empty_cache(); самото натоварване при обучение трябва да завърши при нормалния си пик.

Най-добре е CUDA OOM да се разглежда като проблем с бюджета на паметта, а не като единичен бъг в PyTorch. Измерете пика, първо намалете живия работен набор, след което използвайте смесена точност, акумулиране и контролни точки като умишлени размени. Преминавайте към настройка на аллокатора само когато статистиката на аллокатора сочи фрагментация, и преминавайте към шардиране или различен модел, когато самият модел вече не се побира удобно на един GPU.