Stanford apresenta Onyx, chip para IA esparsa mais rápido que CPUs e 70 vezes mais econômico

Pesquisadores de Stanford apresentaram o Onyx, um acelerador capaz de tirar proveito dos "zeros" dentro dos modelos de IA. A ideia é não gastar energia com multiplicações e somas que, de antemão, não mudam nada e, assim, acelerar os cálculos sem abrir mão de modelos grandes.

Por que os zeros importam

Em redes neurais, dados, pesos e ativações são armazenados como arrays de números: vetores, matrizes e tensores. Em muitos casos, uma parcela significativa desses números é zero ou está tão próxima de zero que pode ser tratada como tal sem perda perceptível de qualidade. Essa propriedade se chama esparsidade.

Se os zeros passam de metade, o modelo já pode se beneficiar de algoritmos especiais: em vez de armazenar e processar a matriz inteira, o sistema guarda apenas os valores úteis e pula as operações vazias.

Para a indústria, isso importa porque os modelos crescem mais rápido do que a infraestrutura fica barata. Mais parâmetros significam mais qualidade, mas também custo maior de execução em tempo, energia e pegada de carbono. O artigo cita o exemplo do Meta Llama com 2 trilhões de parâmetros. Os pesquisadores também lembram os resultados da Cerebras: no Llama 7B, foi possível zerar de 70 a 80 por cento dos parâmetros sem perda de precisão. Isso significa que os modelos grandes já guardam uma reserva de aceleração; falta apenas aprender a usá-la.

Onde a eficiência se perde

O problema é que o hardware popular foi concebido originalmente para computação densa, não para estruturas esparsas. Quando os dados são comprimidos, é preciso guardar metadados junto com os valores não nulos: índices de linhas, colunas e segmentos. O acesso a esses dados se torna indireto e imprevisível: o processador precisa primeiro buscar as coordenadas e só depois o valor. Com isso, parte do tempo vai não para a matemática, mas para vagar pela memória e executar operações auxiliares.

• GPUs lidam muito bem com matrizes densas, mas, com esparsidade aleatória, muitas vezes paralelizam operações inúteis com zeros.
• A esparsidade estruturada nem sempre ajuda, porque exige um padrão rígido de zeragem, por exemplo dois zeros em cada quatro parâmetros vizinhos.
• CPUs são mais flexíveis, mas frequentemente esbarram em falhas do prefetcher e em acessos imprevisíveis à memória.
• Mesmo as bibliotecas sparse não eliminam toda a sobrecarga, porque parte dos recursos vai para a manutenção dos próprios dados.

Os fabricantes já buscam atalhos, mas por enquanto são pontuais. A Apple acelerou os acessos indiretos à memória nos chips A14 e M1, a Cerebras promove a abordagem sparse no seu Wafer Scale Engine e a Meta desenvolve o MTIA. Mas aqui também há limites: algumas soluções funcionam só com esparsidade de pesos, outras dão suporte apenas a operações isoladas, como a multiplicação de matrizes. Para cargas reais de IA, isso é pouco, porque os modelos não consistem em uma única operação, e sim em uma longa cadeia de camadas e transformações diferentes.

Como o Onyx funciona

A equipe de Stanford partiu do nível mais baixo e criou o Onyx, um acelerador programável que trabalha igualmente bem com computação esparsa e densa. Sua base é uma arquitetura CGRA, intermediária entre CPU e FPGA: bem mais flexível que um processador clássico, mas mais eficiente que circuitos totalmente configuráveis em nível de bit.

O Onyx é formado por blocos de computação e blocos de memória que armazenam matrizes comprimidas e as processam diretamente nesse formato, sem expandi-las de volta ao formato denso sem necessidade. O compilador também é essencial: ele traduz expressões como a multiplicação de matriz esparsa por vetor em um grafo de memória e cálculo e depois o distribui pelos blocos do chip.

Segundo Stanford, o Onyx consumiu em média 70 vezes menos energia que uma CPU e executou os cálculos cerca de 8 vezes mais rápido. Na métrica energy-delay product, o ganho chegou a 565 vezes em relação a um Intel Xeon de 12 núcleos com bibliotecas sparse. A próxima geração do Onyx deve incluir suporte a camadas não lineares, normalização, softmax e uma troca mais prática entre os modos sparse e dense.

O que isso significa

A ideia central do artigo não é que surgiu mais um chip de IA, mas que os desenvolvedores começam a otimizar os modelos não só com menor precisão ou menor tamanho, e sim também pela estrutura dos próprios cálculos. Se a abordagem da esparsidade se firmar, os modelos grandes poderão rodar de forma mais barata e rápida, o que significa que o próximo salto em IA pode vir não só de novos modelos, mas de uma nova classe de hardware.

*A Meta é considerada uma organização extremista e está proibida na Federação Russa.