Addestra i modelli con Accelerator

Efficient AI Model Training with PyTorch

Dennis Lee

Data Engineer, Amazon

Trainer e Accelerator

Grafico: facilità d'uso vs. possibilità di personalizzazione per Accelerator e Trainer.

Loop di training personalizzati

Trainer: niente loop di training personalizzati
Alcuni task avanzati di generative AI richiedono due reti

Immagine di una generative adversarial network.

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

Trainer e Accelerator

Grafico: facilità d'uso vs. possibilità di personalizzazione per Accelerator e Trainer.

Modificare un training loop di base

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

Azzera i gradienti
Sposta i dati su un device: .to(device)
Esegui il forward pass
Calcola la cross-entropy loss
Calcola i gradienti nel backward pass
Aggiorna parametri e learning rate

Crea un oggetto Accelerator

Accelerator offre un'interfaccia per il training distribuito

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

device_placement (bool, default True): gestisce il posizionamento sui device

Definisci modello e ottimizzatore

Carica un modello pre-addestrato

from transformers import AutoModelForSequenceClassification

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

Ottimizza i parametri con Adam

from torch.optim import Adam

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

Definisci lo scheduler

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): ottimizzatore PyTorch, ad es. Adam
num_warmup_steps (int): passi per aumentare linearmente lr, imposta int(num_training_steps * 0.1)
num_training_steps (int): passi totali di training, imposta len(train_dataloader) * num_epochs

Prepara il modello per un training efficiente

Il metodo prepare gestisce il posizionamento sui device

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

                        optimizer,

                        dataloader,

                        lr_scheduler)

Costruire un training loop con Accelerator

for batch in dataloader:

    optimizer.zero_grad()

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

Azzera i gradienti
Dati già spostati sul device

Costruire un training loop con Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

Azzera i gradienti
Dati già spostati sul device
Rimuovi le righe che spostano i dati manualmente

Costruire un training loop con Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)

    loss = outputs.loss
    loss.backward()

Azzera i gradienti
Dati già spostati sul device
Rimuovi le righe che spostano i dati manualmente
Esegui il forward pass
Calcola cross-entropy loss e gradienti

Costruire un training loop con Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)
    loss = outputs.loss

    accelerator.backward(loss)

    optimizer.step()
    scheduler.step()

Azzera i gradienti
Dati già spostati sul device
Rimuovi le righe che spostano i dati manualmente
Esegui il forward pass
Calcola cross-entropy loss e gradienti
Sostituisci loss.backward con accelerator
Aggiorna parametri e learning rate

Riepilogo delle modifiche

Prima di Accelerator

Spostare manualmente i dati sui device
- inputs.to(device)
- targets.to(device)
Calcolare i gradienti con loss.backward()

Dopo Accelerator

Posizionamento automatico e data parallelism
- accelerator.prepare(model)
- accelerator.prepare(dataloader)
Sincronizza i gradienti con accelerator.backward(loss)
Loop personalizzabile
Semplice, agnostico all'hardware, scalabile e mantenibile

Ayo berlatih!

Efficient AI Model Training with PyTorch