Massimizza le prestazioni e riduci i costi di formazione per il deep learning con AWS Trainium e Amazon SageMaker

Massimizza le prestazioni e riduci i costi di formazione per il deep learning con AWS Trainium e Amazon SageMaker

Nodo di origine: 2010875

Oggi, decine di migliaia di clienti stanno creando, addestrando e implementando modelli di machine learning (ML) utilizzando Amazon Sage Maker per potenziare le applicazioni che hanno il potenziale per reinventare le loro attività e le esperienze dei clienti. Questi modelli ML sono aumentati in dimensioni e complessità negli ultimi anni, il che ha portato a precisioni all'avanguardia in una vasta gamma di attività e ha anche ridotto i tempi di formazione da giorni a settimane. Di conseguenza, i clienti devono ridimensionare i propri modelli su centinaia o migliaia di acceleratori, il che li rende più costosi da addestrare.

SageMaker è un servizio ML completamente gestito che aiuta sviluppatori e data scientist a creare, addestrare e distribuire facilmente modelli ML. SageMaker fornisce già la scelta più ampia e completa di offerte di elaborazione con acceleratori hardware per la formazione ML, tra cui G5 (Nvidia A10G) istanze e P4d (Nvidia A100) le istanze.

I crescenti requisiti di elaborazione richiedono una potenza di elaborazione più rapida e conveniente. Per ridurre ulteriormente i tempi di addestramento del modello e consentire ai professionisti del machine learning di iterare più velocemente, AWS ha innovato chip, server e connettività del data center. Le nuove istanze Trn1 alimentate da AWSTrainium i chip offrono il miglior rapporto prezzo-prestazioni e l'addestramento del modello ML più veloce su AWS, fornendo fino al 50% in meno di costi per l'addestramento di modelli di deep learning su istanze basate su GPU comparabili senza alcun calo di precisione.

In questo post, mostriamo come massimizzare le prestazioni e ridurre i costi utilizzando le istanze Trn1 con SageMaker.

Panoramica della soluzione

I lavori di formazione SageMaker supportano le istanze ml.trn1, basate su chip Trainium, creati appositamente per applicazioni di formazione ML ad alte prestazioni nel cloud. Puoi utilizzare le istanze ml.trn1 su SageMaker per addestrare l'elaborazione del linguaggio naturale (NLP), la visione artificiale e i modelli di raccomandazione in un vasta gamma di applicazioni, come il riconoscimento vocale, la raccomandazione, il rilevamento di frodi, la classificazione di immagini e video e le previsioni. Le istanze ml.trn1 includono fino a 16 chip Trainium, che è un chip ML di seconda generazione creato da AWS dopo AWS Inferenza. Le istanze ml.trn1 sono le prime Cloud di calcolo elastico di Amazon (Amazon EC2) fino a 800 Gbps di larghezza di banda di rete EFA (Elastic Fabric Adapter). Per un parallelismo efficiente di dati e modelli, ogni istanza ml.trn1.32xl dispone di 512 GB di memoria a larghezza di banda elevata, offre fino a 3.4 petaflop di potenza di elaborazione FP16/BF16 e presenta NeuronLink, un'interconnessione intra-istanza, a larghezza di banda elevata e non bloccante .

Trainium è disponibile in due configurazioni e può essere utilizzato nelle regioni Stati Uniti orientali (Virginia settentrionale) e Stati Uniti occidentali (Oregon).

La tabella seguente riassume le caratteristiche delle istanze Trn1.

Dimensione istanza Trainio
acceleratori
Acceleratore
Memorie
(GB)
CPU virtuali Instance
Memorie
(GiB)
Network NetPoulSafe
Larghezza di banda
(Gbps)
EFA e
RDMA
Assistenza
trn1.2xgrande 1 32 8 32 Fino a 12.5 Non
trn1.32xgrande 16 512 128 512 800
trn1n.32xlarge (prossimamente) 16 512 128 512 1600

Capiamo come utilizzare Trainium con SageMaker con un semplice esempio. Addestreremo un modello di classificazione del testo con SageMaker training e PyTorch utilizzando la Hugging Face Transformers Library.

Utilizziamo il set di dati Amazon Reviews, che consiste in recensioni di amazon.com. I dati coprono un periodo di 18 anni, comprendendo circa 35 milioni di recensioni fino a marzo 2013. Le recensioni includono informazioni sul prodotto e sull'utente, valutazioni e una recensione in chiaro. Il codice seguente è un esempio tratto da AmazonPolarity set di prova:

{
title':'Great CD', 'content':"My lovely Pat has one of the GREAT voices of her generation. I have listened to this CD for YEARS and I still LOVE IT. When I'm in a good mood it makes me feel better. A bad mood just evaporates like sugar in the rain. This CD just oozes LIFE. Vocals are jusat STUUNNING and lyrics just kill. One of life's hidden gems. This is a desert isle CD in my book. Why she never made it big is just beyond me. Everytime I play this, no matter black, white, young, old, male, female EVERYBODY says one thing ""Who was that singing ?""", 'label':1
}

