Addestramento di Deep Learning su AWS Inferentia

Addestramento Deep Learning su AWS Inferentia

Ancora un hack per risparmiare soldi nell’addestramento dei modelli di IA

Foto di Lisheng Chang su Unsplash

Il tema di questo post è il chip di IA di AWS, AWS Inferentia, sviluppato internamente – più precisamente, la seconda generazione AWS Inferentia2. Questo è un seguito del nostro post dell’anno scorso su AWS Trainium e si unisce a una serie di post sul tema degli acceleratori di IA dedicati. Contrariamente ai chip che abbiamo esplorato nei nostri precedenti post della serie, AWS Inferentia è stato progettato per l’inferenza dei modelli di IA ed è specificamente rivolto alle applicazioni di inferenza di deep learning. Tuttavia, il fatto che AWS Inferentia2 e AWS Trainium condividano entrambi l’architettura sottostante NeuronCore-v2 e lo stesso stack software (l’AWS Neuron SDK), pone la domanda: è possibile utilizzare AWS Inferentia anche per i carichi di lavoro di addestramento dell’IA?

Va detto che ci sono alcuni elementi delle specifiche della famiglia di istanze Amazon EC2 Inf2 (che sono alimentate dagli acceleratori AWS Inferentia) che potrebbero renderle meno adatte per alcuni carichi di lavoro di addestramento rispetto alla famiglia di istanze Amazon EC2 Trn1. Ad esempio, sebbene sia Inf2 che Trn1 supportino un’interconnessione NeuronLink-v2 ad alta larghezza di banda e a bassa latenza, i dispositivi Trainium sono collegati in una topologia a toro 2D anziché in una topologia ad anello, il che potrebbe influire sulle prestazioni degli operatori di comunicazione collettiva (vedi qui per ulteriori dettagli). Tuttavia, alcuni carichi di lavoro di addestramento potrebbero non richiedere le caratteristiche uniche dell’architettura Trn1 e potrebbero funzionare altrettanto bene sulle architetture Inf1 e Inf2.

In effetti, la possibilità di addestrare sia sugli acceleratori Trainium che sugli acceleratori Inferentia aumenterebbe notevolmente la varietà di istanze di addestramento a nostra disposizione e la nostra capacità di regolare la scelta dell’istanza di addestramento alle esigenze specifiche di ciascun progetto di deep learning. Nel nostro recente post, Selezione dell’istanza per il deep learning, abbiamo approfondito il valore di avere una vasta varietà di tipi di istanze diversi per l’addestramento di deep learning. Mentre la famiglia Trn1 include solo due tipi di istanze, l’abilitazione dell’addestramento su Inf2 aggiungerebbe altri quattro tipi di istanze. L’inclusione di Inf1 nella combinazione ne aggiungerebbe altri quattro.

La nostra intenzione in questo post è quella di dimostrare l’opportunità di addestrare su AWS Inferentia. Definiremo un modello di visione di prova e confronteremo le prestazioni dell’addestramento su famiglie di istanze Amazon EC2 Trn1 e Amazon EC2 Inf2. Un grande ringraziamento a Ohad Klein e Yitzhak Levi per il loro contributo a questo post.

Avvertenze

  1. Si noti che, al momento della stesura di questo testo, alcune architetture di modelli di DL non sono supportate dal Neuron SDK. Ad esempio, sebbene il modello di inferenza dei modelli CNN sia supportato, l’addestramento dei modelli CNN non è ancora supportato. La documentazione SDK include una matrice di supporto del modello che illustra le funzionalità supportate per architettura di modelli, framework di addestramento (ad esempio, TensorFlow e PyTorch) e versione dell’architettura Neuron.
  2. Gli esperimenti che descriveremo sono stati eseguiti su Amazon EC2 con l’ultima versione dell’AMI Deep Learning per Neuron disponibile al momento della stesura, “Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 20.04) 20230720”, che include la versione 2.8 del Neuron SDK. Essendo che il Neuron SDK è ancora in fase di sviluppo attivo, è probabile che i risultati comparativi che abbiamo ottenuto cambieranno nel tempo. È altamente consigliato riesaminare le conclusioni di questo post con le versioni più aggiornate delle librerie sottostanti.
  3. La nostra intenzione in questo post è quella di dimostrare il potenziale dell’addestramento su istanze alimentate da AWS Inferentia. Si prega di non considerare questo post come un’approvazione all’uso di queste istanze o di qualsiasi altro prodotto che potremmo menzionare. Ci sono molti fattori da considerare nella scelta di un ambiente di addestramento che possono variare notevolmente in base alle particolarità del tuo progetto. In particolare, modelli diversi potrebbero mostrare risultati di prestazioni relative completamente diversi in esecuzione su due tipi di istanze differenti.

