In che modo Amazon Search M5 ha risparmiato il 30% sui costi di formazione LLM utilizzando AWS Trainium

Ripubblicato da Platone

Seguaci: 0

Per decenni, Amazon è stata pioniera e innovatrice del machine learning (ML), offrendo esperienze piacevoli ai propri clienti. Fin dai primi giorni, Amazon ha utilizzato il machine learning per vari casi d'uso come consigli sui libri, ricerca e rilevamento di frodi. Analogamente al resto del settore, i progressi dell’hardware accelerato hanno consentito ai team di Amazon di perseguire architetture modello utilizzando reti neurali e deep learning (DL).

Il programma M5 all'interno di Amazon Search possiede la strategia di apprendimento della scoperta per Amazon e crea modelli su larga scala multilingue, multilocale, multientità, multitasking e multimodali come testo, immagini e video. Il programma M5 ha fornito incorporamenti universali e modelli di base su larga scala a centinaia di team ML in tutta Amazon, mantenendo allo stesso tempo controlli rigorosi sull'ottimizzazione dei costi. Per raggiungere questo obiettivo, il team M5 valuta regolarmente nuove tecniche per ridurre i costi.

Come molte organizzazioni ML, gli acceleratori vengono ampiamente utilizzati per accelerare la formazione e l'inferenza DL. Quando AWS ha lanciato acceleratori appositamente realizzati con la prima versione di AWS Inferenza nel 2020, il team M5s ha iniziato rapidamente a farlo utilizzarli per distribuire in modo più efficiente i carichi di lavoro di produzione, risparmiando sui costi e riducendo la latenza. L'anno scorso, AWS ha lanciato il suo AWSTrainium acceleratori, che ottimizzano le prestazioni per costo per lo sviluppo e la creazione di modelli DL di prossima generazione. In questo post, discutiamo di come M5 sia riuscita a ridurre i costi di addestramento dei propri modelli del 30% e condividiamo alcune delle migliori pratiche apprese lungo il percorso.

Istanze di Trainium

Con i progressi negli acceleratori appositamente realizzati, Amazon fornisce anche acceleratori interessanti sotto forma di AWS Inferentia e Trainium. Come suggerisce il nome, questi chip sono ottimizzati per superare le esigenze rispettivamente dei carichi di lavoro di inferenza e di training. Per l'addestramento su larga scala di modelli di fondazione che raggiungono miliardi di parametri in termini di dimensioni, Trainium Istanze Trn1 e Trn1n sono scelte ideali per le loro caratteristiche. Le istanze Trn1 sono alimentate dallo stato dell'arte NeuronCore-v2e dispongono di una notevole quantità di calcolo e memoria dell'acceleratore. Le istanze Trn1n possono essere scelte anche per una maggiore quantità di larghezza di banda di rete (1,600 Gb), quindi sono ideali per una formazione performante con in mente l'ottimizzazione dei costi.

Per utilizzare gli acceleratori, è necessario un livello software che li supporti. Con i chip Trn e Inf, il SDK di AWS Neuron sblocca gli acceleratori appositamente realizzati da Amazon con l'aiuto di PyTorch XLA. PyTorch XLA converte la modalità desiderosa di PyTorch in un'implementazione basata su grafici in modalità lazy. Questi grafici vengono quindi utilizzati e ulteriormente compilati per essere utilizzati con l'acceleratore. PyTorch Neuron (parte di Neuron SDK) consente agli utenti PyTorch di addestrare i propri modelli su Trainium NeuronCores con poche righe di codice.

Modello e carico di lavoro

Il team M5 forma e distribuisce modelli fondamentali e rappresentazioni universali per assistere vari team in tutta Amazon nel portare gioia Amazon.com clienti. Uno di questi modelli è un modello di codificatore di testo seguito da un percettrone multistrato (MLP) con interazioni di funzionalità esplicite o implicite definite dall'architettura della rete neurale con centinaia di milioni di parametri addestrabili. Questo modello viene addestrato su miliardi di token e viene utilizzato per generare milioni di incorporamenti in un'impostazione di inferenza batch offline. Questi incorporamenti sono input per un servizio Amazon di livello 1 rivolto al cliente.