Per questo post, utilizziamo solo i campi contenuto ed etichetta. Il campo contenuto è una recensione a testo libero e il campo etichetta è un valore binario contenente 1 o 0 rispettivamente per recensioni positive o negative.

Per il nostro algoritmo, utilizziamo BERT, un modello di trasformatore pre-addestrato su un ampio corpus di dati inglesi in modo auto-supervisionato. Questo modello ha principalmente lo scopo di essere messo a punto su attività che utilizzano l'intera frase (potenzialmente mascherata) per prendere decisioni, come la classificazione della sequenza, la classificazione dei token o la risposta alle domande.

Dettagli di implementazione

Cominciamo dando un'occhiata più da vicino ai diversi componenti coinvolti nell'addestramento del modello:

  • AWSTrainium – Al centro, ciascuno Istanza Trainium ha dispositivi Trainium integrati. Trn1.2xlarge ha 1 dispositivo Trainium e Trn1.32xlarge ha 16 dispositivi Trainium. Ogni dispositivo Trainium è costituito da computer (2 NeuronCore-v2), 32 GB di memoria del dispositivo HBM e NeuronLink per una rapida comunicazione tra dispositivi. Ogni NeuronCore-v2 consiste in un'unità di calcolo eterogenea completamente indipendente con motori separati (Tensor/Vector/Scalar/GPSIMD). I GPSIMD sono processori generici completamente programmabili che puoi utilizzare per implementare operatori personalizzati ed eseguirli direttamente sui motori NeuronCore.
  • Formazione su Amazon SageMaker – SageMaker offre un'esperienza di formazione completamente gestita per addestrare facilmente i modelli senza doversi preoccupare dell'infrastruttura. Quando utilizzi SageMaker Training, esegue tutto il necessario per un lavoro di addestramento, come codice, contenitore e dati, in un'infrastruttura di elaborazione separata dall'ambiente di chiamata. Questo ci consente di eseguire esperimenti in parallelo e iterare velocemente. SageMaker fornisce un SDK Python avviare lavori di formazione. L'esempio in questo post utilizza SageMaker Python SDK per attivare il processo di formazione utilizzando Trainium.
  • Neurone AWS – Poiché Trainium NeuronCore ha il proprio motore di calcolo, abbiamo bisogno di un meccanismo per compilare il nostro codice di addestramento. IL Neurone AWS compilatore prende il codice scritto in Pytorch/XLA e lo ottimizza per funzionare su dispositivi Neuron. Il compilatore Neuron è integrato come parte del Deep Learning Container che useremo per addestrare il nostro modello.
  • PyTorch/XLA - Questo Pacchetto Python utilizza il compilatore di deep learning XLA per connettere il framework di deep learning PyTorch e gli acceleratori cloud come Trainium. La creazione di una nuova rete PyTorch o la conversione di una esistente per l'esecuzione su dispositivi XLA richiede solo poche righe di codice specifico per XLA. Vedremo per il nostro caso d'uso quali modifiche dobbiamo apportare.
  • Formazione distribuita – Per eseguire l'addestramento in modo efficiente su più NeuronCores, abbiamo bisogno di un meccanismo per distribuire l'addestramento nei NeuronCores disponibili. SageMaker supporta torchrun con istanze Trainium, che possono essere utilizzate per eseguire più processi equivalenti al numero di NeuronCore nel cluster. Questo viene fatto passando il parametro di distribuzione allo stimatore SageMaker come segue, che avvia un addestramento distribuito in parallelo di dati in cui lo stesso modello viene caricato in diversi NeuronCores che elaborano batch di dati separati:
distribution={"torch_distributed": {"enabled": True}}

Modifiche allo script necessarie per l'esecuzione su Trainium

Diamo un'occhiata alle modifiche al codice necessarie per adottare un normale script PyTorch basato su GPU da eseguire su Trainium. Ad alto livello, dobbiamo apportare le seguenti modifiche:

  1. Sostituisci i dispositivi GPU con dispositivi Pytorch/XLA. Poiché utilizziamo la distribuzione della torcia, dobbiamo inizializzare l'addestramento con XLA come dispositivo come segue:
    device = "xla"
    torch.distributed.init_process_group(device)

  2. Utilizziamo il backend distribuito PyTorch/XLA per collegare le API distribuite PyTorch alla semantica di comunicazione XLA.
  3. Utilizziamo PyTorch/XLA MpDeviceLoader per le pipeline di importazione dei dati. MpDeviceLoader aiuta a migliorare le prestazioni sovrapponendo tre passaggi: traccia, compilazione e caricamento in batch di dati sul dispositivo. Dobbiamo avvolgere il dataloader PyTorch con MpDeviceDataLoader come segue:
    train_device_loader = pl.MpDeviceLoader(train_loader, "xla")

  4. Eseguire la fase di ottimizzazione utilizzando l'API fornita da XLA come mostrato nel codice seguente. Questo consolida i gradienti tra i core ed emette il calcolo del passo del dispositivo XLA.
    torch_xla.core.xla_model.optimizer_step(optimizer)

  5. Mappa le API CUDA (se presenti) alle API PyTorch generiche.
  6. Sostituisci gli ottimizzatori con fusione CUDA (se presenti) con alternative PyTorch generiche.