Modello di Prova

Come negli esperimenti descritti nel nostro post precedente, definiamo un modello di classificazione semplice basato su Vision Transformer (ViT) (utilizzando il pacchetto Python timm versione 0.9.5) insieme a un dataset generato casualmente.

from torch.utils.data import Datasetimport time, osimport torchimport torch_xla.core.xla_model as xmimport torch_xla.distributed.parallel_loader as plfrom timm.models.vision_transformer import VisionTransformer

# utilizzo di dati casualiclass FakeDataset(Dataset):  def __len__(self):    return 1000000  def __getitem__(self, index):    rand_image = torch.randn([3, 224, 224], dtype=torch.float32)    label = torch.tensor(data=[index % 1000], dtype=torch.int64)    return rand_image, labeldef train(batch_size=16, num_workers=4):  # Inizializza il gruppo di processi XLA per l'esecuzione di torchrun  import torch_xla.distributed.xla_backend  torch.distributed.init_process_group('xla')  # multiprocessing: assicurarsi che ogni worker abbia gli stessi pesi iniziali  torch.manual_seed(0)  dataset = FakeDataset()  model = VisionTransformer()  # carica il modello sul dispositivo XLA  device = xm.xla_device()  model = model.to(device)  optimizer = torch.optim.Adam(model.parameters())  data_loader = torch.utils.data.DataLoader(dataset,                         batch_size=batch_size, num_workers=num_workers)  data_loader = pl.MpDeviceLoader(data_loader, device)  loss_function = torch.nn.CrossEntropyLoss()  summ, tsumm = 0, 0  count = 0  for step, (inputs, target) in enumerate(data_loader, start=1):    t0 = time.perf_counter()    inputs = inputs.to(device)    targets = torch.squeeze(target.to(device), -1)    optimizer.zero_grad()    outputs = model(inputs)    loss = loss_function(outputs, targets)    loss.backward()    xm.optimizer_step(optimizer)    batch_time = time.perf_counter() - t0    if idx > 10:  # skip first steps      summ += batch_time      count += 1      t0 = time.perf_counter()    if idx > 500:      break  print(f'tempo medio per step: {summ/count}')if __name__ == '__main__':  os.environ['XLA_USE_BF16'] = '1'  # impostare il numero di worker del dataloader in base al numero di vCPU  # ad esempio, 4 per trn1, 2 per inf2.xlarge, 8 per inf2.12xlarge e inf2.48xlarge  train(num_workers=4)# Comando di inizializzazione:# torchrun --nproc_per_node=2 train.py

Risultati

Nella tabella sottostante confrontiamo la velocità e il rapporto prezzo/prestazioni di vari tipi di istanze Amazon EC2 Trn1 e Amazon EC2 Inf2.

Confronto delle prestazioni del modello di classificazione basato su ViT (dall'autore)

Anche se è evidente che i tipi di istanze alimentate da Trainium offrono una migliore prestazione assoluta (ossia, velocità di addestramento aumentate), l’addestramento sulle istanze alimentate da Inferentia ha portato a una prestazione prezzo migliore del ~39% (per i tipi di istanze a due core) e superiore (per i tipi di istanze più grandi).

Ricordiamo ancora una volta di non prendere decisioni di progettazione basate esclusivamente su questi risultati. Alcune architetture dei modelli potrebbero funzionare con successo sulle istanze Trn1 ma non sulle Inf2. Altri potrebbero avere successo su entrambe ma mostrare risultati di prestazione comparativa molto diversi da quelli mostrati qui.

Si noti che abbiamo omesso il tempo necessario per la compilazione del modello DL. Anche se ciò è richiesto solo la prima volta che il modello viene eseguito, i tempi di compilazione possono essere piuttosto lunghi (ad esempio, oltre dieci minuti per il nostro modello di esempio). Due modi per ridurre il sovraccarico della compilazione del modello sono la compilazione parallela e la compilazione offline. È importante assicurarsi che lo script non includa operazioni (o modifiche al grafo) che attivino frequenti ricompilazioni. Consultare la documentazione del Neuron SDK per ulteriori dettagli.

Sommario

Anche se commercializzato come un chip di inferenza AI, sembra che AWS Inferentia offra un’altra opzione per l’addestramento di modelli di deep learning. Nel nostro precedente post su AWS Trainium abbiamo evidenziato alcune delle sfide che potresti incontrare nell’adattare i tuoi modelli per l’addestramento su un nuovo AI ASIC. La possibilità di addestrare gli stessi modelli anche su tipi di istanze alimentate da AWS Inferentia potrebbe aumentare il potenziale beneficio dei tuoi sforzi.