Add Support for BitNet Architecture Inference

[JoseCarlosGarcia95]

Introduction

Hello! My name is José Carlos, and I hold a degree in Mathematics. Alongside other talented individuals, I co-founded a company where we focus on AI and infrastructure solutions.

I’ve always believed that BitNet is one of the most significant breakthroughs of the past year. I have a personal obsession with this architecture, as it embodies the potential to balance performance and efficiency in language models—a pursuit that deeply motivates me.

---

Changes Made

1. Added Support for BitNet Architecture Inference:

   • Implemented initial support to infer the BitNet architecture within Candle.
   • Note: This implementation does not yet include quantization support for BitNet.

2. New Example Added:

   • Introduced a new example to test and demonstrate the newly added BitNet inference functionality.

3. Supported HF Models:
   The following models were tested successfully:

   • "1bitLLM/bitnet_b1_58-large": BitNet B1 58 Large.
   • "1bitLLM/bitnet_b1_58-3B": BitNet B1 58 3B.

   Future support will be added for:

   • "HF1BitLLM/Llama3-8B-1.58-100B-tokens": Llama 3 (8B, 1.58).

---

Known Limitations

• Current implementation does not support quantization for BitNet.
• Matrix multiplication methods in this implementation are not optimized yet.

I plan to address these in future updates and prepare a PR with:

1. Support for quantizations tailored to BitNet.
2. Optimized methods for matrix multiplication.

---

Roadmap

• Add support for Llama 3 (8B, 1.58).
• Explore and implement methods to enhance matrix multiplication performance.

---

Thank you for considering this PR! Feedback is welcome, and I’m excited to continue contributing to this amazing project. 😊

[LaurentMazare]

The bit-linear tests that you added seem to be broken, would you mind having a look?

[JoseCarlosGarcia95]

The bit-linear tests that you added seem to be broken, would you mind having a look?

Done! @LaurentMazare

[LaurentMazare]

Thanks, could you also provide details on how the model results were lined up with the python implementation? Did you ensure that the logits generated by the candle version are somewhat in line?

[JoseCarlosGarcia95]

@LaurentMazare Everything stems from: The Era of 1-bit LLMs - Training Tips, Code, FAQ

Since the inference for Llama and Linear was already implemented in this project, I used the existing Llama implementation as a foundation and applied the following changes, based on the paper and as seen here:

• Replace Linear with BitLinear in MLP and Attention.
• Add weights and layers for RMSNorm before calling BitLinear, as shown here:
  • Link 1
  • Link 2
• Implement activation_quant and weight_quant equivalently using Candle.
  • Link 3

In principle, unless I’m mistaken, everything seems consistent with the Python implementation. You can test it in situ as follows:

cargo run --example llama-bitnet --features metal

Thank you so much for reviewing!

[LaurentMazare]

In order to check the consistency, the best would be to generate the logits for the same prompt on the candle and python side and check that they are reasonably close. That's what we do for most models before adding them, would you mind giving it a try?

[JoseCarlosGarcia95]

@LaurentMazare

Since this is my first contribution, I don't have much experience with this, but it seems that the logits from the Python implementation and my implementation are similar:

I used the code from this link in Python, and in the code I created, I added a print of the logits (in Candle) and a print of the output.score in Python.

Here’s what I observed:

Candle:
[-3.4375, -10.9765625, 1.171875, -0.4970703, -0.88427734, 2.0449219, 3.1738281, -0.15246582, 0.5522461, 1.4648438, 2.34375, 0.8886719, 4.7539063]

Python:
[-3.458984375, -11.28125, 1.1884765625, -0.52587890625, -0.81103515625, 2.021484375, 3.17578125, -0.173828125, 0.568359375, 1.5283203125, 2.455078125, 0.96044921875, 4.84375]

They seem equivalent, except for numerical error.