Correr IA num CPU: quantização, ONNX e porque é que os modelos pequenos ganham

Porque é que a velocidade importa

Uma API de moderação de chats tem de responder enquanto o utilizador ainda está a olhar para o ecrã: algumas centenas de milissegundos, no máximo. As GPU são rápidas mas caras; os servidores comuns têm CPU. Pôr um modelo rápido num CPU é o que torna um serviço de moderação acessível.

O que calcula um modelo

Uma rede neuronal é sobretudo multiplicação de matrizes: os números que representam o texto são multiplicados pelos pesos do modelo, camada após camada. A velocidade depende de quantas multiplicações há (o tamanho do modelo e o comprimento do texto) e da rapidez de cada uma.

A quantização

Os pesos são normalmente guardados como números decimais de 32 bits. A quantização guarda-os como inteiros de 8 bits (int8), com um fator de escala por linha ou canal:

Formato	Tamanho de um modelo com 560 milhões de parâmetros	Velocidade típica em CPU
float32	cerca de 2,2 GB	referência
int8	cerca de 0,6 GB	muitas vezes 2 a 4 vezes mais rápido

Os CPU modernos têm instruções que multiplicam muitos números de 8 bits ao mesmo tempo (VNNI, AVX-512, Intel AMX), e é daí que vem a velocidade.

O preço é uma perda mínima de precisão. Com a quantização dinâmica por canal, um bom classificador não costuma perder quase nada mensurável, mas as pontuações deslocam-se ligeiramente. Por isso os limiares têm de ser realinhados depois de quantizar: cada limiar é movido para que o modelo quantizado assinale a mesma proporção de mensagens que o original num conjunto de calibração.

ONNX e ONNX Runtime

O ONNX é um formato aberto para guardar um modelo como um grafo de operações, independente do framework com que foi treinado. O ONNX Runtime executa esse grafo e está otimizado para CPU: funde operações, usa todos os núcleos e escolhe as instruções mais rápidas disponíveis.

Um exemplo real

A VigiMod 1.2.1 tem cerca de 560 milhões de parâmetros. Com 4 núcleos de CPU:

Versão	Tempo por par mensagem–categoria
Precisão total	cerca de 113 ms
ONNX int8	cerca de 32 ms (3,5 vezes mais rápido)

Depois de realinhar os limiares, a perda de qualidade medida nos dados de validação ficou abaixo de 0,005 de F1, dentro do ruído.

Outros truques

• Entradas mais curtas: limita o contexto às últimas mensagens que importam.
• Lotes: pontua várias mensagens de uma vez quando chegam juntas.
• Guardar em cache a parte fixa da entrada (por exemplo uma política longa) nos modelos que o permitem.
• Modelo mais pequeno, dados melhores: um modelo pequeno treinado com os exemplos certos muitas vezes ganha a um grande e genérico, por uma fração do custo.

Quando é mesmo preciso uma GPU

Para treinar e para correr modelos muito grandes que geram respostas longas. Um desenho habitual é uma cascata: um modelo pequeno e rápido em todas as mensagens e um maior só nos poucos casos duvidosos.