TrackCore-F: Deploying Transformer-Based Subatomic Particle Tracking on FPGAs

Il paper presenta lo sviluppo di metodologie e strumenti per la sintesi e il deployment su FPGA di modelli Transformer basati sul progetto TrackFormers per il tracciamento di particelle subatomiche, affrontando le sfide legate alla limitatezza delle risorse hardware e alla necessità di partizionamento automatico dei modelli.

L'essenziale

🚀 TrackCore-F: Insegnare ai "Supercomputer Portatili" a vedere l'infinitamente piccolo

Immagina di essere in un enorme stadio pieno di gente che corre in tutte le direzioni. Il tuo compito è seguire ogni singola persona, capire da dove viene, dove sta andando e con chi si sta incontrando, tutto questo in un millisecondo.

Nel mondo della fisica delle particelle, questo è esattamente quello che fanno i rivelatori al CERN (il Large Hadron Collider). Ma c'è un problema: ci sono così tante particelle che i computer normali fanno fatica a tenere il passo. È come se dovessi risolvere un puzzle di un milione di pezzi mentre corri su un tapis roulant.

1. Il Problema: Troppa velocità, troppo calore

Fino a poco tempo fa, per risolvere questi puzzle si usavano i GPU (le schede video potenti dei computer da gioco). Sono veloci, ma sono come dei camioncini: ingombranti, consumano molta energia e non si possono mettere facilmente dentro i rivelatori stessi.

Gli scienziati hanno bisogno di qualcosa di più piccolo, più efficiente e capace di lavorare "sul campo" (direttamente nel rivelatore), senza aspettare che i dati arrivino al centro dati. La soluzione? Gli FPGA.

• Cos'è un FPGA? Immagina un blocco di LEGO elettronico. A differenza di un computer normale (che è come un'auto già costruita e non modificabile), un FPGA è un telaio vuoto dove puoi costruire il motore esatto che ti serve, pezzo per pezzo, per fare un solo compito specifico. È piccolo, veloce e consuma pochissima energia.

2. La Sfida: I "Cervelli" troppo grandi

