Presentando 🤗 Accelerate

Presentando Accelerate 🤗

🤗 Accelerate

Esegui i tuoi script di addestramento PyTorch grezzi su qualsiasi tipo di dispositivo.

La maggior parte delle librerie di alto livello sopra PyTorch fornisce supporto per l’addestramento distribuito e la precisione mista, ma l’astrazione che introducono richiede all’utente di imparare una nuova API se desidera personalizzare il ciclo di addestramento sottostante. 🤗 Accelerate è stato creato per gli utenti di PyTorch che desiderano avere il pieno controllo sui loro cicli di addestramento ma sono riluttanti a scrivere (e mantenere) il codice boilerplate necessario per utilizzare l’addestramento distribuito (per multi-GPU su uno o più nodi, TPU, …) o l’addestramento a precisione mista. I piani futuri includono il supporto per fairscale, deepseed, data-parallelism specifico di AWS SageMaker e parallelismo del modello.

Fornisce due cose: una semplice e coerente API che astrae quel codice boilerplate e un comando di lancio per eseguire facilmente quegli script su vari ambienti.

Integrazione semplice!

Guardiamo prima un esempio:

  import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from accelerate import Accelerator

+ accelerator = Accelerator()
- device = 'cpu'
+ device = accelerator.device

  model = torch.nn.Transformer().to(device)
  optim = torch.optim.Adam(model.parameters())

  dataset = load_dataset('my_dataset')
  data = torch.utils.data.DataLoader(dataset, shuffle=True)

+ model, optim, data = accelerator.prepare(model, optim, data)

  model.train()
  for epoch in range(10):
      for source, targets in data:
          source = source.to(device)
          targets = targets.to(device)

          optimizer.zero_grad()

          output = model(source)
          loss = F.cross_entropy(output, targets)

-         loss.backward()
+         accelerator.backward(loss)

          optimizer.step()

Aggiungendo solo cinque righe di codice a qualsiasi script di addestramento standard di PyTorch, ora è possibile eseguire detto script su qualsiasi tipo di ambiente distribuito, così come con o senza precisione mista. 🤗 Accelerate gestisce anche il posizionamento del dispositivo per te, in modo da poter semplificare ulteriormente il ciclo di addestramento sopra riportato:

  import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from accelerate import Accelerator

+ accelerator = Accelerator()
- device = 'cpu'

- model = torch.nn.Transformer().to(device)
+ model = torch.nn.Transformer()
  optim = torch.optim.Adam(model.parameters())

  dataset = load_dataset('my_dataset')
  data = torch.utils.data.DataLoader(dataset, shuffle=True)

+ model, optim, data = accelerator.prepare(model, optim, data)

  model.train()
  for epoch in range(10):
      for source, targets in data:
-         source = source.to(device)
-         targets = targets.to(device)

          optimizer.zero_grad()

          output = model(source)
          loss = F.cross_entropy(output, targets)

-         loss.backward()
+         accelerator.backward(loss)

          optimizer.step()

In contrasto, ecco i cambiamenti necessari per far funzionare questo codice con l’addestramento distribuito:

+ import os
  import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from torch.utils.data import DistributedSampler
+ from torch.nn.parallel import DistributedDataParallel

+ local_rank = int(os.environ.get("LOCAL_RANK", -1))
- device = 'cpu'
+ device = device = torch.device("cuda", local_rank)

  model = torch.nn.Transformer().to(device)
+ model = DistributedDataParallel(model)  
  optim = torch.optim.Adam(model.parameters())

  dataset = load_dataset('my_dataset')
+ sampler = DistributedSampler(dataset)
- data = torch.utils.data.DataLoader(dataset, shuffle=True)
+ data = torch.utils.data.DataLoader(dataset, sampler=sampler)

  model.train()
  for epoch in range(10):
+     sampler.set_epoch(epoch)  
      for source, targets in data:
          source = source.to(device)
          targets = targets.to(device)

          optimizer.zero_grad()

          output = model(source)
          loss = F.cross_entropy(output, targets)

          loss.backward()

          optimizer.step()

Questi cambiamenti renderanno il tuo script di addestramento funzionante per più GPU, ma il tuo script smetterà di funzionare su CPU o una sola GPU (a meno che tu non inizi ad aggiungere istruzioni if ovunque). Ancora più fastidioso, se volessi testare il tuo script su TPUs dovresti cambiare diverse righe di codice. Stesso discorso per l’addestramento a precisione mista. La promessa di 🤗 Accelerate è:

  • mantenere al minimo le modifiche al tuo ciclo di addestramento in modo da dover imparare il meno possibile.
  • avere le stesse funzioni che funzionano per qualsiasi configurazione distribuita, in modo da dover imparare solo un’API.

Come funziona?

Per capire come funziona la libreria in pratica, diamo un’occhiata a ciascuna riga di codice che dobbiamo aggiungere ad un ciclo di addestramento.

accelerator = Accelerator()

Oltre a fornire l’oggetto principale che utilizzerai, questa riga analizzerà l’ambiente e il tipo di esecuzione dell’addestramento distribuito, eseguendo l’inizializzazione necessaria. Puoi forzare un addestramento sulla CPU o un addestramento a precisione mista passando cpu=True o fp16=True a questa inizializzazione. Entrambe queste opzioni possono anche essere impostate utilizzando il launcher per il tuo script.

model, optim, data = accelerator.prepare(model, optim, data)

