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

Este artigo propõe o desenvolvimento de metodologias e ferramentas para a síntese de modelos Transformer baseados em TrackFormers em FPGAs, visando superar as limitações de recursos e ferramentas atuais para viabilizar a inferência de rastreamento de partículas subatômicas com baixa latência.

O essencial

Imagine que o Grande Colisor de Hádrons (LHC) é uma máquina do tempo e espaço gigantesca que colide partículas a velocidades incríveis. Quando elas colidem, criam uma "tempestade" de milhões de partículas novas. O trabalho dos físicos é reconstruir o caminho que cada uma dessas partículas fez, como se fosse tentar descobrir a rota de milhares de carros em um engarrafamento caótico, apenas olhando para as marcas que eles deixaram no asfalto.

Até agora, fazer esse "rastreamento" era como tentar resolver um quebra-cabeça gigante depois que a festa acabou (depois do experimento), porque os computadores comuns eram muito lentos para acompanhar o ritmo.

Aqui entra o TrackCore-F, um projeto que quer mudar as regras do jogo. Vamos explicar como eles estão fazendo isso usando analogias simples:

1. O Problema: O Cérebro vs. O Motor

Os cientistas estão usando uma tecnologia de Inteligência Artificial chamada Transformers (a mesma base dos grandes modelos de linguagem como o GPT). Eles são ótimos para entender padrões complexos, como rastrear partículas.

• O Desafio: Esses "cérebros" de IA são pesados. Normalmente, eles rodam em placas de vídeo gigantes (GPUs) que consomem muita energia e são grandes demais para caber dentro do detector de partículas, bem perto da colisão.
• A Solução: Eles querem colocar esse cérebro dentro de um FPGA. Pense no FPGA não como um computador comum, mas como um bloco de LEGO eletrônico reconfigurável. Você pode montar a estrutura exata que precisa para aquela tarefa específica, tornando-a super rápida e eficiente em energia.

2. A Missão: Encaixar um Elefante em uma Caixa de Sapatos

O maior problema é que os modelos de IA são grandes (como um elefante), e o FPGA é pequeno (uma caixa de sapatos).

• A Estratégia de "Fatias": Em vez de tentar colocar o elefante inteiro na caixa, os pesquisadores estão cortando o modelo em fatias. Eles pegam uma parte do modelo (uma "fatia" de camadas), a colocam dentro do FPGA para fazer o trabalho pesado lá, e deixam o resto do modelo rodando no computador normal.
• O Fluxo de Trabalho: Eles criaram um "receituário" (fluxo de desenvolvimento) para pegar o modelo treinado, cortá-lo, transformá-lo em código de hardware e montar no FPGA. É como se eles estivessem ensinando o bloco de LEGO a pensar como a IA.

3. O Experimento: Testando no "Banco de Prova"

Eles usaram um kit de desenvolvimento (ZCU102) que é como um "laboratório de testes" para ver se a ideia funciona.

• Os Modelos: Eles testaram dois tipos de "detetives" de IA:
  1. O Classificador (EncCla): Tenta adivinhar a qual grupo cada partícula pertence.
  2. O Regressor (EncReg): Tenta calcular exatamente a trajetória de cada partícula.
• Os Dados: Eles treinaram esses modelos com simulações de partículas, desde cenários simples (poucas partículas) até cenários caóticos (centenas de partículas, como será no futuro do LHC).

4. O Dilema da Precisão vs. Velocidade (Quantização)

Para caber no FPGA, eles tentaram "espremer" os dados. Imagine que você tem uma foto em alta definição (números com muitas casas decimais) e precisa enviá-la por um telefone antigo. Você precisa reduzir a qualidade da imagem para que ela caiba.

• O Resultado: Eles descobriram que, se reduzirem demais a qualidade (transformando números complexos em inteiros simples), o "detetive" começa a cometer erros.
  • Se reduzirem apenas os "pesos" (o conhecimento do modelo), ele ainda funciona bem.
  • Se reduzirem também as "atividades" (o raciocínio em tempo real), ele fica confuso e erra muito.
• A Lição: Às vezes, é melhor ter um modelo um pouco maior e preciso do que um super compacto e confuso.

5. O Que Eles Encontraram?

• Recursos Limitados: O FPGA tem um limite de "espaço de memória" (chamado BRAM). Eles conseguiram colocar uma camada do modelo de IA inteira no chip.
• O Futuro: Com os recursos atuais, eles conseguem colocar até 4 camadas desse modelo. Isso significa que, no futuro, eles poderão colocar partes maiores do "cérebro" da IA diretamente no detector.
• Conclusão: Mesmo que não consigam colocar o modelo inteiro de uma vez, colocar partes dele no FPGA já é um avanço enorme. Isso permite que o rastreamento aconteça quase em tempo real, economizando energia e permitindo que os físicos descubram novas partículas mais rápido.

Resumo em uma Frase

O TrackCore-F está ensinando um bloco de LEGO inteligente (FPGA) a fazer a parte mais difícil de um trabalho de detetive (rastrear partículas), cortando o problema em pedaços menores para que ele possa rodar rápido e eficiente dentro da máquina de física, sem precisar de computadores gigantes.

Resumo técnico

Resumo Técnico: TrackCore-F

1. O Problema

A reconstrução de trajetórias de partículas (tracking) em física de alta energia, especificamente no Grande Colisor de Hádrons (LHC) e na futura fase de Alta Luminosidade (HL-LHC), enfrenta desafios críticos de latência e escalabilidade.

