bitnet_b1_58-3B_bitblas

This is a BitBLAS Implementation for the reproduced 1.58bit model from 1bitLLM/bitnet_b1_58-3B. We replaced the original simulated Int8x3bit Quantized Inference Kernel with BitBLAS INT8xINT2 Kernel. We also evaluated the model's correctness and performance through eval_correctness.py and benchmark_inference_latency.py.

Latest News

• 08/09/2024 ✨: We provide a more efficient implementation for bitnet with vLLM, which should use special model checkpoints, to make the ckpt and study how to deploy, please checkout Make Checkpoints for vLLM.

Make Checkpoints for vLLM

We provide two scripts to make the checkpoints for vLLM. The first script is generate_bitnet_model_native_format.sh, which is used to make a checkpoint with fp16 uncompressed metaadta, the main difference with the original checkpoint is the quant_config.json, which allow vLLM to load the model and execute with a quant extension.

# move to the integration directory
cd /root/to/BitBLAS/integration/BitNet
# make the checkpoint
./maint/generate_bitnet_model_native_format.sh
# the output ckpy will be saved in the `./models/ckpt_bitnet_b1_58-3B` directory

The second script is generate_bitnet_model_bitblas_format.sh, which is used to make a checkpoint with BitBLAS compressed metadata, which can avoid the online dequantize sage for the profiling of vLLM, which lead to more efficient memory utilization.

./maint/generate_bitnet_model_bitblas_format.sh ./models/ckpt_bitnet_b1_58-3B ./models/ckpt_bitnet_b1_58-3B_bitblas
# the output ckpy will be saved in the `./models/ckpt_bitnet_b1_58-3B_bitblas` directory

Finnaly, you can use the ckpt in vLLM with:

cd vllm_workspace
# inference with the ckpt with fp16 uncompressed metadata
python3 inference_with_native_format.py
# inference with the ckpt with BitBLAS compressed metadata
python3 inference_with_bitblas_format.py

BitBLAS Results

Performance

Note: To reproduce the results of BitBLAS, Please checkout the benchmark_inference_latency.py. To reproduce the results of the original model, Please checkout the 1bitLLM/bitnet_b1_58-3B repo.

On-the-Fly GPU Memory Footprint

We measured the GPU memory footprint through the nvidia-smi command. Please checkout nvidia_measure_memory.sh to get the real-time GPU memory usage. And then start a benchmark_model_10k_loops.py workload to measure the overall GPU memory usage.

PPL and Zero-shot Accuracy

The number is Reported from the 1bitLLM/bitnet_b1_58-3B, Please checkout the eval_ppl.py.

PPL and zero-shot accuracy:

The differences between the reported numbers and the reproduced results are possibly variances from the training data processing, seeds, or other random factors.

Citations

@article{ma2024era,
  title={The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits},
  author={Ma, Shuming and Wang, Hongyu and Ma, Lingxiao and Wang, Lei and Wang, Wenhui and Huang, Shaohan and Dong, Li and Wang, Ruiping and Xue, Jilong and Wei, Furu},
  journal={arXiv preprint arXiv:2402.17764},
  year={2024}
}