kyegomez/BitNet

Implementation of "BitNet: Scaling 1-bit Transformers for Large Language Models" in pytorch

What it solves

BitNet is a PyTorch implementation of 1-bit Transformers, designed to significantly reduce the memory and computational requirements of Large Language Models (LLMs). It replaces standard linear projections with binarized layers to enable ultra-compression of model weights.

How it works

The project implements the BitLinear module, which replaces the standard nn.Linear in a Transformer. The process involves a sequence of layer normalization, binarization, and abs-max quantization followed by dequantization. The repository provides tools to swap out linear layers in existing PyTorch or Hugging Face models using replace_linears_in_pytorch_model and replace_linears_in_hf.

Who it’s for

It is intended for AI researchers and developers who want to implement 1-bit LLMs or explore ultra-compression for text, image, or audio processing. Note that models using these layers must be trained or fine-tuned from scratch to function correctly.

Highlights

  • Broad Architecture Support: Includes implementations of BitNetTransformer, BitAttention (using Multi-Grouped Query Attention), BitFeedForward, BitMoE (Mixture of Experts), and BitMamba.
  • Integration Tools: Provides utility functions to easily replace linear layers in standard PyTorch and Hugging Face models.
  • Vision Support: Includes OneBitViT for applying 1-bit compression to Vision Transformers.
  • Performance Optimizations: Features an optimized CUDA kernel for low-bit GEMM operations.

Related

  • Project
    bitsandbytes-foundation/bitsandbytes

    A PyTorch library that enables accessible large language models through 4-bit and 8-bit quantization for both inference and training.

  • Project
    Lightning-AI/pytorch-lightning

    A deep learning framework that abstracts away the engineering boilerplate of PyTorch, enabling researchers to pretrain and finetune AI models at scale across various hardware.

  • Project
    jingyaogong/minimind

    MiniMind is a complete, open-source codebase for training a very small (~64M parameter) large language model entirely from scratch on a tiny budget (about 3 RMB and 2 hours). It covers the full modern LLM lifecycle—pretraining, SFT, LoRA, DPO, RLAIF (PPO/GRPO/CISPO), tool use, and agentic RL—and includes experimental extensions for MoE, vision, and diffusion models.

  • Project
    pytorch-tabular/pytorch_tabular

    A high-level PyTorch framework that provides a unified interface for applying various deep learning architectures to tabular data.

  • Project
    InterDigitalInc/CompressAI

    A PyTorch library and evaluation platform for end-to-end deep learning-based image and video compression research.