llama2.rs: 纯Rust实现的高性能Llama 2推理引擎

llama2.rs: 纯Rust实现的高性能Llama 2推理引擎

llama2.rs是一个由Rust语言编写的Llama 2推理引擎,旨在提供最快速的CPU推理性能。该项目由GitHub用户srush和rachtsingh开发,是对Andrej Karpathy的llama2.c项目的Rust重新实现和扩展。

主要特性

llama2.rs具有以下几个显著特点:

1. 支持4位GPTQ量化,大幅降低模型内存占用
2. 批量预填充prompt tokens,提高推理效率
3. 利用SIMD指令集加速CPU推理
4. 内存映射技术,实现70B参数模型的瞬时加载
5. 静态大小检查,提高内存安全性
6. 支持分组查询注意力(GQA),适用于大规模Llama模型
7. 提供Python调用API

在作者的Intel i9台式机上,llama2.rs可以实现70B参数Llama 2模型1 token/s的推理速度,7B参数模型则可达到9 token/s,相比原始的llama2.c有显著提升。

使用方法

要使用llama2.rs,首先需要安装Rust的nightly工具链:

rustup toolchain install nightly
ulimit -s 10000000 # 增加栈内存限制

然后可以从Hugging Face hub下载预训练模型。例如,以下命令将创建一个70B参数、4位量化、64组大小的模型:

pip install -r requirements.export.txt
python export.py l70b.act64.bin TheBloke/llama-2-70b-Guanaco-QLoRA-GPTQ gptq-4bit-64g-actorder_True

使用cargo编译并运行模型:

cargo run --release --features 70B,group_64,quantized -- -c llama2-70b-q.bin -t 0.0 -s 11 -p "The only thing"

这将输出生成的文本和推理速度。

Python接口

llama2.rs也提供了Python接口。首先需要编译并安装Python包:

cargo build --release --features 7B,group_128,quantized,python
pip install .

然后可以在Python中使用:

import llama2_rs

model = llama2_rs.LlamaModel("lorca13b.act132.bin", False)
tokenizer = llama2_rs.Tokenizer("tokenizer.bin")
random = llama2_rs.Random()
response = llama2_rs.generate(
    model,
    tokenizer,
    "Tell me zero-cost abstractions in Rust ",
    50,
    random, 
    0.0
)

未来计划

llama2.rs的开发团队计划在未来实现以下功能:

1. 使用Triton支持快速GPU处理
2. 支持text-generation-webui项目
3. 完善文档
4. 发布关于GPTQ快速实现方法的博客文章
5. 移除对AutoGPTQ的预加载依赖
6. 直接支持safetensors格式

技术细节

llama2.rs主要依赖以下几个Rust crate:

• memmap2: 用于内存映射
• rayon: 用于并行计算
• clap: 用于命令行参数解析
• pyO3: 用于Python调用
• portable_simd: 用于SIMD加速

为了获得最佳性能,llama2.rs需要针对不同的Llama版本进行重新编译。模型参数在.cargo/config中配置,如果与加载的二进制模型不匹配将会失败。

结语

llama2.rs展示了Rust语言在AI推理引擎开发中的潜力。通过利用Rust的零成本抽象、内存安全和并发特性,llama2.rs在保证安全性的同时实现了卓越的性能。这个项目为在普通CPU上运行大规模语言模型提供了一个高效的解决方案,为AI技术的普及应用铺平了道路。

随着项目的不断发展,我们可以期待llama2.rs在未来带来更多优化和功能,进一步推动AI技术在各种计算环境下的应用。对于有兴趣深入了解或贡献该项目的开发者,可以访问llama2.rs的GitHub仓库获取更多信息。