L'intero esempio, che addestra un modello di classificazione del testo utilizzando SageMaker e Trainium, è disponibile di seguito Repository GitHub. Il file del taccuino Ottimizza i trasformatori per la creazione di modelli di classificazione utilizzando SageMaker e Trainium.ipynb è il punto di ingresso e contiene istruzioni dettagliate per eseguire la formazione.

Test di riferimento

Nel test, abbiamo eseguito due processi di addestramento: uno su ml.trn1.32xlarge e uno su ml.p4d.24xlarge con le stesse dimensioni batch, dati di addestramento e altri iperparametri. Durante i lavori di formazione, abbiamo misurato il tempo fatturabile dei lavori di formazione SageMaker e calcolato il rapporto prezzo-prestazioni moltiplicando il tempo necessario per eseguire i lavori di formazione in ore per il prezzo orario per il tipo di istanza. Abbiamo selezionato il miglior risultato per ogni tipo di istanza tra più esecuzioni di processi.

La tabella seguente riassume i risultati dei nostri benchmark.

Modello Tipo di istanza Prezzo (per nodo * ora) Throughput (iterazioni/sec) Precisione di convalida Tempo fatturabile (sec) Costo della formazione in $
Classificazione di base BERT ml.trn1.32xgrande 24.725 6.64 0.984 6033 41.47
Classificazione di base BERT ml.p4d.24xlarge 37.69 5.44 0.984 6553 68.6

I risultati hanno mostrato che l'istanza Trainium costa meno dell'istanza P4d, fornendo un throughput e una precisione simili durante l'addestramento dello stesso modello con gli stessi dati di input e gli stessi parametri di addestramento. Ciò significa che l'istanza Trainium offre un miglior rapporto prezzo-prestazioni rispetto alle istanze P4D basate su GPU. Con un semplice esempio come questo, possiamo vedere che Trainium offre tempi di addestramento più rapidi di circa il 22% e costi inferiori fino al 50% rispetto alle istanze P4d.

Distribuire il modello addestrato

Dopo aver addestrato il modello, possiamo distribuirlo a vari tipi di istanza come CPU, GPU o AWS Inferentia. Il punto chiave da notare è che il modello addestrato non dipende da hardware specializzato per la distribuzione e l'inferenza. SageMaker fornisce meccanismi per distribuire un modello addestrato utilizzando sia meccanismi in tempo reale che batch. L'esempio del notebook nel repository GitHub contiene il codice per distribuire il modello addestrato come endpoint in tempo reale usando un'istanza ml.c5.xlarge (basata sulla CPU).

Conclusione

In questo post, abbiamo esaminato come utilizzare Trainium e SageMaker per configurare e addestrare rapidamente un modello di classificazione che offre fino al 50% di risparmio sui costi senza compromettere la precisione. Puoi utilizzare Trainium per un'ampia gamma di casi d'uso che implicano il pre-addestramento o la messa a punto di modelli basati su Transformer. Per ulteriori informazioni sul supporto di varie architetture di modello, fare riferimento a Linee guida per l'adattamento dell'architettura del modello.


Informazioni sugli autori

Arun Kumar Lokanatha è un Senior ML Solutions Architect del team di assistenza di Amazon SageMaker. Si concentra sull'aiutare i clienti a creare, addestrare e migrare i carichi di lavoro di produzione ML a SageMaker su larga scala. È specializzato in Deep Learning soprattutto nell'area della PNL e CV. Al di fuori del lavoro, gli piace correre e fare escursioni.

Marco Yu è un ingegnere del software in AWS SageMaker. Si concentra sulla creazione di sistemi di formazione distribuiti su larga scala, sull'ottimizzazione delle prestazioni di formazione e sullo sviluppo di hardware di formazione ml ad alte prestazioni, tra cui SageMaker trainium. Mark ha anche una conoscenza approfondita dell'ottimizzazione dell'infrastruttura di machine learning. Nel tempo libero ama fare escursioni e correre.

Omri Fuchs è un Software Development Manager presso AWS SageMaker. È il leader tecnico responsabile della piattaforma di lavoro di formazione SageMaker, concentrandosi sull'ottimizzazione delle prestazioni di formazione SageMaker e sul miglioramento dell'esperienza di formazione. Ha una passione per la tecnologia ML e AI all'avanguardia. Nel tempo libero ama andare in bicicletta e fare escursioni.

Gal Oshri è un Senior Product Manager nel team Amazon SageMaker. Ha 7 anni di esperienza lavorando su strumenti, framework e servizi di Machine Learning.

Timestamp:

Di più da Apprendimento automatico di AWS