Implantação eficiente de grandes modelos de linguagem de baixa precisão em CPUs
A principal inovação do T-MAC está na adoção de um paradigma de computação baseado em tabelas de consulta (LUT), em vez do tradicional paradigma de multiplicação e acumulação (MAC). O T-MAC utiliza tabelas de consulta para suportar diretamente cálculos de baixa precisão, eliminando assim as operações de dequantização necessárias em outros sistemas e reduzindo significativamente o número de operações de multiplicação e adição.
Experimentos mostraram que o T-MAC apresenta um desempenho excepcional: em um PC Surface AI equipado com o mais recente chipset Qualcomm Snapdragon X E lite, o modelo BitNet-b1.58 de 3B alcança uma taxa de geração de 48 tokens por segundo, o modelo llama de 2 bits e 7B atinge 30 tokens por segundo, e o modelo llama de 4 bits e 7B chega a 20 tokens por segundo.
Isso supera até mesmo o desempenho da NPU!
Ao implantar o modelo llama-2-7b-4bit, embora a NPU possa gerar 10,4 tokens por segundo, a CPU com o T-MAC pode atingir 12,6 tokens por segundo usando apenas dois núcleos, chegando a até 22 tokens por segundo no máximo.
Tudo isso supera em muito a velocidade média de leitura humana e representa um aumento de 4 a 5 vezes em relação ao framework llama.cpp original.
Mesmo em dispositivos de menor potência, como o Raspberry Pi 5, o T-MAC pode atingir uma taxa de geração de 11 tokens por segundo para o BitNet-b1.58 de 3B. O T-MAC também tem vantagens significativas em termos de consumo de energia: para atingir a mesma taxa de geração, o T-MAC requer apenas 1/4 a 1/6 do número de núcleos do llama.cpp original, reduzindo o consumo de energia e deixando recursos computacionais para outras aplicações.
É importante notar que o desempenho computacional do T-MAC aumenta linearmente à medida que o número de bits diminui, um fenômeno difícil de observar em GPUs e NPUs baseadas em dequantização. No entanto, o T-MAC pode alcançar 10 tokens por segundo em um único núcleo e 28 tokens por segundo em quatro núcleos com 2 bits, superando significativamente o desempenho da NPU.
Multiplicação de matrizes sem multiplicação, apenas consulta de tabela (LUT)
Para parâmetros de baixa precisão (pesos), o T-MAC agrupa cada bit separadamente (por exemplo, um grupo de 4 bits), multiplica esses bits pelo vetor de ativação, pré-calcula todas as somas parciais possíveis e as armazena usando LUT.
Em seguida, o T-MAC utiliza operações de deslocamento e acumulação para suportar números de bits escaláveis de 1 a 4. Com esse método, o T-MAC abandona as instruções FMA (multiplicação-adição) ineficientes na CPU em favor das instruções TBL/PSHUF (consulta de tabela) mais eficientes e de menor consumo de energia.
Computação centrada em bits, substituindo a computação centrada em tipos de dados
O cálculo tradicional baseado em dequantização é, na verdade, centrado em tipos de dados, o que requer personalização separada para cada tipo de dado diferente.
Cada combinação de largura de bits de ativação e peso, como W4A16 (peso int4, ativação float16) e W2A8, requer um layout de peso específico e um kernel de computação.
Por exemplo, o layout W3 requer que 2 bits e 1 bit sejam empacotados separadamente e utilizem diferentes métodos de entrelaçamento ou embaralhamento para alinhamento de memória ou decodificação rápida. Em seguida, o kernel de computação correspondente precisa desempacotar esse layout específico para os tipos de dados suportados pelo hardware para execução.
O T-MAC, por outro lado, observa o cálculo de multiplicação de matriz de baixa precisão do ponto de vista dos bits, projetando uma estrutura de dados ideal para um único bit e depois expandindo para 2/3/4 bits através de empilhamento.
Ao mesmo tempo, para vetores de ativação de diferentes precisões (float16/float32/int8), apenas o processo de construção da tabela precisa ser alterado, sem a necessidade de considerar diferentes estruturas de dados durante a consulta da tabela.
Além disso, nos métodos tradicionais baseados em dequantização, ao reduzir de 4 bits para 3/2/1 bits, embora o uso de memória seja menor, a quantidade de cálculos não diminui e o desempenho pode até piorar devido ao aumento da sobrecarga de dequantização.
No entanto, a quantidade de cálculos do T-MAC diminui linearmente com a redução do número de bits, proporcionando melhor aceleração em bits mais baixos e oferecendo uma solução de implantação eficiente para os mais recentes modelos de 1 bit/2 bits, como BitNet e EfficientQAT.
Implementação altamente otimizada de operadores
O cálculo centrado em bits tem muitas vantagens, mas implementá-lo em CPUs ainda apresenta desafios significativos:
(1) Comparado ao acesso contínuo de dados para ativações e pesos, o acesso à tabela é aleatório. A residência da tabela na memória on-chip rápida é crucial para o desempenho final de inferência;
(2) No entanto, a memória on-chip é limitada, e o método de tabela de consulta (LUT) aumenta o uso da memória on-chip em comparação com o mpGEMV tradicional. Isso ocorre porque a tabela de consulta precisa armazenar os resultados da multiplicação do vetor de ativação com todos os padrões de bits possíveis, o que é muito mais do que a própria ativação.
Para abordar isso, os pesquisadores do Microsoft Research Asia investigaram profundamente o fluxo de dados de computação baseado em tabelas e projetaram estruturas de dados e fluxos de computação eficientes para esse paradigma de computação, incluindo:
-
Armazenar LUTs na memória on-chip para aproveitar as instruções vetoriais de consulta de tabela (TBL/PSHUF) da CPU e melhorar o desempenho de acesso aleatório à memória.
-
Alterar a ordem de cálculo do eixo da matriz para maximizar a taxa de reutilização de dados das LUTs limitadas que podem ser colocadas na memória on-chip.
-
Projetar um método de tiling de matriz otimizado especificamente para consulta de tabela, combinado com a busca de parâmetros de tiling ideais usando autotvm.
-
Otimização do layout dos parâmetros (pesos):
a) Reordenação dos pesos para maximizar o acesso contínuo e melhorar a taxa de acertos de cache.
b) Entrelaçamento dos pesos para melhorar a eficiência de decodificação.
- Otimizações específicas para CPUs Intel/ARM, incluindo:
a) Reordenação de registradores para construção rápida de tabelas de consulta.
b) Uso de instruções de média para acumulação rápida de 8 bits.
Os pesquisadores aplicaram várias otimizações passo a passo em uma implementação básica, obtendo uma aceleração significativa em relação aos operadores de baixa precisão do estado da arte: