BitNet

BitNet reemplaza las capas lineales tradicionales en Multi-Head Attention y en las redes feed-forward con capas BitLinear especializadas. Las capas BitLinear cuantizan los pesos usando precisión ternaria (con valores de -1, 0 y 1) y cuantizan las activaciones a una precisión de 8 bits.

Los modelos de BitNet no pueden cuantizarse sobre la marcha. Deben cuantizarse durante el preentrenamiento o el ajuste fino, ya que se trata de una técnica de entrenamiento consciente de la cuantización (QAT). Durante el entrenamiento, los pesos se cuantizan a valores ternarios mediante cuantización simétrica por tensor.

1. Calcule el promedio de los valores absolutos de la matriz de pesos y úselo como escala.
2. Divida los pesos por la escala, redondee los valores, restríngalos entre -1 y 1 y vuelva a escalarlos para continuar en precisión completa.
3. Las activaciones se cuantizan a un ancho de bits especificado (8 bits) mediante cuantización absmax (cuantización simétrica por canal). Esto implica escalar las activaciones a un rango de [−128,127].

Consulte este PR para preentrenar o afinar un modelo de 1,58 bits con Nanotron. Para el afinado, convierta un modelo del formato de Hugging Face al formato de Nanotron. Encuentre los pasos de conversión en este PR.

Cargue un modelo cuantizado de BitNet con from_pretrained().

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

Kernels

@torch.compile se utiliza para descomprimir los pesos y ejecutar la pasada hacia adelante. Su implementación es muy sencilla y ofrece mejoras significativas en la velocidad. Se integrarán núcleos optimizados adicionales en versiones futuras.

Recursos

Lee Fine-tuning LLMs to 1.58bit: extreme quantization made easy para saber más sobre cómo se entrenan y ajustan los modelos de BitNet.