• Limitações Atuais: O rastreamento preciso é tradicionalmente uma tarefa "pós-mortem" (offline) devido às exigências computacionais. Embora algoritmos de Aprendizado de Máquina (ML) tenham melhorado o desempenho, a escala e a frequência dos experimentos futuros podem anular esses ganhos.
• Dependência de Hardware: A execução eficiente de ML depende predominantemente de GPUs, que não são ideais para implantação on-site (no local) devido a restrições de energia e latência.
• Desafio de Implantação: A integração de arquiteturas Transformer (que têm mostrado sucesso em tarefas como jet tagging e reconstrução de trajetórias) em hardware especializado, como FPGAs, é complexa. As ferramentas atuais oferecem suporte limitado ou inexistente para síntese de Transformers, e os recursos de hardware (memória, lógica) são restrições severas para modelos grandes.

2. Metodologia

Os autores propõem uma metodologia estruturada para a síntese e implantação de modelos Transformer baseados no projeto TrackFormers em FPGAs, focando na inferência.

• Modelos Alvo: Dois designs de Encoder-only derivados do TrackFormers:
  1. EncCla (Encoder-Classifier): Classifica coordenadas de impactos (hits) em rótulos pré-definidos. Possui ~1,5 milhão de parâmetros.
  2. EncReg (Encoder-Regressor): Regressa parâmetros de trajetórias potenciais, exigindo um passo de pós-processamento (agrupamento/clustering via HDBSCAN). Possui ~76 mil parâmetros.
• Fluxo de Desenvolvimento (Pipeline):
  1. Conversão: Modelos PyTorch pré-treinados são convertidos para o formato ONNX (Open Neural Network Exchange) para expor operações de baixo nível.
  2. Particionamento: Devido ao tamanho do modelo e limitações de recursos, aplica-se um particionamento manual. O modelo é dividido em três partes: segmento anterior à camada de codificador, a camada de codificador alvo (para FPGA) e o segmento restante.
  3. Síntese de Kernel: O segmento alvo é implementado em C/C++ usando AMD Vitis HLS (High-Level Synthesis). Em vez de implementar blocos de operação ONNX individualmente, desenvolve-se um kernel monolítico para reduzir o uso de recursos.
  4. Integração e Síntese: O kernel é integrado ao System-on-Chip (SoC) Zynq UltraScale+ MPSoC (placa ZCU102) usando o AMD Vivado, conectando-se via barramento AXI4.
  5. Otimização: Uso de pragmas HLS para otimização de desempenho e exploração de quantização (INT8/INT16).

3. Contribuições Chave

• Fluxo de Trabalho Estruturado: Apresentação de um pipeline completo para levar modelos Transformer pré-treinados do PyTorch para a implantação em FPGA, lidando com a conversão ONNX e a síntese HLS.
• Estratégia de Particionamento: Demonstração de que a implantação parcial (apenas camadas específicas no FPGA) é uma alternativa viável e valiosa à implantação monolítica completa, permitindo inferência em hardware mais acessível.
• Análise de Quantização: Investigação detalhada do impacto da quantização (pesos e ativações) na precisão do modelo, revelando que a quantização de ativações é mais prejudicial à acurácia do que a quantização de pesos.
• Mapeamento de Recursos: Avaliação rigorosa do consumo de recursos de hardware para camadas de Transformer em FPGAs, identificando gargalos específicos.

4. Resultados

Os testes foram realizados na placa de avaliação ZCU102 (Zynq UltraScale+ ZU9EG).

• Desempenho de Precisão (Quantização):

  • O modelo base (float) atingiu 97% de precisão.
  • A quantização de pesos para INT8 (mantendo ativações em INT16) reduziu a precisão para 90%.
  • A quantização de ativações para INT8 (mesmo com pesos em INT16) causou uma queda drástica para 71%.
  • A quantização total (INT8/INT8) resultou em 70%.
  • Conclusão: A quantização de ativações é o fator mais crítico para a perda de acurácia.

• Consumo de Recursos (Camada Única de Encoder):

  • BRAM (Memória de Bloco): 38,04% de utilização (347 de 912). Identificado como o fator limitante primário.
  • LUTs (Tabelas de Busca): 24,01% de utilização.
  • DSPs: 2,66% de utilização.
  • Flip-Flops (FFs): 13,53% de utilização.
  • Capacidade de Escala: Com base no uso de BRAM e LUTs de uma única camada, estima-se que no máximo 4 camadas de encoder possam ser implementadas simultaneamente neste FPGA específico para este modelo.

• Latência e Eficiência: A abordagem permite latências online ou pseudo-online, com potencial significativo de eficiência energética comparado a GPUs.

5. Significado e Conclusão

O artigo demonstra que a implantação de redes Transformer complexas em FPGAs para física de partículas é viável, mas exige otimizações cuidadosas além da simples síntese de hardware.

• Viabilidade: A implantação parcial é tão valiosa quanto a completa, pois democratiza o acesso a hardware de inferência de ML em ambientes de física de alta energia.
• Importância da Otimização: As otimizações de algoritmo (como a escolha de quantização e particionamento) têm um impacto maior no resultado final do que a implementação HLS em si.
• Caminho Futuro: O trabalho destaca a necessidade de ferramentas mais automatizadas para partição de modelos e a exploração de memórias externas (DDR) para contornar as limitações de BRAM, permitindo a implantação de modelos maiores no futuro.

Este estudo é um passo fundamental para a transição do rastreamento de partículas de processos offline para sistemas de baixa latência em tempo real, essenciais para a próxima geração de experimentos no LHC.