Utilizza l'infrastruttura per la pipeline di produzione Batch AWS con strategie di coda equa, utilizzando un cluster trn1.32xlarge multinodo abilitato per EFA come calcolo per l'addestramento del modello. Funzionalmente, la pipeline di produzione esegue l'addestramento incrementale del modello, la valutazione del modello addestrato e l'inferenza batch offline sul modello addestrato, il tutto utilizzando PyTorch come libreria DL sottostante.

Obiettivi

Deliziare i nostri clienti è un principio fondamentale. Data la natura della pipeline rivolta al cliente, è fondamentale che tutti gli accordi sul livello di servizio (SLA) siano rispettati senza regressioni. Abbiamo identificato due criteri di accettazione critici per adattare la nostra pipeline di produzione di GPU esistente e trasferirla a Trainium:

Qualità del modello – La qualità dei nostri modelli ha un impatto diretto sull’esperienza del cliente. Richiediamo che ci sia una differenza inferiore allo 0.1% nella qualità del modello tra GPU e Trainium.
Produttività della formazione – Formiamo periodicamente i nostri modelli in modo iterativo per fornire l'esperienza più fresca ai nostri clienti. Richiediamo che la convergenza del modello venga raggiunta entro un periodo di tempo predefinito (ad esempio 1 settimana) per soddisfare i nostri SLA di produzione.

Nelle sezioni seguenti, condividiamo il nostro percorso di lavoro a ritroso partendo da questi criteri e ciò che abbiamo imparato per supportare carichi di lavoro di produzione su scala Amazon.

Script di addestramento

Prima di iniziare con l'addestramento del modello, è necessario apportare modifiche allo script di addestramento per renderlo conforme a XLA. Date le dimensioni del modello, utilizziamo DDP (Distributed Data Parallel) per addestrare il modello. DDP ci consente di aumentare la produttività dell'addestramento del modello aumentando il numero di macchine utilizzate per eseguire l'addestramento del modello, senza alcuna modifica del codice. Abbiamo seguito le istruzioni fornite nel Tutorial di formazione MLP Neuron PyTorch per aggiungere costrutti specifici per XLA nei nostri script di training. Queste modifiche al codice sono semplici da implementare. Di seguito sono riportati alcuni insegnamenti tecnici significativi derivanti dall'esercizio che hanno notevolmente migliorato la produttività del nostro modello:

Posizionamento di xm.mark_step() - xm.mark_step() compila ed esegue i grafici di calcolo raccolti pigramente. Invocare mark_step troppe volte porterà a un numero maggiore di piccoli grafici, mentre invocarlo troppo poche volte porterà a pochi, ma grandi grafici. A seconda dell'applicazione, la produttività e l'implementazione della formazione del modello varieranno in base al posizionamento di xm.mark_step(). La nostra implementazione ne colloca uno xm.mark_step() dopo un passaggio in avanti e all'indietro e uno dopo il passaggio di ottimizzazione.
Wrapping del caricatore dati con caricatore del dispositivo multiprocessore XLA – Questo è un passaggio fondamentale che può essere facilmente trascurato. Il caricatore del dispositivo multiprocessore torch_xla.distributed.parallel_loader.MpDeviceLoader carica i dati di addestramento su ciascun dispositivo XLA con opzioni per precaricare e sovrapporre il caricamento dei dati con le esecuzioni del dispositivo per migliorare la velocità effettiva. Anche il caricatore del dispositivo invoca xm.mark_step() ed è quindi in grado di costruire grafici per il caricamento dei dati sul dispositivo dall'host.

Compilazione per Trainium

