BitNet-Transformers:缩放1比特Transformer实现大型语言模型
BitNet-Transformers是一个开源项目,旨在实现论文《BitNet: Scaling 1-bit Transformers for Large Language Models》中提出的1比特Transformer架构。该项目基于Hugging Face的Transformers库,使用PyTorch实现,并采用了Llama(2)的模型架构。
项目特点
- 实现了1比特权重的BitLinear层,可作为nn.Linear的替代方案
- 大幅降低模型内存占用和能耗,同时保持竞争性能
- 展现出与全精度Transformer类似的扩展性
- 提供了BitLLAMA模型,支持混合精度、8位和1位等多种量化方案
快速开始
- 克隆项目仓库:
git clone https://github.com/beomi/bitnet-transformers
cd bitnet-transformers
- 安装依赖:
pip install -r clm_requirements.txt
- 克隆并安装修改后的Transformers库:
git clone https://github.com/huggingface/transformers
pip install -e transformers
- 更新Llama(2)模型:
rm ./transformers/src/transformers/models/llama/modeling_llama.py
ln -s $(pwd)/bitnet_llama/modeling_llama.py ./transformers/src/transformers/models/llama/modeling_llama.py
模型训练
项目提供了在Wikitext-103数据集上训练BitLLAMA模型的脚本:
./train_wikitext.sh
您可以通过Weights & Biases追踪训练指标。
内存使用对比
在相同配置下(批次大小1,序列长度2048),不同精度的BitLLAMA模型内存占用如下:
- 原始LLAMA (16位): 250MB
- BitLLAMA (混合16位): 200MB
- BitLLAMA (8位): 100MB
- BitLLAMA (1位): TBD
学习资源
未来工作
- 完善BitLinear层的1比特权重实现
- 优化LLamaForCausalLM模型以支持BitLinear层
- 添加语言模型训练的示例代码
- 实现自定义CUDA kernel以提高1比特权重的计算效率
BitNet-Transformers为大型语言模型的高效训练和部署提供了新的可能性。我们鼓励感兴趣的开发者深入探索该项目,为其发展做出贡献。
如需更多信息或讨论,欢迎访问项目GitHub页面并参与讨论。
