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

Cet article présente le développement de méthodologies et d'outils pour synthétiser et déployer des modèles de transformateurs destinés au suivi de particules subatomiques sur des FPGA, en abordant les défis de partitionnement et de contraintes de ressources pour atteindre des latences en ligne.

L'essentiel

🚀 TrackCore-F : Des particules, de l'IA et des puces magiques

Imaginez que vous êtes dans un immense stade de football (le Grand Collisionneur de Hadrons ou LHC). Des joueurs (les particules) courent à une vitesse folle dans des directions opposées et entrent en collision au centre. À chaque collision, des milliers de confettis (d'autres particules) explosent dans toutes les directions.

Le but des physiciens ? Retrouver la trace de chaque confetti pour comprendre d'où il vient et ce qu'il est. C'est ce qu'on appelle le suivi de particules (ou "tracking").

1. Le Problème : Trop de données, trop vite !

Jusqu'à présent, pour analyser ces collisions, les scientifiques prenaient les données, les stockaient, et les analysaient plus tard, comme un détective qui regarde des photos de scène de crime des jours après le crime. C'est trop lent ! Avec les futures expériences, il y aura tellement de collisions que les ordinateurs classiques (les gros serveurs avec des cartes graphiques) ne pourront pas suivre le rythme en temps réel.

Il faut une solution plus rapide, plus petite et moins gourmande en énergie. C'est là qu'intervient la puce FPGA.

L'analogie : Imaginez que les ordinateurs classiques sont comme de grands camions de déménagement : ils sont puissants mais lents et consomment beaucoup de carburant. La puce FPGA, c'est comme un robot artisan ultra-rapide que vous pouvez reconfigurer instantanément pour faire exactement le travail dont vous avez besoin, sur place, avec très peu d'énergie.

2. La Solution : L'Intelligence Artificielle (Transformer)

Pour retrouver les traces des particules, les scientifiques utilisent une forme d'IA très puissante appelée Transformer (la même technologie qui fait fonctionner les chatbots comme moi). Ces "cerveaux" numériques sont excellents pour reconstituer les pistes, mais ils sont lourds et complexes.

Le défi de l'article est le suivant : Comment faire tenir un cerveau d'IA aussi gros et complexe dans un petit robot artisan (la puce FPGA) ?

3. La Méthode : La "Découpe" et le "Montage"

Les chercheurs ont développé une méthode ingénieuse pour y parvenir, qu'ils appellent TrackCore-F.

• L'analogie du Puzzle : Imaginez que votre modèle d'IA est un immense puzzle de 1000 pièces. La puce FPGA est une petite boîte qui ne peut pas contenir tout le puzzle d'un coup.
  • Au lieu de tout essayer de mettre dedans, les chercheurs découpent le puzzle en plusieurs morceaux.
  • Ils mettent un morceau (une partie du calcul) dans la puce FPGA pour qu'il soit traité ultra-rapidement.
  • Le reste du puzzle reste dans l'ordinateur principal.
  • La puce fait son travail, renvoie le résultat, et l'ordinateur continue le reste. C'est comme un chef d'orchestre qui donne une section de la musique à un soliste (la puce) pour qu'il joue sa partie en solo, avant de reprendre l'ensemble.

4. Les Résultats et les Défis

Les chercheurs ont testé cette idée sur une puce réelle (une carte Zynq). Voici ce qu'ils ont découvert :

• C'est possible ! Ils ont réussi à faire fonctionner une partie de l'IA sur la puce.
• Le goulot d'étranglement : La puce a une mémoire très limitée (comme un petit tiroir). Ils ont découvert que la mémoire interne de la puce se remplit très vite. Si on veut mettre plus de "puzzle" (plus de couches de l'IA) dans la puce, il faut soit utiliser de la mémoire externe (ce qui ralentit un peu le processus), soit être très malin pour optimiser l'espace.
• Le piège de la "compression" (Quantification) : Pour gagner de la place, on pourrait essayer de "réduire" la précision des nombres utilisés par l'IA (passer de nombres très précis à des nombres simples, comme passer de "3,14159" à "3").
  • Résultat : Ça économise de la place, mais l'IA devient un peu "bête". Dans leur expérience, trop simplifier les nombres a fait chuter la précision de l'IA de 97% à 70%. C'est comme si un détective arrêtait de regarder les détails fins des empreintes digitales : il trouve plus vite, mais il se trompe souvent de coupable.