Tradizionalmente, il ciclo di sviluppo del modello con GPU prevede l'apporto di modifiche al modello o allo script di training e l'esecuzione diretta sul dispositivo GPU. Gli acceleratori come Trainium che utilizzano XLA richiedono un passaggio aggiuntivo prima che l'addestramento del modello possa essere eseguito sull'acceleratore. I grafici di calcolo XLA possono essere eseguiti solo dopo essere stati compilati. In genere, esistono due modi per eseguire questa compilazione: Ahead of Time (AOT), in cui si tracciano e compilano prima tutti i grafici e poi li si esegue, o Just In Time (JIT), in cui i grafici vengono tracciati, compilati ed eseguiti man mano che vengono eseguiti. si incontrano. Neuron SDK fornisce entrambi questi elementi immediatamente. In genere, la compilazione AOT viene eseguita per prima. I grafici vengono quindi eseguiti dopo questa compilazione. Se vengono rilevati nuovi grafici, il runtime Neuron richiama una compilazione JIT prima di eseguirli. Per eseguire la compilazione AOT, Neuron SDK fornisce neuron_parallel_compile, un'utilità di compilazione che estrae i grafici da un'esecuzione di prova dello script di training ed esegue la compilazione AOT parallela.

Un aspetto importante della compilazione AOT è garantire che non vengano creati nuovi grafici di calcolo nel corso della formazione. Una fonte di nuovi grafici di calcolo (e quindi di ricompilazioni) sono le forme dinamiche dei batch di addestramento durante l'addestramento del modello. Abbiamo scoperto che l'utilizzo di forme statiche e batch di dimensioni fisse elimina la compilazione dei tempi di addestramento e migliora notevolmente la produttività dell'addestramento senza alcun effetto sulla precisione del modello. Applicando tali vincoli all'addestramento, abbiamo osservato che sono necessari solo 4-5 passaggi di addestramento del modello, un passaggio di convalida del modello e il checkpoint del modello una volta per tracciare tutti i grafici durante la compilazione AOT. È importante notare che Neuron SDK è in continua evoluzione e in futuro supporterà anche forme dinamiche.

Inoltre, i grafici compilati vengono archiviati nel file Cache persistente dei neuroni su disco o in un Servizio di archiviazione semplice Amazon (Amazon S3) secchio. Ciò è particolarmente utile per i carichi di lavoro di produzione in cui l'architettura del modello e la configurazione del training non cambiano. Pertanto, il sovraccarico della compilazione viene sostenuto una sola volta. Usare la cache è semplice come impostare un flag di ambiente:

export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

Anche il compilatore Neuron ne fornisce tre opzioni di ottimizzazione a livello di compilatore (O1, O2, O3) per bilanciare il tempo di compilazione e la velocità di esecuzione del modello. O1 consente ottimizzazioni fondamentali sul grafico di calcolo e riduce al minimo il tempo di compilazione, O3 fornisce una migliore velocità effettiva di esecuzione del modello al costo di un tempo di compilazione più elevato e O2 (opzione predefinita) è un equilibrio tra i due. Per il nostro caso d'uso, abbiamo utilizzato l'ottimizzazione O1 e osservato una riduzione dell'86% del tempo di compilazione senza alcuna modifica ai parametri di accuratezza del modello, osservando al contempo una riduzione del throughput di circa il 5–7% rispetto all'ottimizzazione predefinita (O2). A seconda del caso d'uso è possibile scegliere diversi livelli di ottimizzazione.

Per riassumere, abbiamo utilizzato i seguenti flag per la compilazione:

NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

Compatibilità dei checkpoint

Una volta completata con successo la compilazione, possiamo procedere al training dei nostri modelli su Trainium. Come accennato in precedenza, addestriamo i nostri modelli in modo incrementale, il che significa che carichiamo un checkpoint del modello precedentemente addestrato e continuiamo l'addestramento con nuovi dati. PyTorch e PyTorch XLA consentono una transizione fluida tra acceleratori attraverso l'interoperabilità dei checkpoint. La flessibilità di spostarsi tra GPU e Trainium ci ha permesso di caricare senza problemi il modello GPU precedente e di eseguire l'addestramento sulle macchine Trainium. Questo è stato fondamentale per garantire che potessimo inizializzare il nostro modello con il miglior modello addestrato in precedenza senza tempi di inattività della produzione o perdita di precisione del modello.

