Come co-progettare l’architettura del software/hardware per l’IA/ML in una nuova era?

Come collaborare nella progettazione dell'architettura del software/hardware per l'IA/ML nella nuova era?

Una visione olistica del design di un’architettura efficiente per AI/ML

Le tecnologie di intelligenza artificiale generativa all’avanguardia in ambito visione artificiale, elaborazione del linguaggio naturale, ecc. stanno esplodendo di recente con scoperte scientifiche su innovative architetture di modelli, tra cui la diffusione stabile, la resa neurale (NeRF), il testo in 3D, i grandi modelli linguistici (LLM) e altro ancora. Queste tecnologie avanzate richiedono reti di intelligenza artificiale più complesse e richiedono risorse di calcolo ed architetture energeticamente efficienti di ordini di grandezza maggiori.

Per soddisfare il requisito architettonico sopra menzionato, è essenziale colmare il divario tra hardware(HW) e software(SW)/design degli algoritmi. La co-progettazione coinvolge un processo iterativo di progettazione, test e perfezionamento di componenti HW e SW fino a quando il sistema soddisfa i requisiti di prestazione desiderati. Tuttavia, SW e HW sono tradizionalmente progettati in modo indipendente poiché gli sviluppatori SW raramente devono pensare a quale HW utilizzare e l’HW viene tipicamente progettato per supportare una vasta gamma di SW. Sono disponibili studi molto limitati sulla co-progettazione SW/HW, che è fortemente necessaria per supportare carichi di lavoro AI in modo efficiente.

Fonte: Immagine dell'autore.

Dobbiamo progettare un’architettura efficiente, consapevole di SW/algoritmi e un SW/algoritmo consapevole dell’HW, in modo che possano interagire strettamente per sfruttare al massimo le limitate risorse di calcolo a nostra disposizione. Come possiamo fare? Di seguito è riportata una metodologia in evoluzione che può essere utilizzata come riferimento se si progetta nuove funzionalità HW/SW per l’AI.

1. Identificare carichi di lavoro AI proxy.

Per avviare l’analisi, abbiamo bisogno di modelli AI proxy e definire un elenco di priorità per approfondimenti successivi. Ci sono diverse risorse a cui fare riferimento, tra cui l’ultima ricerca scientifica (CVPR, Siggraph, laboratori di ricerca delle grandi aziende tecnologiche) con codice open source, feedback o richieste dei clienti, tendenze del settore, ecc. Filtrare diversi modelli rappresentativi in base al proprio giudizio esperto. Questo passaggio è cruciale poiché li userai per progettare la tua ‘futura architettura’.

2. Analisi completa dell’architettura del modello.

Assicurati di analizzare in modo completo l’architettura del modello per capirne funzionalità e innovazione e scompongila in dettagliate parti il più possibile. Ci sono nuovi operatori non supportati nello stack tecnologico attuale? Dove si trovano i livelli computazionalmente intensivi? È un modello con trasferimenti di dati (memoria) pesanti? Di che tipo di dati è richiesta e quali tecniche di quantizzazione possono essere applicate senza compromettere l’accuratezza? Quali parti del modello possono essere accelerate dall’HW e quali sono le potenziali ottimizzazioni delle prestazioni?

Ad esempio, nella resa neurale, il modello richiede sia la resa che il calcolo (moltiplicazione di matrici) che lavorano in parallelo, è necessario verificare se lo stack SW attuale supporta la resa e il calcolo contemporaneamente. Nei LLM, la dimensione della cache chiave-valore (KV) aumenta rispetto alla lunghezza della sequenza di input, è fondamentale comprendere il requisito di memoria e l’ottimizzazione potenziale della gerarchia di trasferimento dei dati/memoria per gestire una grande cache KV.

Sinistra: Disposizione della memoria su NVIDIA A100 durante la distribuzione di un LLM con 13 miliardi di parametri; Destra: Utilizzo della memoria e prestazioni in rapporto alla dimensione del batch. Fonte dell'immagine: Efficient Memory Management for Large Language Model Serving with Paged Attention [1]

3. Abilitazione e prototipazione SW

Scarica il codice open source per il modello identificato al Passaggio 2 e eseguilo sul framework SW/HW ‘target’. Questo passaggio non è semplice, soprattutto per modelli nuovi/disruptive. Poiché l’obiettivo è abilitare una soluzione funzionante per l’analisi delle prestazioni, non è necessario fornire codice di qualità prodotto in questa fase. Una correzione approssimativa sul SW senza ottimizzazione delle prestazioni è accettabile per passare al Passaggio 4. Un passaggio importante consiste nel convertire il modello pre-addestrato nel framework di sviluppo (Pytorch) in un nuovo formato richiesto dal nuovo framework di destinazione.

