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

Este artículo presenta el desarrollo de metodologías y herramientas para sintetizar modelos de transformadores basados en el proyecto TrackFormers en FPGAs, abordando las limitaciones de recursos y la falta de soporte actual para implementar el seguimiento de partículas subatómicas con latencia en línea.

Lo esencial

¡Hola! Imagina que el Gran Colisionador de Hadrones (LHC) es como una ciudad gigante y caótica donde miles de partículas (los "ciudadanos") chocan entre sí a velocidades increíbles. El trabajo de los físicos es como el de un detective que debe seguir el rastro de cada ciudadano para saber quién es y hacia dónde va. A esto le llamamos "rastreo de partículas".

Hasta ahora, para hacer este trabajo, los científicos usaban computadoras muy potentes (como las que tienen tarjetas gráficas de videojuegos) que tardaban un poco en procesar toda la información. Pero el LHC va a volverse aún más rápido y producirá más datos que nunca. ¡Las computadoras actuales se quedarían cortas!

Aquí es donde entra el proyecto TrackCore-F. Su misión es construir un "detective robótico" super rápido que viva dentro de un chip especial llamado FPGA (una especie de cerebro electrónico reconfigurable que se puede instalar directamente en el laboratorio).

Aquí te explico cómo lo hacen, usando analogías sencillas:

1. El Detective Inteligente (La IA Transformer)

Antes, los detectives usaban reglas fijas. Ahora, usan una tecnología llamada Transformer (la misma que usan los chatbots como yo).

• La analogía: Imagina que en lugar de mirar una sola huella dactilar, el detective mira toda la escena del crimen de un solo vistazo y entiende la historia completa al instante.
• En este papel, usan dos tipos de detectives:
  • El Clasificador (EncCla): Mira las huellas y dice: "¡Esta huella pertenece al grupo A, esta al grupo B!".
  • El Regresor (EncReg): Mira las huellas y calcula: "Esta huella viene de un coche que iba a esta velocidad y en esta dirección".

2. El Reto: Meter un Elefante en un Frasco (FPGAs)

El problema es que estos "detectives" (los modelos de IA) son muy grandes y complejos, pero el chip FPGA (el robot) es pequeño y tiene recursos limitados. Es como intentar meter un elefante en un frasco de mermelada.

• La solución: En lugar de meter al elefante entero, lo cortan en trozos.
• La analogía: Imagina que tienes una receta de cocina muy larga. No puedes cocinarla toda en una sola sartén pequeña. Así que cocinas la primera parte en la sartén (el chip FPGA) y luego pasas el plato a la mesa para terminar el resto.
• El equipo creó una herramienta para cortar el modelo de IA en "rebanadas". Una rebanada se instala en el chip para hacer el trabajo pesado y rápido, y el resto se deja para la computadora normal.

3. El Proceso de Construcción (La Fábrica)

Para convertir el cerebro de IA en un chip físico, usan una cadena de montaje especial:

1. Traducción: Primero toman el modelo (escrito en un lenguaje de programación llamado PyTorch) y lo traducen a un idioma universal (ONNX).
2. Diseño: Usan herramientas (Vitis HLS) para convertir ese diseño en instrucciones que el chip puede entender (como convertir un plano de arquitectura en ladrillos reales).
3. Ensamblaje: Finalmente, usan un software (Vivado) para pegar todas las piezas y crear el "bitstream" (el código final que hace funcionar el chip).

4. El Dilema de la Precisión vs. Velocidad (Cuantización)

Para que el chip sea más rápido y gaste menos energía, los científicos intentaron "simplificar" los números que usa el detective.

• La analogía: Imagina que el detective mide la distancia con una regla de milímetros (precisión alta). Para ir más rápido, deciden usar una regla de centímetros (precisión baja).
• El resultado: ¡Funcionó rápido, pero el detective empezó a cometer errores!
  • Si simplifican solo los "pesos" (la memoria del detective), sigue siendo bastante bueno.
  • Si simplifican también las "activaciones" (lo que el detective piensa en el momento), ¡se vuelve torpe y pierde muchos casos!
  • Lección: No puedes sacrificar demasiada precisión solo por velocidad; el detective necesita ver bien para no confundirse.

5. ¿Cuánto cabe en el chip?

El chip tiene un espacio limitado de "estanterías" (memoria) y "herramientas" (lógica).

• Descubrieron que la "estantería" (memoria BRAM) es lo que más se llena.
• Con el chip que probaron (ZCU102), solo podían instalar 4 capas de este detective inteligente. Si intentaran poner más, el chip se quedaría sin espacio.
• La solución creativa: Si necesitan más, pueden conectar el chip a una memoria externa (como un disco duro), pero eso haría que el detective tuviera que caminar más lejos para buscar sus herramientas, lo que lo haría un poco más lento.

En Resumen

Este paper nos cuenta cómo los científicos están aprendiendo a cortar y adaptar sus inteligencias artificiales más avanzadas para que quepan en chips pequeños y rápidos instalados directamente en los detectores de partículas.

No es solo poner un cerebro gigante en un chip pequeño; es como reconstruir ese cerebro pieza por pieza para que quepa en un reloj de pulsera, asegurándose de que, aunque sea pequeño, siga siendo lo suficientemente inteligente para encontrar nuevas partículas en el caos del universo. ¡Es un paso gigante para que la física de partículas sea más rápida y eficiente!

Resumen técnico

1. El Problema

El seguimiento de partículas en física de altas energías (como en el Gran Colisionador de Hadrones, LHC) es una tarea computacionalmente intensiva que tradicionalmente se realiza post-mortem (después de la toma de datos). Con la llegada de la fase de Alta Luminosidad (HL-LHC), la escala y frecuencia de los experimentos aumentarán, haciendo insostenible el procesamiento actual.