Poiché il modello GPU è stato salvato utilizzando le utilità di salvataggio del modello PyTorch standard, siamo stati in grado di utilizzare l'utilità di caricamento del checkpoint PyTorch per caricare il modello GPU sui dispositivi Trainium.

Ad esempio, su GPU/CPU, puoi salvare il modello con il seguente codice:

torch.save(model.state_dict(), PATH)

Quindi carichi nuovamente il modello su Trainium:

import torch_xla.core.xla_model as xm
xla_device = xm.xla_device()
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(xla_device)

Allo stesso modo, puoi salvare il modello su Trainium con il seguente codice:

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH)

E carica nuovamente il modello su GPU/CPU:

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

Infatti, poiché utilizziamo DDP per l'addestramento del modello, il caricamento del modello è indipendente dal numero di macchine utilizzate per addestrare il checkpoint precedente. Ciò ci consente di ridimensionare orizzontalmente la flotta Trn1 senza modifiche al codice o effetti negativi sull’addestramento del modello. Questi checkpoint basati su PyTorch possono essere utilizzati direttamente o addirittura scritti con script per casi d'uso di inferenza su AWS Inferentia2 o altri acceleratori.

Stabilità operativa

Non sarà mai sottolineato abbastanza che l’esecuzione dei carichi di lavoro in produzione richiede il rispetto di più SLA. Per il nostro caso d'uso, a parte gli SLA relativi alla qualità del modello e al throughput dell'addestramento, è fondamentale che la pipeline di produzione sia operativamente stabile, il che significa tempi di inattività e interruzioni minimi durante l'addestramento, la valutazione e l'inferenza del modello.

Come per la pipeline esistente basata su GPU, abbiamo aggiunto numerosi meccanismi per rendere la pipeline operativamente stabile. Prima di iniziare l'addestramento del modello, eseguiamo numerosi test di integrità per valutare lo stato delle macchine. Questi test generalmente includono semplici operazioni tensoriali per verificare lo stato dei dispositivi acceleratori. Abbiamo osservato che per la formazione distribuita è importante eseguire test per verificare anche la comunicazione collettiva tra le istanze. Abbiamo usato il Suite di test NCCOM dall'SDK Neuron per raggiungere questo obiettivo, eseguendo una varietà di operazioni come all-gather, all-reduce e reduce-scatter.

Anche dopo aver seguito i suggerimenti che abbiamo menzionato, abbiamo osservato che i problemi transitori sono inevitabili in qualsiasi pipeline, indipendentemente dall’acceleratore sottostante. Per creare resilienza in qualsiasi pipeline di formazione, consigliamo di creare meccanismi di ripetizione per risolvere questi potenziali problemi. Noi usiamo Tentativi automatizzati di AWS Batch per riprovare i processi che riscontrano un errore temporaneo durante l'addestramento del modello. Questi riavvii possono essere costosi se si verifica un errore verso la fine della formazione. Per contrastare questo problema, abbiamo adattato i nostri script di addestramento per caricare un checkpoint del modello precedentemente addestrato e continuare l'addestramento da quel punto. Con questa funzionalità, siamo in grado di riavviare in modo aggressivo i lavori di formazione non riusciti con un sovraccarico minimo.

Con questi meccanismi di resilienza in atto, siamo stati in grado di raggiungere tassi di successo del 98.5% per i nostri carichi di lavoro su Trn1, paragonabili ai tassi di successo della nostra pipeline GPU esistente.

Risultati

Per convalidare l'accuratezza dei nostri modelli, abbiamo inizializzato due modelli dallo stesso checkpoint GPU e ne abbiamo addestrato uno su Trainium e l'altro su una GPU comparabile. Entrambi i modelli sono stati addestrati con gli stessi iperparametri di addestramento. Il set di dati utilizzato per il calcolo delle metriche è un set di dati di controllo e valutiamo l'accuratezza del modello su questo set di dati ogni N passaggi globali. L'asse X è il passo globale e l'asse Y è la precisione del modello. Abbiamo osservato una differenza inferiore allo 0.1% nell'accuratezza del modello in ogni punto del grafico seguente.

