$ show kyegomez/BitNet
BitNet is a PyTorch library implementing the BitNet architecture for 1-bit transformers, as described in the paper "BitNet: Scaling 1-bit Transformers for Large Language Models". The library provides drop-in replacements for standard linear layers with BitLinear layers that use quantized weights and activations to enable efficient computation. It includes components such as BitMGQA (Multi-Grouped Query Attention), BitFeedForward, and full transformer models (BitNetTransformer), as well as adapters for other architectures like Mamba and Mixture of Experts. The library also offers utilities to convert existing Hugging Face or PyTorch models to use BitLinear layers and includes an optimized CUDA kernel for low-precision GEMM operations to accelerate inference. The architecture is modular, with the core BitLinear layer encapsulated in `bitnet/bitlinear.py`, which can be used independently. Higher-level components like attention and feed-forward networks are built using BitLinear and are designed to be reusable in other projects. The library supports training and inference workflows, as demonstrated by the provided example scripts and training routine for the enwik8 dataset. With active maintenance and a focus on performance through CUDA optimization, BitNet serves as a specialized tool for researchers and practitioners interested in extremely low-precision deep learning models.
The LLM identified 5 files worth highlighting for reuse:
Implements the BitLinear layer which quantizes weights and activations to 1-bit using the straight-through estimator. It is self-contained (depends only on PyTorch and the zeta library for RMSNorm) and can be dropped into any neural network as a replacement for nn.Linear. Demonstrates efficient low-precision computation.
Constructs a feed-forward network using BitLinear layers, with optional GLU, Swish, and post-activation layer normalization. It is a reusable building block for transformer-like architectures, demonstrating how to compose BitLinear layers into higher-order components.
Implements Multi-Grouped Query Attention (MGQA) using BitLinear for projections. It reduces the number of key-value heads for efficiency while maintaining performance. The module is self-contained and can be integrated into transformer architectures, showcasing an optimized attention mechanism.
Provides functions to recursively replace all nn.Linear layers in a PyTorch or Hugging Face model with BitLinear layers. This utility enables easy conversion of existing models to the BitNet architecture, demonstrating a non-intrusive model adaptation pattern.
Implements a CUDA kernel for low-precision GEMM operations used to accelerate BitLinear layers. The kernel performs matrix multiplication in half precision with scaling factors, demonstrating high-performance GPU programming for efficient inference.