Riepilogo Tecnico: DCC – Compilazione Centrata sui Dati dei Kernel di Machine Learning per Architetture Processing-In-Memory

1. Enunciato del Problema

Le architetture Processing-In-Memory (PIM) offrono una soluzione promettente al collo di bottiglia della memoria nel Machine Learning (ML) collocando unità di elaborazione vicino ai banchi DRAM. Tuttavia, esiste un disallineamento fondamentale tra i requisiti di layout dei dati dei processori Host ad alte prestazioni (ad esempio, GPU) e i core PIM.

• Requisito Host: Per sfruttare il parallelismo a livello di banco e l'ampia larghezza di banda, i processori Host richiedono che gli elementi di dati consecutivi siano distribuiti su più banchi DRAM.
• Requisito PIM: I core PIM possono accedere ai dati solo dai loro banchi locali. Per massimizzare la larghezza di banda locale e abilitare un accesso efficiente a più elementi, gli elementi consecutivi devono essere posizionati all'interno dello stesso banco.

Ciò rende necessarie ridisposizioni dei dati (in input e in output) eseguite tramite il bus di memoria Host, che comporta costi significativi di movimento dei dati che possono dominare il tempo di esecuzione end-to-end. Gli approcci di compilazione attuali per i sistemi PIM soffrono di due limitazioni principali:

1. Mancanza di Ottimizzazione Sistematica: Gli strumenti esistenti spesso mirano a back-end specifici (ad esempio, UPMEM) o a kernel specifici, mancando di supporto per carichi di lavoro ML diversificati e per dispositivi PIM multipli (ad esempio, Samsung HBM-PIM, SK Hynix GDDR6-AiM).
2. Sintonizzazione Centrata sul Calcolo: I compilatori esistenti (ad esempio, ATiM) ottimizzano gli schedule di calcolo indipendentemente dai costi di ridisposizione dei dati. Gli autori dimostrano che le trasformazioni di calcolo e le ridisposizioni dei dati sono interdipendenti; ottimizzarle in isolamento produce prestazioni subottimali perché uno schedule di calcolo efficiente potrebbe richiedere ridisposizioni dei dati costose, mentre uno schedule di calcolo meno efficiente potrebbe abilitare un movimento dei dati più economico.

2. Metodologia: DCC

Gli autori propongono DCC, il primo compilatore ML centrato sui dati per sistemi PIM. DCC ottimizza congiuntamente le strategie di ridisposizione dei dati e gli schedule di calcolo a livello di ciclo all'interno di un processo di sintonizzazione unificato. Il sistema comprende quattro componenti chiave:

2.1 Astrazione a Più Livelli

DCC introduce un'astrazione generica che mappa la gerarchia di memoria PIM in una gerarchia di calcolo, disaccoppiando il compilatore dalle semantiche specifiche del back-end. Questa astrazione opera a tre livelli:

• Livello di Sistema: L'xPU Host (ad esempio, GPU), lo spazio di memoria Host e i dispositivi PIM.
• Livello di Gruppo PIM: Gruppi di core PIM (e i loro banchi locali) appartenenti allo stesso canale di memoria o rank. Ciò consente ottimizzazioni a grana grossa e broadcast di istruzioni a livello di gruppo.
• Livello di Core PIM: Unità di elaborazione individuali e i loro banchi locali, abilitando ottimizzazioni a grana fine per l'accesso alla memoria e il layout dei dati locali.

2.2 Generatore di Schedule Centrato sui Dati

A differenza dei compilatori tradizionali che iniziano con trasformazioni di ciclo, DCC adotta una strategia data-first:

1. Costruzione dell'Insieme di Dimensioni: Le dimensioni dei tensori sono correlate alle variabili di ciclo per formare "insiemi di dimensioni complesse".
2. Costruzione dei Tile di Dati: Il sistema genera esaustivamente candidati partizioni di dati (tile) attraverso gruppi e core PIM.
3. Potatura delle Bozze: Tre regole eliminano le bozze poco promettenti:
   • Ridondanza: Rimozione delle bozze con partizioni di dati e mappature identiche.
   • Allineamento SIMD: Potatura delle bozze in cui le dimensioni delle partizioni non sono allineate alla larghezza SIMD del back-end (ad esempio, 16-way).
   • Equivalenza delle Prestazioni: Mantenimento di una sola bozza se più bozze producono le stesse prestazioni nel caso peggiore per core.
4. Mappatura: Le bozze rimanenti sono mappate su tile di calcolo, determinando gli intervalli di ciclo e gli indici dei tensori per ogni core.

2.3 Ottimizzatore di Codice Specifico per PIM

Questo componente applica ottimizzazioni adattate ai vincoli PIM:

