BitNet-Transformers 入门学习指南 - 1比特Transformer缩放大型语言模型

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位等多种量化方案

快速开始

1. 克隆项目仓库:

 git clone https://github.com/beomi/bitnet-transformers
 cd bitnet-transformers

2. 安装依赖:

 pip install -r clm_requirements.txt

3. 克隆并安装修改后的Transformers库:

 git clone https://github.com/huggingface/transformers
 pip install -e transformers

4. 更新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

学习资源

1. 项目GitHub仓库
2. BitNet论文
3. Hugging Face Transformers文档
4. PyTorch文档

未来工作

• 完善BitLinear层的1比特权重实现
• 优化LLamaForCausalLM模型以支持BitLinear层
• 添加语言模型训练的示例代码
• 实现自定义CUDA kernel以提高1比特权重的计算效率

BitNet-Transformers为大型语言模型的高效训练和部署提供了新的可能性。我们鼓励感兴趣的开发者深入探索该项目,为其发展做出贡献。

如需更多信息或讨论,欢迎访问项目GitHub页面并参与讨论。