Triton for MTIA: Bridging the Programming Model Gaps for Custom AI Accelerators
Este artigo apresenta a primeira implantação em escala de produção da linguagem de programação Triton no acelerador customizado MTIA-2i da Meta, demonstrando que um novo backend de compilador e extensões mínimas de linguagem permitem o desenvolvimento de kernels de alto desempenho e eficientes que reduzem a lacuna entre frameworks de ML e hardware customizado.
Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo
Imagine que você esteja tentando construir um sistema de trens massivo e de alta velocidade para transportar passageiros (dados) através de um continente. Durante anos, os únicos trens disponíveis eram os elétricos padrão, de alta tecnologia, que todos sabiam dirigir (GPUs). Eles tinham portas automáticas, GPS integrado e uma viagem suave. Mas recentemente, um novo tipo de motor de trem foi inventado (aceleradores de IA customizados). Esses novos motores são incrivelmente poderosos e construídos especificamente para carga pesada, mas são estranhos. Eles não têm portas automáticas; você tem que abri-las e fechá-las manualmente. Eles não têm GPS; você tem que desenhar o mapa você mesmo. E eles não rodam nos mesmos trilhos dos trens antigos.
O problema é que os engenheiros que constroem os vagões de trem (os desenvolvedores de software) estão acostumados com os trens elétricos antigos e fáceis de dirigir. Se eles quiserem usar os novos motores supervelozes, terão que aprender uma maneira completamente nova e difícil de dirigir, ou terão que construir um vagão de trem totalmente novo para cada nova rota. Isso atrasa tudo. A grande questão é: podemos ensinar o estilo de condução antigo e fácil de usar a funcionar nestes novos motores estranhos sem perder a vantagem de velocidade? Este artigo explora exatamente isso, tentando preencher a lacuna entre o mundo familiar do software de IA padrão e o hardware estranho e customizado construído por empresas como a Meta para executar os seus modelos de IA.
A História: Ensinando um Novo Motor a Falar uma Língua Antiga
A Meta, a empresa por trás do Facebook e do Instagram, tem construído os seus próprios chips especiais chamados MTIA-2i para executar os seus modelos de IA de forma mais rápida e barata. Pense no MTIA-2i como um motor de carro de corrida construído à medida. É incrível no que faz, mas é construído de forma diferente dos motores padrão (GPUs) com os quais a maioria dos programadores de IA está habituada.
No mundo da IA, existe uma linguagem de programação popular chamada Triton. Pode pensar no Triton como um "controlo remoto universal" para a IA. Em vez de escrever código complexo e confuso para dizer a uma GPU padrão como mover dados, os desenvolvedores escrevem código simples e limpo em Triton, e um compilador inteligente traduz isso para a linguagem de máquina que a GPU entende. É como falar inglês para um tradutor que sabe como falar com a máquina.
No entanto, o Triton foi originalmente desenhado apenas para motores de GPU padrão. Quando a Meta tentou utilizá-lo no seu motor de carro de corrida customizado MTIA-2i, não funcionou. O "controlo remoto" não encaixava nos botões do novo motor. O motor MTIA-2i trabalha de forma assíncrona (ele recebe ordens e executa-as mais tarde), gere a sua própria memória de uma forma estranha de "buffer circular" e precisa de instruções muito específicas para manter tudo a fluir suavemente. O código Triton padrão simplesmente não conseguia lidar com estas peculiaridades.
O Que o Artigo Faz: Construindo um Novo Adaptador
A equipa da Meta decidiu resolver isto. Eles não se limitaram a forçar o controlo remoto antigo a funcionar; construíram um novo backend de compilador especificamente para o MTIA-2i. Este novo sistema atua como um adaptador superinteligente. Ele pega no código Triton simples e fácil de ler que os desenvolvedores escrevem e traduz-no perfeitamente nas instruções complexas de baixo nível que o motor MTIA-2i necessita.
Aqui está como eles fizeram funcionar, usando algumas analogias divertidas:
- A Linha de Buffet "FIFO": Nas GPUs padrão, o computador gere a memória automaticamente como um buffet inteligente onde os pratos são agarrados e substituídos instantaneamente. No MTIA-2i, é mais como uma linha de montagem manual com um buffer "First-In, First-Out" (FIFO). O novo compilador aprendeu a gerir esta linha perfeitamente, garantindo que os dados cheguem exatamente quando a máquina precisa deles, sem que o programador tenha de empurrar manualmente cada um dos pratos.
- A Cozinha de Dois Chefs: O chip MTIA-2i tem dois "núcleos" (como dois chefs numa cozinha) que podem trabalhar ao mesmo tempo. O compilador descobriu como dividir as tarefas de cozedura entre eles para que não batam uns nos outros nem fiquem à espera. Ele permite até que programadores especialistas digam aos chefs exatamente quem deve picar as cebolas e quem deve mexer a sopa, se quiserem ser mais sofisticados.
- A Rota de Entrega em "Zigue-zague": Imagine que tem de entregar encomendas a 64 casas. Se apenas for pelas ruas em linha reta (linear), as primeiras casas ficam prontas rapidamente, mas as últimas esperam para sempre. A equipa descobriu que uma rota de entrega em "zigue-zague" (indo e voltando) equilibrava muito melhor o trabalho. Eles adicionaram uma funcionalidade ao Triton para que os programadores pudessem escolher facilmente esta rota mais inteligente.
Os Resultados: Rápido, Flexível e Pronto para o Mundo Real
A equipa não construiu isto apenas num laboratório; eles colocaram-no a trabalhar no mundo real. Eles implementaram com sucesso estes kernels Triton no ambiente de produção em 60 tipos diferentes de modelos de IA.
Eis o que descobriram:
- Velocidade: O código escrito na linguagem simples Triton correu tão rápido como o código escrito por especialistas em C++ de baixo nível (a linguagem "nativa" do chip). De facto, para General Matrix Multiplication (GEMM) — uma tarefa matemática pesada — o código Triton atingiu mais de 80% da velocidade máxima teórica do chip.
- Adoção: Num único trimestre, o número de tipos de modelos que utilizam Triton triplicou.
- Impacto: O Triton agora lida com 50% das camadas e 47% do tempo de execução não-GEMM destes modelos. Isso é quase metade do trabalho sendo feito por este novo sistema de fácil utilização!
Por Que Isto Importa
O artigo argumenta que não é necessário escolher entre "fácil de escrever" e "super rápido". Antes disto, se quisesse utilizar um chip customizado como o MTIA, tinha de ser um mago da codificação de baixo nível e escrever tudo do zero. Se quisesse utilizar uma linguagem de alto nível como o Triton, ficava preso às GPUs padrão.
Ao colmatar esta lacuna, a Meta demonstrou que as linguagens de alto nível podem ser adaptadas para funcionar em quase qualquer hardware customizado. Isto significa que, no futuro, à medida que as empresas constroem chips de IA ainda mais estranhos e especializados, os desenvolvedores não precisarão de aprender uma linguagem nova e difícil para cada um deles. Podem continuar a usar as ferramentas familiares e poderosas que já conhecem, enquanto o compilador trata dos detalhes complicados do novo hardware.
O artigo também nota que tentaram anteriormente uma linguagem diferente e de muito baixo nível chamada KNYFE, mas abandonaram-na porque era demasiado difícil para a maioria das pessoas aprender. Eles provaram que manter uma linguagem flexível e de alto nível como o Triton, mesmo com alguns pequenos ajustes customizados, é a estratégia vencedora para manter o desenvolvimento de IA rápido e eficiente.
Em suma, o artigo demonstra que, com o "truque de magia" do compilador certo, podemos fazer com que o "controlo remoto universal" funcione em quase qualquer "televisão", mesmo naquelas estranhas e construídas à medida, sem sacrificar a qualidade da imagem.
Afogado em artigos na sua área?
Receba digests diários dos artigos mais recentes que correspondam às suas palavras-chave de pesquisa — com resumos técnicos, no seu idioma.