Inoltre, per valutare il rapporto costo-efficacia dell'addestramento del modello, preferiamo confrontare il tempo impiegato per raggiungere la convergenza del modello. Riteniamo che ciò fornisca una visione più pratica del risparmio sui costi rispetto a misure quali costo per token, FLOPS/dollaro raggiunti e altri fattori. Considerando il tempo di allenamento di trn1.32xl e comparabile Cloud di calcolo elastico di Amazon (Amazon EC2), abbiamo osservato che Trainium offre un costo inferiore fino al 30% per la convergenza dei modelli.

Conclusione

Sono molti i fattori da considerare quando si valutano diversi acceleratori per i carichi di lavoro DL. Alcuni dei più importanti sono la qualità del modello, la produttività, i costi e la disponibilità. È fondamentale garantire che la qualità del modello e la produttività non vengano sacrificati in base all'acceleratore scelto.

Grazie alla nostra partnership e collaborazione con il team Annapurna Neuron, il team Amazon Search M5 è stato in grado di risparmiare fino al 30% sui costi passando a Trainium. Il team è in grado di utilizzare Trainium e raggiungere la qualità del modello e la parità di rendimento con acceleratori comparabili sul mercato. L'interoperabilità dei checkpoint e le modifiche minime al codice con il supporto per XLA hanno consentito a M5 di scegliere tra più acceleratori per i propri carichi di lavoro. Ciò ha consentito al team M5 di sfruttare l'ampia potenza di elaborazione di Trainium e di creare soluzioni indipendenti dall'acceleratore per soddisfare i clienti di Amazon.com. Da un punto di vista operativo, Trainium si è dimostrato in grado di supportare servizi di primo livello su scala Amazon. Il team M1 continua a spostare più carichi di lavoro su Trainium per fornire i migliori modelli per Amazon ai costi più bassi.

In sintesi, il team M5 è stato in grado di eseguire una formazione ML economicamente vantaggiosa e di livello produttivo aggiungendo Trainium alla flotta di acceleratori. Ti invitiamo a dare un'occhiata a Trainium e ad altri dispositivi Neuron come AWS Inferentia per sfruttare i vantaggi del silicio Amazon appositamente creato per i carichi di lavoro ML. Inizia facilmente con uno dei tanti tutorial con diversi modelli, come Lama 2, disponibile su Trainium.

Informazioni sugli autori

Abhinandan Patni è un Senior Software Engineer presso Amazon Search. Si concentra sulla creazione di sistemi e strumenti per la formazione di deep learning distribuito scalabile e l'inferenza in tempo reale.

James Park è un Solutions Architect presso Amazon Web Services. Lavora con Amazon.com per progettare, costruire e distribuire soluzioni tecnologiche su AWS e ha un interesse particolare per l'intelligenza artificiale e l'apprendimento automatico. Nel tempo libero gli piace scoprire nuove culture, nuove esperienze e rimanere aggiornato sulle ultime tendenze tecnologiche. Lo puoi trovare su LinkedIn.

Jerry Mannil è un ingegnere informatico presso Amazon Search. Lavora per migliorare l'efficienza, la robustezza e la scalabilità dell'infrastruttura di formazione distribuita.

Ken Su è un ingegnere informatico presso Amazon Search. Lavora per migliorare l'efficienza della formazione e il flusso di lavoro di formazione distribuito scalabile. Fuori dal lavoro gli piacciono le escursioni e il tennis.

RJ è un ingegnere in Amazon. Costruisce e ottimizza sistemi per sistemi distribuiti per la formazione e lavora sull'ottimizzazione dell'adozione di sistemi per ridurre la latenza per ML Inference. Fuori dal lavoro, sta esplorando l'utilizzo dell'intelligenza artificiale generativa per creare ricette alimentari.