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

Dieser Beitrag stellt Methoden und Werkzeuge zur Synthese von Transformer-Modellen für die Subatomare Teilchenspurverfolgung auf FPGAs vor, wobei insbesondere die Herausforderungen bei der Partitionierung großer Modelle und die begrenzte Tool-Unterstützung adressiert werden.

Das Wesentliche

Die große Jagd nach den kleinsten Teilchen

Stell dir vor, der Large Hadron Collider (LHC) ist eine riesige, superschnelle Rennbahn, auf der winzige Teilchen wie Autos in entgegengesetzte Richtungen rasen und dann frontal zusammenstoßen. Bei jedem Crash entstehen tausende neuer „Autos" (Teilchen), die in alle Richtungen davonfliegen.

Die Aufgabe der Detektoren (wie ATLAS oder CMS) ist es, diesen Chaos zu sortieren. Sie müssen genau nachverfolgen, welches Teilchen wohin geflogen ist. Das nennt man Spurverfolgung (Tracking).

Bisher war das wie ein riesiges Puzzle, das man erst nach dem Rennen am Computer zusammengefügt hat. Das dauert aber zu lange. Die Physiker wollen das aber live während des Rennens tun, um sofort zu wissen, ob etwas Spannendes passiert ist.

Das Problem: Der Computer ist zu langsam, die Hardware zu klein

Um diese riesigen Datenmengen live zu verarbeiten, braucht man extrem schnelle Rechenkraft. Normalerweise nutzt man dafür riesige Grafikkarten (GPUs), die wie ein ganzer Server-Raum aussehen. Aber diese sind zu groß, zu teuer und verbrauchen zu viel Strom, um sie direkt in das Teilchen-Experiment zu stellen.

Die Lösung? FPGAs.
Stell dir einen FPGA nicht als fertigen Computer vor, sondern als einen riesigen, leeren Baukasten aus Legosteinen. Du kannst diese Steine genau so zusammenstecken, wie du sie für deine spezielle Aufgabe brauchst. Sie sind klein, schnell und extrem stromsparend.

Die Herausforderung: Der Transformer ist ein Riese

In den letzten Jahren haben Wissenschaftler eine sehr clevere Methode entwickelt, um diese Teilchen zu verfolgen: Transformer-Modelle (die gleiche KI-Technologie, die auch Chatbots wie ich antreibt). Diese Modelle sind aber wie riesige Elefanten. Sie brauchen viel Platz und viel Kraft.

Das Problem: Ein FPGA ist wie ein kleiner Garten. Ein ganzer Elefant (das große KI-Modell) passt da einfach nicht rein.

Die Lösung von TrackCore-F: Der clevere Umzug

Die Autoren des Papers (Arjan, Uraz und Amirreza) haben einen neuen Weg gefunden, um diese KI auf den kleinen FPGA zu bringen. Hier ist ihr Plan, einfach erklärt:

