Fine-tuning LLMs to 1.58-bit: Extreme Quantization with BitNet

Hugging Face has developed a method to fine-tune existing Large Language Models (LLMs) to 1.58-bit precision using the BitNet architecture. This approach allows models to represent parameters using only three values (-1, 0, 1), drastically reducing computational and energy costs without requiring the massive budget typically needed to pre-train a 1-bit model from scratch.

BitNet Architecture and 1.58-bit Quantization

BitNet replaces standard Linear layers in Multi-Head Attention and Feed-Forward Networks with BitLinear layers. These layers utilize ternary precision for weights and 8-bit precision for activations.

The Computation Paradigm

Unlike standard LLMs (e.g., Llama) that rely on FP16 addition and multiplication, BitNet b1.58 uses INT8 addition for matrix multiplication. This shift in computation theoretically reduces the energy consumption for matrix multiplication by 71.4 times compared to a Llama baseline.

Training with Straight-Through Estimators (STE)

Because the round() function used for ternary quantization is non-differentiable, BitNet employs a Straight Through Estimator (STE). The STE approximates the gradient of the rounding operation as 1, allowing gradients to flow through the operation as if it were an identity function, which enables standard gradient-based optimization.

Quantization Mechanics

  • Weights: Quantized using symmetric per-tensor quantization. The scale is the inverse of the mean absolute value of the weight matrix. Weights are scaled, rounded, clamped between -1 and 1, and then rescaled.
  • Activations: Quantized to 8-bit precision using absmax per-token quantization, scaling values into the range [-128, 127]. Layer Normalization (LN) is applied before activation quantization to maintain output variance.

Fine-tuning Existing Models to 1.58-bit

Hugging Face successfully fine-tuned a Llama 3 8B model to 1.58-bit precision. Initial experiments revealed that abruptly introducing BitLinear layers caused the model to lose nearly all its pre-trained information, resulting in a loss spike.

Dynamic Warmup Quantization

To prevent the loss of prior knowledge, Hugging Face implemented a dynamic $\lambda$ (lambda) value to introduce quantization gradually:

$$\lambda = \min\left(\frac{\text{training_step}}{\text{total_training_steps}}, 1\right)$$

By scaling the difference between the original and quantized values by $\lambda$, the model transitions from full precision ($\lambda=0$) to full quantization ($\lambda=1$). This linear scheduler led to better convergence and a perplexity of approximately 4 on the TinyStories dataset.

Scaling and Generalization

To ensure the model retained general knowledge and didn't overfit to small datasets, the team scaled training to the FineWeb-edu dataset. Using a learning rate of 1e-4 and a batch size of 2 million tokens over 10 billion tokens, the model achieved a WikiText perplexity of 12.2.

Further scaling to 100 billion tokens showed that while the model performed closely to the original Llama 3 8B in some metrics, it generally lagged slightly behind the full-precision baseline.

Performance Benchmarks and Results

Models fine-tuned with the 1.58-bit architecture were released under the HF1BitLLM organization.

Key Findings

  • Competitive Performance: After fine-tuning on 10 billion tokens, the 1.58-bit Llama 3 8B model outperformed the BitNet 7B model (trained on 100B tokens) and the FBI LLM (distilled on 1.26 trillion tokens).
  • MMLU Benchmarks: The developed 8B models surpassed the Llama 1 7B model in MMLU benchmarks.
  • Model Size: Packing weights into an int8 tensor reduced the parameter count from 8B to 2.8B.

Inference Optimization and Custom Kernels

To realize the speed and memory benefits of 1.58-bit weights, Hugging Face implemented custom CUDA and Triton kernels to handle on-the-fly weight unpacking during matrix multiplication.

Tiled Matrix Multiplication

To overcome memory bandwidth bottlenecks and redundant data access, the team used tiling. This technique divides matrices into smaller sub-matrices (tiles) that fit into the GPU's fast shared memory, reducing the frequency of slow global memory accesses.

Kernel Benchmarking

  • Triton vs. Torch: The custom Triton kernel achieved performance approximately equal to @torch.compile with BF16 precision.
  • BitBlas: The team found that BitBlas, a mixed-precision software library, outperformed both the custom Triton kernel and Torch's matmul function in low precision, although it introduced higher initial loading times due to kernel compilation.

Integration with Transformers

Integration is handled via a new "bitnet" quantization method in the transformers library. Standard Linear layers are replaced with specialized BitLinear layers. The API remains unchanged, allowing users to load the models using AutoModelForCausalLM.from_pretrained.

Sources