Questo è il nucleo principale dell’API e preparerà i tre principali tipi di oggetti: modelli (torch.nn.Module), ottimizzatori (torch.optim.Optimizer) e dataloader (torch.data.dataloader.DataLoader).

Modello

La preparazione del modello include l’involucro nel contenitore appropriato (ad esempio DistributedDataParallel) e il posizionamento sul dispositivo appropriato. Come con un addestramento distribuito normale, dovrai svolgere il modello per salvarlo o accedere ai suoi metodi specifici, che possono essere fatti con accelerator.unwrap_model(model).

Ottimizzatore

L’ottimizzatore è anche incluso in un contenitore speciale che eseguirà le operazioni necessarie per far funzionare la precisione mista. Gestirà anche correttamente il posizionamento del dispositivo dello stato dict se non è vuoto o caricato da un checkpoint.

Dataloader

Qui è dove la maggior parte della magia è nascosta. Come hai visto nell’esempio di codice, la libreria non si basa su un DistributedSampler, in realtà funzionerà con qualsiasi campionatore tu possa passare al tuo dataloader (se hai mai dovuto scrivere una versione distribuita del tuo campionatore personalizzato, non c’è più bisogno di farlo!). Il dataloader è incluso in un contenitore che prenderà solo gli indici rilevanti per il processo corrente nel campionatore (o salterà i batch per gli altri processi se si utilizza un IterableDataset) e metterà i batch sul dispositivo appropriato.

Per far funzionare questo, Accelerate fornisce una funzione di utilità che sincronizzerà i generatori di numeri casuali su ciascuno dei processi eseguiti durante l’addestramento distribuito. Per impostazione predefinita, sincronizza solo il generator del tuo campionatore, quindi la tua data augmentation sarà diversa su ogni processo, ma il mescolamento casuale sarà lo stesso. Naturalmente, puoi utilizzare questa utility per sincronizzare più RNG se ne hai bisogno.

accelerator.backward(loss)

Questa ultima riga aggiunge i passaggi necessari per il passaggio all’indietro (principalmente per la precisione mista, ma altre integrazioni richiederanno un comportamento personalizzato qui).

Cosa succede con la valutazione?

La valutazione può essere eseguita normalmente su tutti i processi, oppure, se vuoi che venga eseguita solo sul processo principale, puoi utilizzare il comodo test:

if accelerator.is_main_process():
    # Loop di valutazione

Ma puoi anche eseguire facilmente una valutazione distribuita utilizzando Accelerate, ecco cosa dovresti aggiungere al tuo loop di valutazione:

+ eval_dataloader = accelerator.prepare(eval_dataloader)
  predictions, labels = [], []
  for source, targets in eval_dataloader:
      with torch.no_grad():
          output = model(source)

-     predictions.append(output.cpu().numpy())
-     labels.append(targets.cpu().numpy())
+     predictions.append(accelerator.gather(output).cpu().numpy())
+     labels.append(accelerator.gather(targets).cpu().numpy())

  predictions = np.concatenate(predictions)
  labels = np.concatenate(labels)

+ predictions = predictions[:len(eval_dataloader.dataset)]
+ labels = label[:len(eval_dataloader.dataset)]

  metric_compute(predictions, labels)

Come per l’addestramento, devi aggiungere una riga per preparare il tuo dataloader di valutazione. Quindi puoi semplicemente utilizzare accelerator.gather per raccogliere tra processi i tensori di previsioni e etichette. L’ultima riga da aggiungere tronca le previsioni e le etichette al numero di esempi nel tuo dataset perché il dataloader di valutazione preparato restituirà alcuni elementi in più per assicurarsi che tutti i batch abbiano la stessa dimensione su ciascun processo.

Un launcher per dominarli tutti

Gli script che utilizzano Accelerate saranno completamente compatibili con i tuoi launcher tradizionali, come torch.distributed.launch. Tuttavia, ricordare tutti gli argomenti può essere un po’ fastidioso e quando hai configurato la tua istanza con 4 GPU, eseguirai la maggior parte dei tuoi allenamenti utilizzandole tutte. Accelerate viene fornito con una comoda interfaccia a riga di comando (CLI) che funziona in due fasi:

accelerate config

Questo avvierà un piccolo questionario sulla tua configurazione, che creerà un file di configurazione che puoi modificare con tutte le impostazioni predefinite per i tuoi comandi di allenamento. Quindi

accelerate launch path_to_script.py --args_to_the_script

avvierà il tuo script di allenamento utilizzando queste impostazioni predefinite. L’unica cosa che devi fare è fornire tutti gli argomenti necessari per il tuo script di allenamento.

Per rendere questo launcher ancora più fantastico, puoi usarlo per creare un’istanza AWS utilizzando SageMaker. Dai un’occhiata a questa guida per scoprire come!

Come partecipare?

Per iniziare, basta eseguire pip install accelerate o consultare la documentazione per altre opzioni di installazione.

Accelerate è un progetto completamente open source, puoi trovarlo su GitHub, dare un’occhiata alla sua documentazione o dare un’occhiata ai nostri esempi di base. Facci sapere se hai qualche problema o se desideri che la libreria supporti una determinata funzionalità. Per tutte le domande, il forum è il posto giusto da consultare!

Per esempi più complessi in situazioni specifiche, puoi dare un’occhiata agli esempi ufficiali di Transformers. Ogni cartella contiene un file run_task_no_trainer.py che sfrutta la libreria Accelerate!