Per riconoscere le particelle, gli scienziati usano l'Intelligenza Artificiale, in particolare una tecnologia chiamata Transformer (la stessa tecnologia che fa funzionare i chatbot intelligenti come me!).
Il problema è che questi "cervelli" digitali sono enormi. Metterli dentro un piccolo blocco di LEGO (l'FPGA) è come cercare di far entrare un elefante in una Fiat 500. Non ci stanno!

Inoltre, gli strumenti per trasformare questi "cervelli" complessi in circuiti elettronici sono ancora molto primitivi. È come se avessi un progetto architettonico per un grattacielo, ma avessi solo un martello e un chiodo per costruirlo.

3. La Soluzione di TrackCore-F: Tagliare e Incollare

Il team di ricerca (dall'Università di Twente, in Olanda) ha sviluppato un nuovo metodo, chiamato TrackCore-F, per risolvere questo problema.

Ecco come funziona, passo dopo passo:

• Il Taglio (Partitioning): Invece di cercare di mettere tutto il "cervello" (il modello AI) dentro l'FPGA, lo tagliano a fette. Immagina di prendere una torta gigante e di metterne solo una fetta nel tuo piccolo frigorifero (l'FPGA).
• Il Lavoro di Squadra: La fetta di torta che sta nel frigorifero fa il lavoro pesante e veloce. Le altre fette restano nel "frigorifero grande" (il processore normale del computer). Si passano i dati avanti e indietro in modo che il lavoro sembri continuo.
• Il Risultato: Riescono a far girare queste intelligenze artificiali direttamente dentro i rivelatori, ottenendo risultati quasi istantanei (latenza ultra-bassa).

4. I Risultati: Velocità contro Precisione

Gli scienziati hanno provato a "semplificare" i numeri usati dal computer (una tecnica chiamata quantizzazione) per farli entrare meglio nell'FPGA.

• L'analogia: È come se invece di scrivere "3,14159265..." (precisione massima), scrivessimo solo "3" (precisione bassa). Risparmi molto spazio, ma potresti sbagliare il calcolo.
• La scoperta: Hanno scoperto che se semplifichi troppo i numeri, l'AI diventa confusa e sbaglia a riconoscere le particelle. È meglio usare numeri un po' più precisi, anche se occupano più spazio.

Hanno anche scoperto che il vero collo di bottiglia non è la potenza di calcolo, ma la memoria interna (i "cassetti" dove tenere i dati). Il loro prototipo riesce a gestire circa 4 "strati" di questo cervello artificiale prima di riempire tutti i cassetti disponibili.

5. Perché è importante?

Questo lavoro è fondamentale per il futuro della fisica.
Immagina che il CERN stia per diventare ancora più potente (la fase "High-Luminosity"). Ci saranno così tante particelle che i computer attuali si schianteranno.
Con TrackCore-F, gli scienziati potranno mettere questi "cervelli" intelligenti direttamente dentro i rivelatori.

• Vantaggio: Non dovremo più aspettare ore per analizzare i dati. Potremo vedere le particelle in tempo reale, proprio mentre accadono.
• Analogia finale: È come passare da un investigatore che deve portare le foto in laboratorio per svilupparle (lento), a un detective che ha una fotocamera istantanea e un computer tascabile che risolve il caso sul posto in un secondo.

In sintesi: Gli scienziati hanno imparato a "scomporre" i cervelli artificiali più complessi per farli entrare in piccoli chip programmabili, permettendo alla fisica di guardare l'universo a velocità supersonica.

Sommario tecnico

Titolo: TrackCore-F: Implementazione di Tracciamento di Particelle Subatomiche basato su Transformer su FPGA

1. Il Problema

Il tracciamento delle particelle (tracking) negli esperimenti di fisica delle alte energie, come quelli condotti al Large Hadron Collider (LHC), è diventato un compito computazionalmente oneroso. Sebbene l'uso di architetture di Machine Learning (ML), in particolare i Transformer, abbia portato a soluzioni efficaci per compiti come il "jet tagging" e la ricostruzione delle tracce, la loro implementazione in tempo reale (online o pseudo-online) rimane una sfida.

• Vincoli di Latenza: Le attuali soluzioni basate su GPU, sebbene potenti, non sempre garantiscono le latenze necessarie per l'elaborazione online, specialmente con l'aumento della luminosità previsto per la fase HL-LHC (High-Luminosity LHC).
• Limitazioni Hardware: L'uso di acceleratori hardware specializzati come le FPGA (Field-Programmable Gate Array) offre vantaggi in termini di efficienza energetica e possibilità di deployment on-site, ma il supporto degli strumenti attuali per sintetizzare modelli Transformer su FPGA è limitato o inesistente.
• Vincoli di Risorse: I modelli Transformer tendono ad essere grandi. Le risorse hardware delle FPGA (memoria, logica, DSP) sono limitate, rendendo difficile il deployment di modelli completi senza una partizione intelligente o una sintesi automatizzata.

2. Metodologia

Gli autori propongono un flusso di lavoro strutturato per la sintesi e il deployment di modelli Transformer pre-addestrati su FPGA, focalizzandosi sull'inferenza.

• Modelli Utilizzati: Lo studio si basa su due architetture Transformer derivate dal progetto TrackFormers:
  1. EncCla (Encoder-Classifier): Un modello "encoder-only" che associa gli hit (segnali dei rivelatori) a etichette di classe predefinite. Ha circa 1,5 milioni di parametri.
  2. EncReg (Encoder-Regressor): Un modello "encoder-only" che regredisce i parametri delle tracce potenziali, richiedendo un successivo clustering (es. HDBSCAN). Ha circa 76.000 parametri.
• Dataset: I modelli sono stati addestrati su 5 dataset che variano in complessità, da tracce lineari ed elicoidali semplici (simulazioni REDVID) a dataset più complessi derivati dalla sfida TrackML Kaggle.
• Flusso di Sviluppo (Toolchain):
  1. PyTorch: I modelli sono forniti in formato PyTorch.
  2. ONNX: Conversione in formato Open Neural Network Exchange per esporre le operazioni a basso livello e facilitare la quantizzazione.
  3. AMD Vitis HLS: Utilizzato per sintetizzare kernel hardware in C/C++ o OpenCL a partire da descrizioni software.
  4. AMD Vivado: Integrato per configurare l'IP su Zynq MPSoC e generare il bitstream finale.
  5. PYNQ: Ambiente Python per l'interazione con la logica programmabile (PL) dello Zynq.
• Strategia di Deployment:
  • Viene adottato un approccio di partizionamento manuale. Il modello viene diviso in tre parti: segmenti prima del layer encoder, il layer encoder da implementare sull'FPGA, e il segmento residuo.
  • Viene sviluppato un kernel monolitico per il layer encoder (inclusi Multi-Head Attention, Add & Norm, Feed-Forward) in C/C++ per Vitis HLS.
  • I pesi vengono trasferiti al kernel tramite interfaccia AXI4-Lite.
  • L'obiettivo è creare un flusso che permetta sia il deployment monolitico (se le risorse lo consentono) che quello partizionato.

3. Risultati Chiave

• Consumo di Risorse (FPGA ZCU102):

  • Il test è stato eseguito su un kit di valutazione ARM Zynq UltraScale+ MPSoC (ZU9EG).
  • L'implementazione di un singolo layer encoder ha consumato:
    • BRAM (Block RAM): 38,04% delle risorse disponibili (347 su 912).
    • LUT: 24,01%.
    • DSP: 2,66%.
    • Flip-Flops: 13,53%.
  • Collo di bottiglia: La risorsa limitante è la BRAM. Tuttavia, l'uso di memoria DDR esterna potrebbe ridurre l'utilizzo della BRAM a scapito di una maggiore latenza di trasferimento.
  • Scalabilità: Dati i consumi attuali, è stimabile che su questa FPGA specifica si possano implementare al massimo 4 layer encoder per questo modello.

• Effetti della Quantizzazione:

  • È stata analizzata l'impatto della quantizzazione (INT8/INT16) su pesi e attivazioni per migliorare le prestazioni.
  • Risultato critico: La quantizzazione delle attivazioni ha un impatto devastante sulla precisione del modello.
    • Modello base (Float): Accuratezza 0,97.
    • Pesi INT8 + Attivazioni INT16: Accuratezza scende a 0,71.
    • Pesi INT8 + Attivazioni INT8: Accuratezza scende a 0,70.
  • La quantizzazione dei soli pesi (mantenendo le attivazioni in precisione più alta) ha un impatto minore, ma le configurazioni miste non sono supportate nativamente dagli strumenti ONNX Static quantization attuali, rendendo l'implementazione complessa.

• Performance: Sebbene non vengano riportati numeri assoluti di throughput in questo abstract preliminare, l'obiettivo è raggiungere latenze competitive o migliori rispetto alle GPU, sfruttando l'efficienza energetica delle FPGA.

4. Contributi Principali

1. Flusso di Sviluppo Riproducibile: Definizione di una pipeline di strumenti (PyTorch -> ONNX -> Vitis HLS -> Vivado) per il deployment di modelli Transformer su FPGA, un'area attualmente poco supportata.
2. Strategia di Partizionamento: Dimostrazione della fattibilità del deployment parziale (slice del modello) su FPGA, permettendo l'inferenza su hardware più accessibile anche quando il modello completo non entra nella memoria on-chip.
3. Analisi delle Risorse: Valutazione dettagliata del consumo di risorse (BRAM, LUT, DSP) per un layer encoder Transformer, fornendo linee guida per la scalabilità su FPGA specifiche (ZCU102).
4. Valutazione della Quantizzazione: Evidenziazione del trade-off critico tra efficienza hardware e accuratezza, mostrando come la quantizzazione delle attivazioni possa degradare significativamente le prestazioni nei task di fisica delle particelle.

5. Significato e Impatto

Questo lavoro è significativo per la comunità della fisica delle alte energie perché:

• Abilita l'Online Tracking: Sposta il tracciamento assistito da ML da un compito "post-mortem" (offline) verso soluzioni online o pseudo-online, cruciali per gestire i dati ad alta luminosità dell'HL-LHC.
• Efficienza Energetica: Promuove l'uso di FPGA come alternativa sostenibile alle GPU per l'elaborazione in tempo reale nei rivelatori di particelle.
• Guida per l'Hardware: Fornisce dati empirici sui limiti delle risorse FPGA per i Transformer, guidando i futuri sviluppi di ottimizzazione (es. pruning, ottimizzazione di funzioni non lineari come SoftMax) e di strumenti di sintesi automatica.
• Accessibilità: Dimostra che anche un deployment parziale può essere prezioso, democratizzando l'accesso a soluzioni ML avanzate su hardware embedded più economico e diffuso.

In conclusione, il paper stabilisce le basi metodologiche per integrare modelli Transformer complessi nell'ecosistema hardware delle FPGA, evidenziando sia le potenzialità che le sfide attuali legate alla memoria e alla precisione numerica.