Extropic revela chip Z1 para inferência de transformers mais eficiente

O anúncio

A Extropic apresentou o Z1T, sua primeira família de modelos parecidos com Transformers projetada para chips probabilísticos esparsos como o Z1. A empresa reporta ganhos de eficiência energética superiores a 100x em relação a GPUs na inferência, acompanhados de uma lei de escalonamento empírica própria. Também abriu o código das receitas de treinamento e publicou os pesos de um dos maiores treinamentos do Z1T.

Motivação

A Extropic argumenta que a inferência de Transformers domina as cargas de trabalho dos aceleradores modernos, com demanda de mercado que gerou um dos maiores investimentos industriais da história recente e forte aumento no consumo de energia. Embora a tese de longo prazo da empresa seja criar novos modelos nativos de hardware probabilístico, o trabalho atual atende à necessidade urgente de reduzir a energia da inferência de Transformers.

Como funciona o hardware

O Z1 é um chip CMOS sub-limiar que modela bits probabilísticos (pbits), amostrando de um modelo de Ising programável via amostragem de Gibbs a 50 MHz. A conectividade física é fixa no silício: cada nó tem grau 16, com 269.568 pbits e 2.135.904 arestas de acoplamento por chip. Essa esparsidade é a chave da eficiência energética e um limite físico do hardware - diferente das GPUs, otimizadas para multiplicações de matrizes densas, nas quais operações esparsas não geram aceleração proporcional (uma multiplicação entre matrizes com 99% de esparsidade em GPU não é 100x mais rápida).

Arquitetura e codificação

Para rodar redes neurais no chip, os autores se inspiram na linhagem das máquinas de Boltzmann: em vez de amostrar diretamente, extraem médias estatísticas do computador termodinâmico para produzir vetores clássicos. Números contínuos são quantizados via codificação dy4p (4 pbits com escalas potências de dois); ao aumentar o número de amostras em runtime, é possível obter precisão efetiva maior que os 4 bits físicos, incluindo bits fracionários.

O Z1T adapta os primitivos do Transformer: RMSNorm é substituído por Dynamic Tanh (DyT), e a atenção softmax densa dá lugar à atenção convolucional com portas (GCA), com projeções esparsas de 4 entradas. O MLP é compilado como unidades tanh-lineares esparsas no tecido do chip.

Inferência desagregada

A decodificação é dividida entre substratos heterogêneos: operações esparsas e tanh-lineares rodam no Z1, enquanto FPGAs coprocessadores na mesma placa executam embeddings, residuais, pooling e o readout de vocabulário. Os autores ressaltam que o companheiro poderia ser GPUs ou outros aceleradores (XPUs), algo que explorarão em trabalhos futuros.

Escalonamento e desempenho

Os experimentos de escalonamento, treinados em OpenWebText com o tokenizador BPE do GPT-2, mostram que o Z1T precisa de cerca de uma ordem de magnitude mais FLOPs que um GPT-2 para atingir a mesma perda - mas as operações no Z1 são cerca de três ordens de magnitude mais eficientes energeticamente, resultando em ganho líquido de duas ordens de magnitude. Experimentos com conectividade variável (de 5% a 99,8% de esparsidade) indicam que FLOPs densos são mais eficientes por parâmetro, mas o custo energético por FLOP difere entre substratos.

Nas estimativas de energia, o sistema Z1T consome 294,52 nJ por token (8,74 nJ no Z1 e 285,78 nJ no FPGA), contra 4,09 µJ a 40,9 µJ do H100 dependendo da utilização - entre cerca de 14x e 139x de ganho, chegando a aproximadamente 4.680x nas camadas que rodam só no Z1. Em latência, o Z1T estima 58,8 µs por token (~17.000 tokens/s) contra 102 µs do H100 com torch.compile. Os autores advertem que as comparações com H100 usam um modelo pequeno com MFU de 0,006%, que batch no GPU tornaria a H100 substancialmente mais eficiente, e que o pré-fill ainda é melhor feito em GPU. As estimativas excluem a leitura densa final de logits e assumem contagem de chips não modelada.

Perspectivas

O FPGA consome mais de 95% da energia do sistema, o que indica espaço para melhorias: ao mover mais operações para um substrato CMOS sub-limiar projetado para modelos como o Z1T, a empresa sugere que se poderia alcançar até 1000x mais eficiência energética que GPUs. O trabalho é descrito como o início da co-design entre modelos parecidos com Transformers e hardware probabilístico.