Träna modeller med Accelerator

Effektiv AI-modellträning med PyTorch

Dennis Lee

Data Engineer, Amazon

Trainer och Accelerator

Diagram som jämför användarvänlighet och anpassningsbarhet för Accelerator och Trainer.

Anpassade träningsloopar

Trainer: inga anpassade träningsloopar
Vissa avancerade uppgifter inom generativ AI kräver två nätverk

Bild av ett generativt adversariellt nätverk.

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

Trainer och Accelerator

Diagram som jämför användarvänlighet och anpassningsbarhet för Accelerator och Trainer.

Modifiera en grundläggande träningsloop

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()

Nollställ gradienterna
Flytta data till angiven enhet: .to(device)
Utför framåtpass
Beräkna korsentriförlust
Beräkna gradienter i bakåtpass
Uppdatera modellparametrar och inlärningshastighet

Skapa ett Accelerator-objekt

Accelerator tillhandahåller ett gränssnitt för distribuerad träning

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

device_placement (bool, standard True): Hanterar enhetsplacering automatiskt

Definiera modell och optimerare

Läs in en förtränad modell

from transformers import AutoModelForSequenceClassification

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

Optimera modellparametrar med Adam

from torch.optim import Adam

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

Definiera schemaläggaren

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): PyTorch-optimerare, t.ex. Adam
num_warmup_steps (int): steg för linjär ökning av lr, sätts till int(num_training_steps * 0.1)
num_training_steps (int): totalt antal träningssteg, sätts till len(train_dataloader) * num_epochs

Förbered modellen för effektiv träning

Metoden prepare hanterar enhetsplacering

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

                        optimizer,

                        dataloader,

                        lr_scheduler)

Bygga en träningsloop med Accelerator

for batch in dataloader:

    optimizer.zero_grad()

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

Nollställ gradienterna
Data har tidigare flyttats till enheten

Bygga en träningsloop med Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

Nollställ gradienterna
Data har tidigare flyttats till enheten
Ta bort rader som manuellt flyttar data

Bygga en träningsloop med Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)

    loss = outputs.loss
    loss.backward()

Nollställ gradienterna
Data har tidigare flyttats till enheten
Ta bort rader som manuellt flyttar data
Utför ett framåtpass
Beräkna korsentriförlust och gradienter

Bygga en träningsloop med Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)
    loss = outputs.loss

    accelerator.backward(loss)

    optimizer.step()
    scheduler.step()

Nollställ gradienterna
Data har tidigare flyttats till enheten
Ta bort rader som manuellt flyttar data
Utför ett framåtpass
Beräkna korsentriförlust och gradienter
Ersätt loss.backward med accelerator
Uppdatera modellparametrar och inlärningshastighet

Sammanfattning av ändringar

Före Accelerator

Data måste flyttas manuellt till enheter
- inputs.to(device)
- targets.to(device)
Beräkna gradienter med loss.backward()

Efter Accelerator

Automatisk enhetsplacering och dataparallellism
- accelerator.prepare(model)
- accelerator.prepare(dataloader)
Gradientssynkronisering med accelerator.backward(loss)
Anpassningsbar loop
Användarvänlig, hårdvaruoberoende, skalbar och underhållbar

Lass uns üben!

Effektiv AI-modellträning med PyTorch