BitNet-Transformers

BitNet-Transformers 项目介绍

项目概述

BitNet-Transformers 是一个依赖 Huggingface Transformers 的创新项目，旨在实现“一位变换器”以用于大型语言模型。这个项目利用 PyTorch 和 Llama(2) 架构，为 Transformer 模型带来了全新的优化方法。

项目背景

在深度学习领域中，大型语言模型（LLM）的出现带来了处理自然语言任务时的显著进步。然而，这些模型通常需要大量的计算资源和存储空间。因此，优化这些资源的使用，尤其是在模型权重的存储和处理上，是一个值得研究的方向。BitNet 通过引入“1 位”权重方法，极大地减少了模型的内存消耗。

环境准备

为了在本地环境中启动和运行 BitNet-Transformers 项目，用户首先需要克隆项目的 GitHub 仓库并安装相关依赖项。使用以下命令可完成开发环境的配置：

# 克隆项目仓库
git clone https://github.com/beomi/bitnet-transformers
cd bitnet-transformers

# 安装必要依赖
pip install -r clm_requirements.txt

# 克隆 Huggingface 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 等工具来追踪模型指标的变化。

GPU 内存使用比较

BitNet 变种模型在不同的位数精度下，显著优化了 GPU 内存的使用：

• 原始 LLAMA - 16 位：模型权重需要约 250MB 的 GPU 内存。
• BitLLAMA - 混合 16 位：模型权重占用约 200MB，采用混合位数策略。
• BitLLAMA - 8 位：进一步减少到 100MB，通过即需使用的方式灵活切换位数。
• BitLLAMA - 1 位：使用 1 位权重进一步降低内存需求。

这些优化措施使得在相同硬件条件下，可以训练更大的模型或运行更多的任务。

项目进展与计划

目前，项目已经完成了以下功能的实现：

• 添加 BitLinear 层
• 集成 LLamaForCausalLM 模型
• 支持 1 位权重的保存

未来的工作计划包括优化 BitLinear 层的权重存储方法以及利用自定义 CUDA 核心提升性能。

BitNet-Transformers 项目通过创新的权重存储和计算方法，为大规模模型的轻量级部署和训练提供了新的思路。该项目具有开源性质，为研究人员和开发者提供了一个探索和实现新型模型优化技术的平台。