• Limitaciones de Hardware: Aunque las GPU dominan el despliegue de modelos de Machine Learning (ML), no son ideales para latencias en línea o pseudo-en línea debido a su consumo energético y falta de integración en el sitio.
• Desafío de los Transformadores: Las arquitecturas Transformer han demostrado ser superiores en tareas como la reconstrucción de trayectorias (tracking), pero su despliegue en hardware especializado, específicamente en FPGAs (Field-Programmable Gate Arrays), es complejo.
• Barreras Técnicas: Las herramientas actuales de síntesis tienen un soporte limitado o nulo para modelos Transformer grandes. Además, los recursos de las FPGAs (memoria BRAM, lógica LUT, DSP) son restrictivos, obligando a particionar modelos grandes de manera manual y no automatizada.

2. Metodología

Los autores proponen un flujo de trabajo estructurado para la síntesis y despliegue de modelos Transformer pre-entrenados en FPGAs, enfocándose en la inferencia.

• Modelos Seleccionados: Se utilizan dos arquitecturas derivadas del proyecto TrackFormers:
  1. EncCla (Encoder-Classifier): Un diseño de solo codificador que asocia hitos a etiquetas de clase predefinidas (aprox. 1.5M de parámetros).
  2. EncReg (Encoder-Regressor): Un diseño de solo codificador que regresa parámetros de trayectorias potenciales, requiriendo un paso de post-procesamiento con clustering (HDBSCAN) (aprox. 76k parámetros).
• Flujo de Desarrollo:
  1. Entrenamiento y Conversión: Los modelos se entrenan en PyTorch y se convierten a formato ONNX para exponer operaciones de bajo nivel.
  2. Particionamiento Manual: Dado que los modelos completos no caben en la FPGA, se divide el modelo en tres partes: segmento previo, capa de codificador (kernel) y segmento restante.
  3. Síntesis de Kernel: Se utiliza AMD Vitis HLS (C/C++) para desarrollar el kernel de la capa de codificador (Atención Multi-Cabeza, Normalización, Capa Feed-Forward). Se utilizan pragmas de HLS para optimización.
  4. Integración: El IP generado se integra en AMD Vivado con el sistema MPSoC (Zynq UltraScale+) y se conecta mediante buses AXI4.
  5. Hardware: Se utiliza la placa de evaluación ZCU102 (ZU9EG) con un procesador Cortex-A53 y lógica programable.
• Cuantización: Se investigó el uso de cuantización (INT8/INT16) para reducir recursos, evaluando su impacto en la precisión.

3. Contribuciones Clave

• Flujo de Despliegue Híbrido: Presentación de un flujo de trabajo viable que combina herramientas de software (PyTorch, ONNX) y hardware (Vitis HLS, Vivado) para implementar selectivamente capas de Transformadores en FPGA.
• Estrategia de Particionamiento: Demostración de que el despliegue parcial (solo capas críticas en FPGA) es una alternativa valiosa y factible al despliegue monolítico completo, permitiendo inferencia en hardware más accesible.
• Análisis de Cuantización: Evaluación detallada de cómo la cuantización de pesos y activaciones afecta la precisión del modelo en el contexto de seguimiento de partículas.
• Optimización de Recursos: Identificación de los cuellos de botella específicos en FPGAs para modelos Transformer, más allá de la simple lógica.

4. Resultados

• Precisión vs. Cuantización:
  • El modelo base tiene una precisión del 0.97.
  • La cuantización de activaciones tiene un impacto devastador en la precisión (caída a 0.71-0.70 si se cuantizan a INT8), mientras que la cuantización solo de pesos tiene un efecto menor (0.90).
  • Esto sugiere que mantener las activaciones en punto flotante (o alta precisión) es crucial para este tipo de modelos.
• Consumo de Recursos (ZCU102):
  • Para una sola capa de codificador, el uso de BRAM (memoria de bloque) es el factor limitante, alcanzando un 38.04% de utilización (347 de 912 bloques).
  • El uso de LUTs es del 24.01% y DSPs del 2.66%.
  • Conclusión de Escalabilidad: Dado el uso de BRAM, se estima que es posible implementar un máximo de 4 capas de codificador en esta FPGA específica para este modelo antes de agotar los recursos.
• Latencia y Eficiencia: Aunque no se reportan números exactos de latencia en el resumen, el objetivo es lograr latencias en línea o pseudo-en línea, aprovechando la eficiencia energética de las FPGA frente a las GPU.

5. Significancia

Este trabajo es fundamental para la física de altas energías futura (HL-LHC) por varias razones:

• Viabilidad de ML en Línea: Demuestra que es posible llevar algoritmos avanzados de ML (Transformadores) a hardware embebido y de baja latencia (FPGAs), lo cual es esencial para la toma de decisiones en tiempo real en los detectores.
• Eficiencia Energética: Las FPGA ofrecen un potencial de eficiencia energética muy superior al de las GPU, un factor crítico para instalaciones de gran escala.
• Guía para la Industria: Proporciona un marco metodológico para superar las limitaciones de las herramientas actuales de síntesis, mostrando que el particionamiento manual y la optimización de la implementación HLS son tan importantes como el diseño del modelo mismo.
• Advertencia sobre Cuantización: Alerta a la comunidad de que las técnicas estándar de cuantización (INT8) pueden degradar severamente el rendimiento en tareas de física de precisión, requiriendo estrategias híbridas o de precisión mixta.

En resumen, el artículo establece un paso crucial hacia la implementación práctica de IA avanzada en la infraestructura de hardware de los futuros experimentos de física de partículas.