torch.onnx.export(model, dummy_input, "resnet50.onnx", verbose=False, input_names=input_names, outputnames=output_names, export_params=True)

Tuttavia, ci sono spesso casi in cui è necessario un notevole sforzo di supporto. Ad esempio, per eseguire modelli di rendering differenziabili, è necessario il supporto di autograd. È molto probabile che questa funzionalità non sia pronta nel nuovo framework e richieda mesi di lavoro da parte del team di sviluppo. Un altro esempio è la quantizzazione GPTQ per LLMs, che potrebbe non essere supportata inizialmente nel framework di inferenza. Invece di attendere il team di ingegneria, gli architetti possono eseguire il carico di lavoro sul sistema Nvidia per l’analisi delle prestazioni, poiché Nvidia è la scelta di HW per lo sviluppo accademico. Questo consente lo sviluppo di un elenco di requisiti SW basato sulle lacune osservate durante l’abilitazione del SW.

4. Analisi delle prestazioni e innovazione architettonica.

Ci sono numerose metriche per valutare le prestazioni di un modello di intelligenza artificiale. Di seguito sono riportate le principali considerazioni da fare.

4.1 FLOPs (Operazioni in virgola mobile) e MACs (Operazioni di moltiplicazione-somma).

Queste metriche vengono comunemente utilizzate per calcolare la complessità computazionale dei modelli di apprendimento profondo. Forniscono un modo rapido e semplice per comprendere il numero di operazioni aritmetiche richieste. I FLOPs possono essere calcolati mediante metodi come l’analisi dei paper, i report Vtune o strumenti come flops-counter.pytorch e pytorch-OpCounter.

4.2 Occupazione della memoria e larghezza di banda (BW)

La memoria occupata principalmente dai pesi (parametri di rete) e dai dati di input. Ad esempio, un modello Llama con 13 miliardi di parametri in FP16 consuma circa 13 * 2 (FP16 = 2 byte) = 26 GB di memoria (l’input è trascurabile poiché i pesi occupano molto più spazio). Un altro fattore chiave per i LLMs è la dimensione della cache KV. La cache KV occupa fino al 30% della memoria totale ed è dinamica (fare riferimento all’immagine nella fase 2). I modelli di grandi dimensioni sono di solito vincolati dalla memoria poiché la velocità dipende da quanto rapidamente vengono spostati i dati dalla memoria di sistema alla memoria locale, o dalla memoria locale ai buffer/locali. La larghezza di banda della memoria disponibile è molto migliore per predire la latenza dell’inferenza (tempo di generazione dei token per i LLM) rispetto al picco di potenza di calcolo TOPS. Un indicatore delle prestazioni è l’utilizzo della larghezza di banda della memoria (MBU) definito come larghezza di banda effettiva/larghezza di banda massima. Idealmente, un MBU vicino al 100% indica che la larghezza di banda della memoria è completamente utilizzata.

Una memoria sufficiente non è sufficiente!

Nvidia sta aumentando i FLOPs di ordini di grandezza ma non la larghezza di banda della memoria. Fonte: Substack/SemiAnalysis

Poiché la memoria rappresenta un collo di bottiglia, è necessaria l’esplorazione di tecniche di compressione avanzate del modello e tecnologie di memoria/caching. Alcuni lavori pionieristici sono elencati di seguito:

  • MemGPT: utilizza risorse di gerarchia della memoria a diversi livelli, come una combinazione di RAM veloci e piccole e memorie di archiviazione grandi e lente. Le informazioni devono essere trasferite esplicitamente tra di esse. [2]
  • Quantizzazione a bassa precisione (GPTQ, AWQ, GGML) per ridurre l’occupazione della memoria dei modelli
  • Calcolo in memoria (PIM): riduce il consumo energetico e migliora le prestazioni eliminando la necessità di spostamento dei dati.

4.3 Latenza/throughput.

Nella visione artificiale, la latenza è il tempo necessario per generare un frame. Nel contesto dei LLMs, è il tempo tra il primo token e la generazione del token successivo. Il throughput è il numero di token/frame al secondo. La latenza è una metrica critica per misurare le prestazioni del sistema di intelligenza artificiale ed è un fattore composto dalle prestazioni SW/HW. Ci sono diverse strategie di ottimizzazione da considerare, per citarne alcune:

  • Ottimizzazione delle operazioni vincolate dalla larghezza di banda come normalizzazioni, operazioni punto per punto, SoftMax e ReLU. Si stima che le normalizzazioni e le operazioni punto per punto consumino circa il 40% in più di tempo di esecuzione rispetto alle moltiplicazioni matriciali, pur raggiungendo rispettivamente 250 volte e 700 volte meno FLOPS delle moltiplicazioni matriciali. Per risolvere il problema, è possibile utilizzare la fusione del kernel per fondere più operatori al fine di risparmiare sui costi di trasferimento dei dati o sostituire operatori costosi (softmax) con operatori leggeri (ReLU).