1. Das Schneiden (Partitionierung)
Da der ganze Elefant nicht in den Garten passt, schneiden sie ihn in handliche Stücke. Sie nehmen das große KI-Modell und teilen es in Teile auf. Ein Teil läuft auf dem normalen Computer (dem „Gehirn" des Systems), und ein spezieller, sehr schneller Teil wird direkt in die Legosteine des FPGA eingebaut.

• Analogie: Stell dir vor, du hast einen riesigen Kochrezept. Du machst den ganzen Salat selbst, aber den komplizierten Teil, das Schneiden der Zwiebeln, gibst du einem speziellen Roboter-Messer (dem FPGA), das das extrem schnell erledigt, bevor du den Rest fertigstellst.

2. Der Bau (Synthese)
Sie haben einen neuen Bauplan entwickelt. Sie nehmen das KI-Modell, schneiden ein Stück heraus (z. B. eine Schicht des neuronalen Netzwerks) und programmieren die Legosteine so um, dass sie genau diese eine Schicht ausführen. Das passiert mit speziellen Werkzeugen (Vitis HLS und Vivado), die aus Computercode direkt die Schaltungen für den FPGA bauen.

3. Das Ergebnis: Schnell und sparsam
Sie haben getestet, wie gut das funktioniert:

• Platz: Ein einziger „Elefanten-Teil" (eine Encoder-Schicht) braucht schon fast 40 % des verfügbaren Speichers auf dem FPGA. Das zeigt, wie voll die Kiste ist. Man könnte also maximal vier solcher Teile auf einmal unterbringen.
• Genauigkeit: Um Platz zu sparen, haben sie versucht, die Zahlen im Modell zu vereinfachen (von „genauen Dezimalzahlen" auf „ganze Zahlen" zu runden, ähnlich wie man beim Einkaufen auf ganze Euro rundet). Das Ergebnis war ernüchternd: Je mehr man rundete, desto mehr verlor die KI ihre Treffsicherheit. Ein zu stark vereinfachtes Modell verpasst die Spuren der Teilchen.

Warum ist das wichtig?

Bisher mussten Physiker warten, bis die Daten vom Supercomputer analysiert waren. Mit TrackCore-F können sie die KI direkt in das Experiment einbauen.

• Vorteil: Die Ergebnisse kommen sofort (Echtzeit).
• Vorteil: Es braucht viel weniger Strom.
• Vorteil: Es ist möglich, auch auf kleinerer Hardware zu arbeiten, nicht nur auf riesigen Servern.

Fazit

Das Paper zeigt, wie man einen riesigen KI-Elefanten in kleine, handliche Stücke schneidet, um ihn in einen kleinen, effizienten FPGA-Garten zu pflanzen. Es ist noch nicht perfekt (man verliert etwas Genauigkeit, wenn man zu stark vereinfacht), aber es ist ein riesiger Schritt in Richtung einer Zukunft, in der Teilchenphysiker ihre Experimente live und direkt am Ort des Geschehens analysieren können.

Technische Zusammenfassung

1. Problemstellung

Die Teilchenspurverfolgung (Tracking) in der Hochenergiephysik, insbesondere am Large Hadron Collider (LHC) und zukünftig am High-Luminosity LHC (HL-LHC), stellt immense Anforderungen an die Rechenleistung. Traditionell ist die präzise Spurverfolgung eine nachgelagerte („post-mortem") Aufgabe aufgrund der hohen Rechenkomplexität.

• Herausforderung: Der Einsatz von Machine-Learning (ML)-Modellen, speziell Transformern, hat zwar die Genauigkeit und Effizienz der Offline-Analyse verbessert, führt jedoch zu einem Anstieg der Datenmenge und Frequenz, der die Vorteile teilweise wieder aufhebt.
• Hardware-Limitierung: Der aktuelle Standard für ML-Deployment sind GPUs, die jedoch oft nicht für den Echtzeit-Einsatz („online" oder „pseudo-online") an den Detektoren geeignet sind.
• Ziel: Die Entwicklung von Lösungen für Feld-programmierbare Gate-Arrays (FPGAs), die eine energieeffiziente, on-site-Deployment mit niedrigen Latenzzeiten ermöglichen.
• Spezifisches Hindernis: Die direkte Synthese großer Transformer-Modelle auf FPGAs ist durch begrenzte Hardware-Ressourcen (BRAM, LUTs, DSPs) und das Fehlen geeigneter automatisierter Tools für die Partitionierung erschwert.

2. Methodik und Ansatz

Das Paper beschreibt einen strukturierten Entwicklungs-Workflow für die Implementierung von Transformer-basierten Tracking-Modellen auf einem AMD Zynq UltraScale+ MPSoC (ZCU102).

• Modelle: Es werden zwei Architekturen aus dem Projekt TrackFormers untersucht:
  1. EncCla (Encoder-Classifier): Ein reiner Encoder, der Hit-Koordinaten in vordefinierte Klassen (Track-Parameter-Bins) klassifiziert. (~1,5 Mio. Parameter).
  2. EncReg (Encoder-Regressor): Ein reiner Encoder, der Track-Parameter direkt regressiert. Eine nachgelagerte Clustering-Phase (HDBSCAN) ist notwendig. (~76.500 Parameter).
• Entwicklungs-Flow:
  1. Konvertierung: Trainierte PyTorch-Modelle werden in das ONNX-Format konvertiert, um Low-Level-Operationen und Gewichte zu extrahieren.
  2. Partitionierung: Da das gesamte Modell oft zu groß für die FPGA-Ressourcen ist, erfolgt eine manuelle Partitionierung. Das Modell wird in drei Teile gesplittet:
     • Segment vor dem Encoder-Layer.
     • Der Encoder-Layer selbst (wird als Hardware-Kernel auf dem FPGA implementiert).
     • Restliches Modell.
  3. Synthese: Der ausgewählte Layer wird in C/C++ mit AMD Vitis HLS (High-Level Synthesis) entwickelt und in RTL (Register Transfer Level) für das FPGA synthetisiert.
  4. Integration: Der generierte IP-Core wird in Vivado mit dem MPSoC-System integriert (Verbindung über AXI4-Bus) und ein Bitstream wird generiert.
  5. Optimierung: Es werden HLS-Pragmas zur Leistungssteigerung eingesetzt. Statt einzelne ONNX-Operationen zu implementieren, wird angestrebt, einen monolithischen Kernel für den gesamten Layer zu erstellen, um Ressourcen zu sparen.

3. Schlüsselbeiträge

• Workflow-Etablierung: Präsentation eines funktionsfähigen Workflows für die partielle oder vollständige Synthese von Transformer-Modellen auf FPGAs, der die Lücke zwischen ML-Training und Hardware-Deployment schließt.
• Strategie zur Partitionierung: Demonstration, dass eine Aufteilung des Modells (Partial Deployment) eine valide Alternative zum Full Deployment ist, um auch auf ressourcenbeschränkter Hardware Inference durchzuführen.
• Analyse der Quantisierung: Untersuchung des Einflusses von Quantisierung (INT8 vs. INT16) auf die Genauigkeit und Ressourcennutzung.
• Ressourcen-Benchmarking: Detaillierte Analyse der FPGA-Ressourcennutzung für einen einzelnen Encoder-Layer.

4. Ergebnisse

• Genauigkeit vs. Quantisierung:
  • Das Basismodell (FP16/FP32) erreicht eine Vorhersagegenauigkeit von 0,97.
  • Eine Quantisierung der Aktivierungen hat einen drastischen negativen Einfluss auf die Genauigkeit (Absinken auf 0,71 bei INT8-Aktivierungen), während die Quantisierung nur der Gewichte weniger kritisch ist.
  • Die Kombination aus INT8-Gewichten und INT8-Aktivierungen führt zu einem Genauigkeitsverlust auf 0,70. Dies zeigt, dass aggressive Quantisierung für diese Aufgabe derzeit problematisch ist.
• Ressourcennutzung (ZCU102):
  • Ein einzelner Encoder-Layer verbraucht 38,04 % der verfügbaren Block-RAMs (BRAMs), was den primären Engpass darstellt.
  • Die Nutzung von LUTs (Look-Up Tables) liegt bei 24,01 %.
  • Fazit: Aufgrund der BRAM-Begrenzung können maximal 4 Encoder-Layer gleichzeitig auf diesem spezifischen FPGA implementiert werden. Die Nutzung von externem DDR-Speicher könnte die BRAM-Nutzung senken, würde aber die Latenz erhöhen.
• Performance-Potenzial: Der Ansatz zeigt, dass Transformer-basierte Tracking-Algorithmen auf FPGAs realisierbar sind, wobei die Optimierung der Datenfluss-Architektur und die Wahl der Partitionierung wichtiger sind als die reine HLS-Implementierung.

5. Bedeutung und Ausblick

Das Paper unterstreicht die Notwendigkeit, ML-basierte Tracking-Lösungen für die nächste Generation von Teilchenbeschleunigern (HL-LHC) auf FPGAs zu portieren, um Echtzeit-Entscheidungen und Energieeffizienz zu gewährleisten.

• Praktische Relevanz: Die Ergebnisse zeigen, dass FPGAs eine vielversprechende Alternative zu GPUs für den Einsatz direkt am Detektor sind, auch wenn die Modellgröße durch Partitionierung angepasst werden muss.
• Herausforderungen: Die hohe Sensitivität der Transformer-Architekturen gegenüber der Quantisierung der Aktivierungen ist eine wichtige Erkenntnis; zukünftige Arbeiten müssen hier robustere Quantisierungsmethoden oder spezielle Hardware-Optimierungen entwickeln.
• Zukunft: Der entwickelte Workflow bietet eine Basis für die Automatisierung der Partitionierung und die Erweiterung auf komplexere Modelle, was entscheidend für die zukünftige Datenverarbeitung in der Teilchenphysik sein wird.