BitNet

BitNet substitui as camadas lineares tradicionais em Multi-Head Attention e redes feed-forward por camadas BitLinear especializadas. As camadas BitLinear quantizam os pesos usando precisão ternária (com valores de -1, 0 e 1) e quantizam as ativações para precisão de 8 bits.

Os modelos BitNet não podem ser quantizados em tempo de execução. Eles precisam ser quantizados durante o pré-treinamento ou o ajuste fino, pois se trata de uma técnica de Treinamento Consciente de Quantização (QAT). Durante o treinamento, os pesos são quantizados para valores ternários com quantização simétrica por tensor.

1. Calcule a média dos valores absolutos da matriz de pesos e use-a como escala.
2. Divida os pesos pela escala, arredonde os valores, restrinja-os entre -1 e 1 e reescale-os para continuar em precisão total.
3. As ativações são quantizadas para uma largura de bits especificada (8 bits) usando quantização absmax (quantização simétrica por canal). Isso envolve escalar as ativações para um intervalo de [−128,127].

Consulte este PR para pré-treinar ou ajustar finamente um modelo de 1,58 bits com Nanotron. Para o ajuste fino, converta um modelo do Hugging Face para o formato Nanotron. Encontre as etapas de conversão neste PR.

Carregue um modelo quantizado BitNet com from_pretrained().

from transformers import AutoModelForCausalLM
path = "/path/to/model"
model = AutoModelForCausalLM.from_pretrained(path, device_map="auto")

Kernels

@torch.compile é usado para descompactar os pesos e executar a passagem direta. A implementação é muito direta e proporciona melhorias significativas de velocidade. Kernels otimizados adicionais serão integrados em versões futuras.

Recursos

Leia Fine-tuning LLMs to 1.58bit: extreme quantization made easy para saber mais sobre como os modelos BitNet são treinados e ajustados.