Trénování modelů s Acceleratorem

Efficient AI Model Training with PyTorch

Dennis Lee

Data Engineer, Amazon

Trainer a Accelerator

Graf porovnávající snadnost použití a možnosti přizpůsobení pro Accelerator a Trainer.

Vlastní trénovací smyčky

Trainer: bez vlastních trénovacích smyček
Některé pokročilé úlohy v generativní AI vyžadují dvě sítě

Obrázek generativní adverzární sítě.

¹ https://www.aitude.com/basics-of-generative-adversarial-network-model/

Trainer a Accelerator

Graf porovnávající snadnost použití a možnosti přizpůsobení pro Accelerator a Trainer.

Úprava základní trénovací smyčky

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

    outputs = model(inputs)

    loss = outputs.loss

    loss.backward()

    optimizer.step()
    scheduler.step()

Vynulování gradientů
Přesun dat na zadané zařízení: .to(device)
Dopředný průchod
Výpočet ztráty křížové entropie
Výpočet gradientů zpětným průchodem
Aktualizace parametrů modelu a rychlosti učení

Vytvoření objektu Accelerator

Accelerator poskytuje rozhraní pro distribuované trénování

from accelerate import Accelerator
accelerator = Accelerator(
    device_placement=True
)

device_placement (bool, výchozí True): Automatické přiřazení zařízení

Definice modelu a optimalizátoru

Načtení předtrénovaného modelu

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-cased", return_dict=True)

Optimalizace parametrů modelu pomocí Adam

from torch.optim import Adam

optimizer = Adam(params=model.parameters(), lr=2e-5)

Definice plánovače

from transformers import get_linear_schedule_with_warmup

lr_scheduler = get_linear_schedule_with_warmup(
      optimizer=optimizer,

      num_warmup_steps=num_warmup_steps,

      num_training_steps=num_training_steps)

optimizer (obj): PyTorchový optimalizátor, např. Adam
num_warmup_steps (int): kroky lineárního zvyšování lr, nastaveno na int(num_training_steps * 0.1)
num_training_steps (int): celkový počet trénovacích kroků, nastaveno na len(train_dataloader) * num_epochs

Příprava modelu pro efektivní trénování

Metoda prepare zajišťuje přiřazení zařízení

model, optimizer, dataloader, lr_scheduler = \
    accelerator.prepare(model,

                        optimizer,

                        dataloader,

                        lr_scheduler)

Sestavení trénovací smyčky s Acceleratorem

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

Vynulování gradientů
Dříve přesunuta data na zařízení

Sestavení trénovací smyčky s Acceleratorem

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

Vynulování gradientů
Dříve přesunuta data na zařízení
Odebrání řádků s ručním přesunem dat

Sestavení trénovací smyčky s Acceleratorem

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)

    loss = outputs.loss
    loss.backward()

Vynulování gradientů
Dříve přesunuta data na zařízení
Odebrání řádků s ručním přesunem dat
Dopředný průchod
Výpočet ztráty křížové entropie a gradientů

Sestavení trénovací smyčky s Acceleratorem

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)
    loss = outputs.loss

    accelerator.backward(loss)

    optimizer.step()
    scheduler.step()

Vynulování gradientů
Dříve přesunuta data na zařízení
Odebrání řádků s ručním přesunem dat
Dopředný průchod
Výpočet ztráty křížové entropie a gradientů
Nahrazení loss.backward pomocí accelerator
Aktualizace parametrů modelu a rychlosti učení

Přehled změn

Před Acceleratorem

Ruční přesun dat na zařízení
- inputs.to(device)
- targets.to(device)
Výpočet gradientů pomocí loss.backward()

Po Acceleratoru

Automatické přiřazení zařízení a datový paralelismus
- accelerator.prepare(model)
- accelerator.prepare(dataloader)
Synchronizace gradientů pomocí accelerator.backward(loss)
Přizpůsobitelná smyčka
Uživatelsky přívětivý, hardwarově nezávislý, škálovatelný a udržovatelný

Lass uns üben!

Efficient AI Model Training with PyTorch