Proporzioni delle classi di operatori. Fonte: Data movement is all you need
  • Architettura hardware specializzata. L’integrazione di hardware specializzato (AVX, GPU, TPU, NPU) può portare a significativi incrementi di velocità e risparmi energetici, il che è particolarmente importante per applicazioni che richiedono elaborazione in tempo reale su dispositivi limitati. Ad esempio, le istruzioni Intel AVX possono garantire fino a 60.000 volte più velocità rispetto al codice Python nativo.
Incrementi di velocità ottenuti dall'ingegnerizzazione delle prestazioni di un programma che moltiplica due matrici 4096x4096. Fonte: C'è molto spazio in alto: cosa guiderà le prestazioni del computer dopo la legge di Moore?[3]

Le Tensor cores su grafiche Nvidia (V100, A100, H100, ecc.) possono moltiplicare e sommare due matrici FP16 e/o FP32 in un solo ciclo di clock, a differenza delle Cuda cores che possono eseguire solo 1 operazione per ciclo. Tuttavia, l’utilizzo dei tensor core è molto basso (dal 3% al 9% per l’addestramento end-to-end), risultando in un alto costo energetico e una bassa performance. Sono in corso ricerche attive per migliorare l’utilizzo dell’array sistolico (FlexSA, SA multidirezionale, ecc.) che spiegherò nella prossima serie di articoli.

Array sistolico. Fonte: Telesens

Inoltre, poiché la memoria e il traffico dei dati costituiscono sempre un limite per i modelli AI di grandi dimensioni, è cruciale esplorare architetture avanzate che considerino una memoria più grande e efficiente a livello di chip. Un esempio è il design di memoria centrale Cerebras, in cui la memoria viene indirizzata in modo indipendente per ogni core.

Architettura di memoria Cerebras
  • Ci sono molte altre ottimizzazioni: parallelismo, cache KV quantization per LLMs, attivazione sparso e ottimizzazione end-to-end – spiegherò ulteriormente nei prossimi articoli

4.4 Potenza ed efficienza energetica

La potenza è un altro fattore che dobbiamo considerare, soprattutto per scenari utente a basso consumo energetico. È sempre un compromesso tra prestazioni e potenza. Come illustrato di seguito, l’operazione di accesso alla memoria richiede un ordine di grandezza di energia maggiore rispetto alle operazioni di calcolo. Ridurre il trasferimento di memoria è fortemente richiesto per risparmiare energia.

Costo energetico per calcolo e memoria. Fonte: Song H from Stanford [4]

Conclusioni

Ciò sono le principali metriche che dobbiamo misurare per le prestazioni dei modelli AI. Utilizzando profilers di prestazioni come Vtune, Nsight o altri strumenti di modellazione, gli architetti possono approfondire i dettagli delle prestazioni (traffico di calcolo/memoria layer per layer, efficienza del calcolo, utilizzo della banda, ecc.) e identificare i punti critici sfruttando tracce. Molto spesso, le prestazioni possono essere inaspettatamente basse a causa di inefficienze del software. È un processo iterativo per co-progettare il SW/HW AI basato su modelli di proxy.

È un’azione collaborativa per condurre un’analisi architettonica approfondita dei modelli di intelligenza artificiale emergenti. Il ciclo di indagine potrebbe essere piuttosto lungo, e ci sono sempre compromessi tra tempi di produzione e approfondimento dell’analisi. È necessaria allo stesso tempo una forte competenza nelle tecnologie HW, SW e AI per progettare soluzioni di intelligenza artificiale efficienti.

Riferimenti

[1] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, Ion Stoica, Efficient Memory Management for Large Language Model Serving with Paged Attention, 2023, arxiv

[2] Charles Packer, Vivian Fang, Shishir G. Patil, Kevin Lin, Sarah Wooders, Joseph E. Gonzalez, MemGPT: Towards LLMs as Operating Systems, 2023, arxiv

[3] Charles Leiserson, Neil Thompson, Joel Emer, Bradley Kuszmaul, Butler Lampson, Daniel Sanchez e Tao Schardl, There’s plenty of room at the Top: What will drive computer performance after Moore’s law? 2020, Science

[4] Song Han, Jeff Pool, John Tran, William J. Dally, Learning both Weights and Connections for Efficient Neural Networks, 2015, arxiv