Tokenizador do BERT na Hugging Face fica 50 vezes mais rápido ao eliminar trabalho desnecessário com ASCII

O tokenizador transforma uma string nos números que a rede neural aceita. Na biblioteca tokenizers, da Hugging Face, essa etapa estava lenta sem necessidade nos modelos BERT. A equipe fez o profiling do código, encontrou trabalho supérfluo e o removeu. O tokenizador do BERT agora é cerca de 50 vezes mais rápido que na versão 0.23. O número vem de duas etapas: o ganho de 8,6 vezes obtido na preparação do lançamento da v1.0.0 e mais 5,7 vezes de um novo ajuste. Esse ajuste ainda está em um pull request. Em paralelo, os tokenizadores BPE clássicos, nos quais roda a maioria dos LLMs, ficaram até 30 vezes mais rápidos.

Por que o velho BERT ainda importa

O BERT foi lançado em 2018, mas seus descendentes continuam por toda parte: embeddings, busca, RAG, rerankers, classificadores. Os três modelos mais baixados do Hugging Face Hub usam o normalizador do BERT:

• all-MiniLM-L6-v2 — 218 milhões de downloads por mês; • o cross-encoder ms-marco-MiniLM-L6-v2 — 79 milhões; • bge-small-en-v1.5 — 59 milhões.

O próprio bert-base-uncased aparece em sexto lugar, com 34 milhões. Entre os 10 000 modelos mais populares, o BPE soma 842 milhões de downloads por mês, e o WordPiece (algoritmo do BERT), 627 milhões. Mais da metade dos downloads de BPE vai para tokenizadores com normalização Unicode, na maioria das vezes NFC, como em Qwen ou ModernBERT. As melhorias aceleram esse caso também.

Onde o tempo se perdia

A normalização reduz diferentes formas de escrever o mesmo caractere a uma forma única. Por exemplo, "é" pode ser escrito como um único ponto de código ou como a letra "e" com um sinal de acento separado. Sem normalização, o modelo gera tokens diferentes para os dois. O BERT sem distinção de maiúsculas e minúsculas tem cinco passos: limpeza de caracteres de controle e de espaço, espaços ao redor de ideogramas chineses, decomposição NFD, remoção de diacríticos e conversão para minúsculas.

O flame graph (gráfico em que a largura do bloco mostra a fatia de tempo da função) deu um quadro inequívoco. O normalizador consumia 88% do tempo, enquanto o próprio WordPiece, que escolhe os tokens, ficava com apenas 8%. Cerca de dois terços de todo o tempo iam para a função binary_search_by, que verificava se o caractere era um sinal combinante. Trata-se de uma busca binária em um array de 1567 caracteres (6 KB), com cerca de 11 comparações e ramificações imprevisíveis para o processador. A busca rodava para cada caractere, inclusive para os simples "s", "l", "o" e "w" latinos.

Três truques preguiçosos

O ponto de partida foi estudar os dados. Textos reais têm muito ASCII, os caracteres de uma mesma escrita aparecem juntos e a maior parte do texto já vem normalizada. Disso saíram três técnicas.

• Pular ASCII com SIMD. Em UTF-8, um byte é ASCII se o bit mais significativo for zero. O processador carrega 16 bytes, reúne os bits mais significativos em um número com uma única instrução e compara esse número com zero.

Se for zero, o bloco inteiro dispensa decodificação, consulta a tabelas e até cópia. Se não for, a contagem de bits zero localiza o primeiro byte não latino. A conversão de letras ASCII maiúsculas para minúsculas também é feita em lote: operação OR com 0x20.

• Cache por blocos de pontos de código. O cirílico ocupa 256 pontos de código consecutivos (U+0400–U+04FF). Na primeira vez que encontra um caractere do bloco, o algoritmo já calcula a resposta para "a normalização pode afetá-lo?"

nas 256 posições. Depois disso, basta uma leitura de memória. O bloco ocupa quatro linhas de cache de 64 bytes e permanece no cache mais rápido, junto ao núcleo.

• Cópia preguiçosa. Se o texto não mudou, a string original é devolvida sem alocar memória. Na primeira alteração, a memória é alocada uma única vez, e os trechos intocados são copiados em blocos grandes.

O que deu certo e onde não ajudou

Em texto em inglês, a velocidade subiu de 23 para 151 MB/s em relação ao branch main, e a fatia do normalizador caiu de 88% para 13%. Os tokens coincidem com os anteriores em todos os corpora de teste. Quanto mais ASCII no texto, maior o ganho: inglês e código-fonte aceleraram mais. Em russo, grego e árabe, o ASCII representa cerca de 11%, mas o cache de blocos trouxe ganho de 2,2 a 3,2 vezes. No ModernBERT com NFC, o inglês quase não mudou (1,06 vez), porque o ASCII já era pulado antes, enquanto código, árabe e russo ficaram cerca de duas vezes mais rápidos.

A abordagem tem pontos fracos:

• O texto em chinês não ficou mais rápido: há pouco ASCII nele, e o BERT de qualquer forma insere espaços ao redor de cada ideograma. • Strings de até 32 bytes podem ser processadas mais devagar, porque o custo adicional supera o ganho. No pior caso (StripAccents em 8 bytes de chinês), a velocidade cai para 0,35 da anterior. • Considerando só o normalizador, houve lentidão acima da margem de erro em 55 de 320 casos, 19 deles com entradas de até 32 bytes. • O BERT com distinção de maiúsculas e minúsculas processa chinês de 8% a 15% mais devagar em entradas a partir de 64 bytes.

Para engenheiros, a conclusão é simples. Se você processa milhões de documentos com MiniLM ou bge para busca e RAG, a tokenização na CPU pode ser um item relevante de custo, e depois do lançamento da v1.0.0 ela ficará muito mais barata. Quem tem muitas strings curtas ou em chinês no fluxo deve medir a velocidade nos próprios dados.