bitnet_cpu

KO | EN

gitlite — search

#python #open-source #react #javascript #web-development #embedded-systems #java #desktop-application #gui #mcp #terminal #typescript

★ 56 Open GitHub ↗

Experiments with BitNet inference on CPU

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# BitNet CPU

This repo provides some tests to see how well the BitNet inference idea works on a CPU.  This is based on the open-source reproduction here: https://huggingface.co/1bitLLM/bitnet_b1_58-3B/ which is reproducing the results from this series of Microsoft papers:

* https://github.com/microsoft/unilm/blob/master/bitnet/The-Era-of-1-bit-LLMs__Training_Tips_Code_FAQ.pdf
* "The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits" https://arxiv.org/abs/2402.17764
* "BitNet: Scaling 1-bit Transformers for Large Language Models" https://arxiv.org/abs/2310.11453
* Maybe more here? https://github.com/microsoft/unilm/tree/master/bitnet

## Setup

```bash
git clone https://github.com/catid/bitnet_cpu.git
cd bitnet_cpu
mkdir build
cd build
cmake ..
```

### Linux build:

```bash
make -j
```

### Windows Visual Studio 2022 build:

Follow the setup instructions using Git Bash: https://git-scm.com/

Open the .sln file in the ./build/ folder.

Select Release mode and build!

## CPU Approach 1: AVX-512

The best SIMD intrinsics for this do not exist on most Intel computers, but you can use AVX512 _mm256_mask_add_epi16() and _mm256_mask_sub_epi16() on Intel servers to speed this up a lot.  With those intrinsics, the weights can be 2 bits per parameter, and the model might be competitive for memory/compute efficiency.  I optimized the reference inference kernel using AVX-512 on a rented Intel Xeon W-2295.  This is using OMP for multi-thread optimization, and some careful unrolling for pipeline parallelism.

```bash
catid@project512:~/sources/bitnet_cpu/build$ ./tests/math_test 
Running non-AVX-512 version
Test case 1 passed!
Test case 2 passed!
Test case 3 passed!
Test case 4 passed!
Running AVX-512 version
Tests passed!

catid@project512:~/sources/bitnet_cpu/build$ ./benchmark_model 
Preallocating buffers...
Benchmarking model...
Warmup took 131 milliseconds
Benchmark Results:
Number of Layers: 182
Number of Benchmark Iterations: 1
Average Time per Iteration: 25 milliseconds
```

This would be about 1000/25 = ~40 tokens per second on CPUs with AVX-512.

Here's on my AMD Ryzen 9 7950X Windows PC:

```bash
mrcat@kuang3 MINGW64 ~/sources/bitnet_cpu/build (master)
$ ./tests/Release/math_test.exe
Running non-AVX-512 version
Test case 1 passed!
Test case 2 passed!
Test case 3 passed!
Test case 4 passed!
Running AVX-512 version
Tests passed!

mrcat@kuang3 MINGW64 ~/sources/bitnet_cpu/build (master)
$ ./Release/benchmark_model.exe
Preallocating buffers...
Benchmarking model...
Warmup took 21 milliseconds
Benchmark Results:
Number of Layers: 182
Number of Benchmark Iterations: 100
Average Time per Iteration: 18.86 milliseconds
```

This is ~50 tokens/second with AVX-512.

Here's on the same machine with AVX2 instead:

```bash
mrcat@kuang3 MINGW64 ~/sources/bitnet_cpu/build (master)
$ ./Release/benchmark_model.exe
Preallocating buffers...
Benchmarking model...
Warmup took 63 milliseconds
Benchmark Results:
Number of Layers: 182
Number of Benchmark Iterations: 100
Average Time per Iteration: 63.91 milliseconds
```

So we have 15 tokens/second without AVX-512.

Compare this to AMD Ryzen 5 7535HS CPU achieves about 7.4 tokens/second for Gemma 3B, and it's clear that the BitNet inference kernel is competitive with 8-bit inference on CPU in some cases.

## CPU Approach 2: AVX2

Using the _mm256_sign_epi8() intrinsic, we can use 1 byte per parameter (8-bit quant) to implement an efficient BitNet kernel that is portable to most Intel CPUs using AVX2.  This is using the same unrolling/multi-threading optimizations as the AVX-512 version.

Running on an 12th Gen Intel(R) Core(TM) i9-12900K with 20 threads:

```
(base) ➜  build git:(master) ✗ ./tests/math_test
Running non-AVX-512 version
Test case 1 passed!
Test case 2 passed!
Test case 3 passed!
Test case 4 passed!
Tests passed!

(base) ➜  build git:(master) ✗ ./benchmark_model
Preallocating buffers...
Benchmarking model...
Warmup took 35 milliseconds
Benchmark Results:
Number of Layers: 182
Number of Benchmark Iterations: 100
Average Time per Iteration: 34.85 milliseconds
```

We are getting about 1000/35 = ~28 tokens per second on most Intel CPUs with AVX2.

## CPU Approach 3: LOL

I also tried converting the model weights to C++ code, which was kind of funny.  I put the results in the `lol_bitnet_to_cpp` folder.  This experiment obviously failed: The C++ files are 200 MB and the compiler takes more than 10 hours to compile each one.

## Thanks

@lithium0003 on Github suggested using the _mm256_sign_epi8() intrinsic to greatly speed up the AVX2 version.

bitnet_cpu

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README