• Ottimizzazione della Ridisposizione dei Dati: Sfrutta la memoria on-chip controllabile (ad esempio, memoria condivisa GPU) per riorganizzare i dati. Esegue letture seriali dai canali di memoria Host seguite da scritture parallele sui canali PIM, massimizzando il parallelismo dei canali sia per le operazioni di lettura che di scrittura.
• Ottimizzazione del Codice di Calcolo: Utilizza una generazione gerarchica dei comandi. Prioritizza i comandi di tipo SIMD a livello di gruppo (broadcast di una singola istruzione a tutti i core di un gruppo) per ridurre l'overhead del bus dei comandi, ricorrendo a comandi a livello di banco solo quando necessario.

2.4 Predittore di Prestazioni Accoppiato

Per evitare il costo proibitivo del profiling di migliaia di bozze candidate, DCC impiega un predittore basato sull'apprendimento (XGBoost).

• Funzione: Stima congiuntamente i costi di ridisposizione dei dati e i tempi di esecuzione del calcolo per selezionare la configurazione end-to-end con le migliori prestazioni.
• Addestramento: Il modello è addestrato offline su un sottoinsieme diversificato di bozze.
• Runtime: Per dimensioni di tensori statiche, lo schedule migliore viene pre-selezionato. Per dimensioni dinamiche (ad esempio, negli LLM), il predittore stima le prestazioni in tempo reale.

2.5 Interfaccia di Programmazione

DCC fornisce un'interfaccia Python compatibile con PyTorch. Gli utenti annotano i kernel con @DCC_kernel. Il compilatore analizza queste definizioni di alto livello, genera la Rappresentazione Intermedia (IR) coinvolgendo tile di dati e di calcolo, ed emette istruzioni di back-end di basso livello.

3. Risultati Chiave

Gli autori hanno valutato DCC su due back-end PIM all'avanguardia: Samsung HBM-PIM e AttAcc (un prototipo di ricerca con unità GEMV/Softmax specializzate), confrontandolo con l'esecuzione solo GPU e con implementazioni PIM di base.

• Kernel ML Individuali:
  • HBM-PIM: DCC raggiunge un speedup fino a 7,68× (media 2,21×) rispetto all'esecuzione GPU.
  • AttAcc: DCC raggiunge un speedup fino a 13,17× (media 3,92×) rispetto all'esecuzione GPU.
  • Confronto con PIM di Base: DCC migliora l'implementazione di base AttAcc fino a 2,74× (media 1,75×) e HBM-PIM fino a 2,09× (media 1,53×) per kernel specifici come GEMV e RED.
• Inferenza End-to-End di LLM:
  • Su AttAcc, DCC accelera l'inferenza GPT-3 e LLaMA-2 di una media di 4,52× (fino a 7,71× per LLaMA-2) rispetto all'esecuzione solo GPU.
  • DCC supera significativamente l'implementazione di base AttAcc (che utilizza tiling fisso) abilitando l'esecuzione efficiente dei layer di generazione e proiezione QKV, che in precedenza erano un collo di bottiglia sulla GPU.
• Efficienza Energetica: DCC riduce il consumo energetico di una media di 1,24× rispetto al sistema di base AttAcc a causa della riduzione del movimento dei dati.
• Overhead di Compilazione: Il tempo di compilazione offline è basso (secondi) e lo schema di potatura riduce lo spazio di ricerca di circa 9×. Il predittore accoppiato raggiunge un'alta accuratezza (97,37% sugli LLM) e riduce il tempo di compilazione di circa 3,8× rispetto alla ricerca esaustiva.

4. Significato e Affermazioni

Il documento afferma che DCC è il primo compilatore ML centrato sui dati per sistemi PIM che:

1. Ottimizza Congiuntamente la ridisposizione dei dati e gli schedule di calcolo, dimostrando che questi due fattori sono interdipendenti e devono essere sintonizzati insieme per evitare prestazioni subottimali.
2. Supporta Multipli Back-end attraverso un'astrazione a più livelli, abilitando la portabilità su diverse architetture PIM (ad esempio, HBM-PIM, AttAcc) senza richiedere sintonizzazione manuale specifica per il back-end.
3. Abilita Alte Prestazioni per kernel ML ad alta intensità di memoria e inferenza su larga scala di LLM, superando significativamente sia l'esecuzione solo GPU che gli approcci di compilazione PIM esistenti.
4. Migliora la Programmabilità fornendo un'interfaccia Python di alto livello che astrae la complessità della gestione del layout dei dati di basso livello e della pianificazione delle istruzioni per i dispositivi PIM.

Gli autori concludono che DCC fornisce un percorso pratico ed efficace per l'accelerazione ML eterogenea, bilanciando alte prestazioni e programmabilità negli emergenti ecosistemi PIM. Il codice è open-source per facilitare ulteriori ricerche.