5. Pourquoi c'est important ?

Ce travail est une étape cruciale pour le futur de la physique.

• Avant : On attendait des jours pour analyser les collisions.
• Avec TrackCore-F : On pourrait analyser les collisions en temps réel, directement sur place dans le laboratoire.
• Bénéfice : On économise de l'énergie, on va plus vite, et on peut découvrir de nouvelles particules plus rapidement, comme si on avait un détective qui résout le crime pendant que le suspect est encore sur les lieux.

En résumé

Les auteurs ont créé une "boîte à outils" pour transformer de l'intelligence artificielle lourde en un logiciel léger capable de courir sur des puces électroniques spécialisées. C'est comme apprendre à un éléphant (l'IA) à danser sur une table de ping-pong (la puce FPGA) : c'est difficile, ça demande de la gymnastique et de la précision, mais si on y arrive, le spectacle est incroyable !

Résumé technique

1. Problématique et Contexte

La reconstruction de traces de particules dans les expériences de physique des hautes énergies (comme au LHC) est une tâche computationnelle intensive. Traditionnellement, cette tâche est effectuée « post-mortem » (après l'expérience) en raison des contraintes de temps de calcul. Cependant, avec l'arrivée de la phase à haute luminosité du LHC (HL-LHC), le volume de données et la fréquence des expériences augmentent, rendant le traitement en temps réel (ou quasi réel) indispensable.

Bien que les architectures basées sur l'apprentissage automatique (ML), et plus particulièrement les Transformers, aient démontré des performances supérieures pour des tâches comme l'étiquetage de jets et la reconstruction de traces, leur déploiement sur du matériel embarqué reste un défi majeur. Les solutions actuelles reposent souvent sur des GPU, qui ne sont pas toujours adaptés aux contraintes de latence et d'efficacité énergétique des détecteurs en ligne. Les FPGA (Field-Programmable Gate Arrays) offrent une alternative prometteuse pour une accélération matérielle efficace et une faible latence, mais leur utilisation pour les Transformers est entravée par :

• Le manque d'outils de support automatisés pour la synthèse de ces modèles.
• Les contraintes sévères des ressources matérielles (mémoire, logique).
• La difficulté de partitionner les grands modèles pour qu'ils tiennent sur une puce FPGA.

2. Méthodologie

Les auteurs proposent un flux de développement structuré pour déployer des modèles Transformers pré-entraînés sur un FPGA, en se concentrant sur l'inférence.

Modèles et Données :

• Architectures : Deux modèles issus du projet TrackFormers sont utilisés :
  • EncCla (Encoder-Classifier) : Un Transformer uniquement encodeur qui prédit l'association des points de détection (hits) à des étiquettes de classes prédéfinies (environ 1,5 million de paramètres).
  • EncReg (Encoder-Regressor) : Un Transformer uniquement encodeur qui régresse les paramètres des pistes potentielles, nécessitant un post-traitement par clustering (environ 76 000 paramètres).
• Données : Les modèles sont entraînés sur cinq jeux de données simulés, allant de traces linéaires simples à des configurations complexes de 200 à 500 pistes par événement (simulations REDVID et données TrackML).

Flux de Déploiement et Outils :
Le workflow utilise une chaîne d'outils spécifique pour convertir un modèle PyTorch en un noyau matériel synthétisable :

1. PyTorch & ONNX : Conversion du modèle pré-entraîné vers le format ONNX pour exposer les opérations de bas niveau.
2. Partitionnement Manuel : En raison de la taille des modèles, une partition manuelle est appliquée. Le modèle est divisé en trois segments : la partie avant la première couche d'encodage, la couche d'encodage cible (déployée sur le FPGA), et le reste du modèle.
3. Synthèse HLS (High-Level Synthesis) : La couche d'encodage cible est implémentée en C/C++ sous AMD Vitis HLS. Les poids sont transférés via une interface AXI4-Lite. Les opérations clés (Multi-Head Attention, Normalisation, Feed-Forward) sont optimisées avec des pragmas HLS.
4. Intégration et Synthèse : Le noyau IP est intégré dans l'environnement AMD Vivado avec le SoC Zynq MPSoC (ZCU102), connecté via le bus AXI4, et un flux de bits (bitstream) est généré.
5. Quantification : Une étude est menée sur l'impact de la quantification (INT8/INT16) des poids et des activations pour réduire la consommation de ressources.

3. Résultats Clés

Performance et Précision (Quantification) :
L'étude de la quantification révèle un compromis critique entre la réduction des ressources et la précision du modèle :

• Le modèle de base (précision 0,97) voit sa précision chuter drastiquement si les activations sont quantifiées en INT8 (précision tombant à 0,71).
• La quantification des poids seule a un impact moindre.
• La combinaison INT8 pour les poids et les activations donne une précision de 0,70.
• Conclusion : La quantification des activations est très préjudiciable à la précision pour ces modèles de suivi.

Consommation des Ressources (FPGA ZCU102) :
L'implémentation d'une seule couche d'encodage sur le FPGA Zynq UltraScale+ ZU9EG montre :

• BRAM (Mémoire Bloquée) : 38,04 % d'utilisation (347 sur 912). C'est la ressource limitante principale.
• LUT (Tableaux de Recherche) : 24,01 % d'utilisation.
• DSP : 2,66 % d'utilisation.
• FF (Flip-Flops) : 13,53 % d'utilisation.

Analyse de Scalabilité :
Bien que la BRAM soit le facteur limitant actuel, elle peut être réduite en utilisant de la mémoire DDR externe (au prix d'une latence accrue). Une fois la BRAM optimisée, les LUT deviennent le facteur limitant. Les auteurs estiment qu'au maximum 4 couches d'encodage peuvent être implémentées sur ce FPGA spécifique pour ce modèle.

4. Contributions Principales

1. Flux de travail de déploiement : Définition d'une méthodologie complète pour convertir des modèles Transformers PyTorch/ONNX en noyaux FPGA, en passant par la quantification et le partitionnement manuel.
2. Analyse du partitionnement : Démonstration qu'un déploiement partiel (une seule couche sur FPGA, le reste sur CPU) est une stratégie viable et précieuse pour accéder à des matériels plus abordables, même si elle est moins performante qu'un déploiement monolithique.
3. Évaluation de la quantification : Mise en évidence du fait que la quantification des activations dégrade sévèrement la précision des modèles de suivi de particules, suggérant que des techniques de quantification plus avancées ou une quantification sélective sont nécessaires.
4. Benchmark matériel : Fourniture de métriques précises sur l'utilisation des ressources FPGA pour les architectures Transformers appliquées à la physique des hautes énergies.

5. Signification et Impact

Ce travail est significatif car il comble le fossé entre les modèles d'IA de pointe (Transformers) et les contraintes matérielles rigides des expériences de physique des particules.

• Faisabilité : Il prouve que l'inférence de modèles complexes de reconstruction de traces est possible sur des FPGA, ouvrant la voie à un traitement en ligne (online) plus efficace.
• Efficacité Énergétique : Le déploiement sur FPGA promet une bien meilleure efficacité énergétique que les solutions GPU, un critère crucial pour les détecteurs embarqués.
• Direction Future : Le papier souligne que l'optimisation du modèle (quantification, partitionnement) est souvent plus critique que l'implémentation HLS elle-même. Il appelle à un développement d'outils automatisés pour gérer le partitionnement des grands modèles et à des recherches sur des méthodes de quantification préservant la précision.

En résumé, TrackCore-F établit une base solide pour l'adoption des Transformers dans le traitement en temps réel des données du LHC, tout en identifiant clairement les verrous technologiques restants.