gpt-fast

简单高效的原生 PyTorch 文本生成 transformer。

特点:

1. 极低延迟
2. 少于1000行 Python 代码
3. 除 PyTorch 和 sentencepiece 外无其他依赖
4. int8/int4 量化
5. 推测性解码
6. 张量并行
7. 支持 Nvidia 和 AMD GPU

这并非旨在成为一个"框架"或"库" - 而是为了展示使用原生 PyTorch 可以获得怎样的性能 :) 请随意复制粘贴和分支使用。

关于此代码库内容的深入解析，请查看这篇博客文章。

支持的模型

LLaMA 系列

请查看本页面其余部分了解 LLaMA 系列模型的基准测试。

Mixtral 8x7B

我们还支持 Mixtral 8x7B，这是一个高质量的稀疏混合专家模型(MoE)，平均令牌生成速率为:

注意，基准测试在8个 A100-80GB GPU 上运行，功率限制为 330W，采用混合立方网格拓扑。所有基准测试均在批次大小=1的情况下运行，使得报告的令牌/秒数等同于"令牌/秒/用户"。此外，它们使用非常小的提示长度（仅5个令牌）运行。

有关 Mixtral 8x7B 的更多详情，请查看此页面或这个说明。

示例

为保持仓库的简洁性，以下是您可以作为PR对 gpt-fast 进行的各种扩展示例。

• Google Gemma
• xAI Grok-1
• Databricks DBRX

社区

社区中受 gpt-fast 启发的项目:

• gpt-blazing: 将相同的性能优化策略应用于更多模型（如 baichuan2）。
• gptfast: 将部分性能优化应用于所有 Huggingface 模型。
• gpt-accelera: 扩展 gpt-fast 以支持 SFT/RM/PPO 训练和批处理推理，以优化吞吐量。

安装

下载 PyTorch nightly 版本

安装所需的包:

pip install -r requirements.txt

要下载 llama 模型，请访问 https://huggingface.co/meta-llama/Llama-2-7b 并完成获取访问权限的步骤。 然后使用 huggingface-cli login 登录。

下载权重

已测试/支持的模型

tinyllamas/stories{15,42,100}
openlm-research/open_llama_7b
meta-llama/Llama-2-7b-chat-hf
meta-llama/Llama-2-13b-chat-hf
meta-llama/Llama-2-70b-chat-hf
codellama/CodeLlama-7b-Python-hf
codellama/CodeLlama-34b-Python-hf
mistralai/Mistral-7B-v0.1
mistralai/Mistral-7B-Instruct-v0.1
mistralai/Mistral-7B-Instruct-v0.2
meta-llama/Meta-Llama-3-8B

例如，要转换 Llama-2-7b-chat-hf

export MODEL_REPO=meta-llama/Llama-2-7b-chat-hf
./scripts/prepare.sh $MODEL_REPO

基准测试

基准测试在8个 A100-80GB GPU 上运行，功率限制为 330W，采用混合立方网格拓扑。注意，所有基准测试均在批次大小=1的情况下运行，使得报告的令牌/秒数等同于"令牌/秒/用户"。此外，它们使用非常小的提示长度（仅5个令牌）运行。

推测性采样

验证器: Llama-70B (int4), 草稿: Llama-7B (int4): 48.4 令牌/秒

张量并行

张量并行 + 量化

AMD

在 MI-250x 的一个 GCD 上运行的基准测试。

生成文本

模型定义在 model.py 中，生成代码在 generate.py 中。

python generate.py --compile --checkpoint_path checkpoints/$MODEL_REPO/model.pth --prompt "Hello, my name is"

为了稍微提高一些性能，你还可以使用 --compile_prefill 编译预填充。不过这会增加编译时间。

量化

通过以下方式选择要使用的设备

# 当前支持的设备：cuda, cpu
export DEVICE=cuda

Int8 仅权重量化

要生成此版本的模型

# 在 checkpoints/$MODEL_REPO/model_int8.pth 输出模型
python quantize.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --mode int8

要使用 int8 运行，只需将 int8 检查点传递给 generate.py。

python generate.py --compile --checkpoint_path checkpoints/$MODEL_REPO/model_int8.pth --device $DEVICE

Int4 仅权重量化

要生成 int4 版本的模型

# 在 checkpoints/$MODEL_REPO/model_int4.g32.$DEVICE.pth 输出模型
python quantize.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --mode int4 --groupsize 32

要使用 int4 运行，只需将 int4 检查点传递给 generate.py。

python generate.py --checkpoint_path checkpoints/$MODEL_REPO/model_int4.g32.pth --compile

投机采样

要使用投机采样生成（DRAFT_MODEL_REPO 应指向比 MODEL_REPO 更小的模型）。

在这个例子中，"更小"的模型只是模型的 int8 量化版本。

export DRAFT_MODEL_REPO=meta-llama/Llama-2-7b-chat-hf
python generate.py --compile --checkpoint_path checkpoints/$MODEL_REPO/model.pth --draft_checkpoint_path checkpoints/$DRAFT_MODEL_REPO/model_int8.pth

注意：在 A100 80GB 上运行，功率限制为 330 瓦。根据经验，峰值带宽约为 1700 GB/s。

张量并行

ENABLE_INTRA_NODE_COMM=1 torchrun --standalone --nproc_per_node=2 generate.py --compile --checkpoint_path checkpoints/$MODEL_REPO/model.pth

实验性功能

评估

我们使用 EleutherAI 评估工具来评估模型的准确性。要评估准确性，请确保已安装评估工具，并将模型检查点和所需任务传递给 eval.py。

python eval.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --compile --tasks hellaswag winogrande

注意：gpt-fast 目前不支持生成式任务

评估工具的安装说明：https://github.com/EleutherAI/lm-evaluation-harness/tree/master#install

GPTQ

我们有一个纯 PyTorch 实现的 GPTQ，它利用 torch._dynamo.export 访问模型结构。你可以通过使用相同的命令生成 GPTQ 量化版本的 int4 量化，但在量化模式中添加 'gptq'，即：

# 在 checkpoints/$MODEL_REPO/model_int4-gptq.g32.pth 输出模型
python quantize.py --mode int4-gptq --calibration_tasks wikitext --calibration_seq_length 2048

然后你可以像上面一样用这个模型进行评估或生成文本。

许可证

gpt-fast 根据 BSD 3 许可证发布。

致谢

感谢：

• Lightning AI 对 PyTorch 的支持，以及在闪光注意力、int8 量化和 LoRA 微调方面的工作。
• GGML 推动了 LLM 的快速本地设备推理。
• Karpathy 带头开发简单、可解释和快速的 LLM 实现。
• MLC-LLM 推动了异构硬件上 4 位量化的性能。