gpt-fast

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

特点：

1. 极低延迟
2. 代码量少于 1000 行 Python
3. 仅依赖 PyTorch 和 sentencepiece
4. int8/int4 量化
5. 推测解码
6. 张量并行
7. 支持 Nvidia 和 AMD 显卡

这不是一个“框架”或“库”——它的目的是展示使用原生 PyTorch 可以达到怎样的性能 :) 欢迎随意复制、粘贴和 fork。

如需深入了解此代码库的内容，请参阅这篇博客文章。

支持的模型

LLaMA 系列

请查看本页关于 LLaMA 系列模型的基准测试。

Mixtral 8x7B

我们还支持了高质量的稀疏专家混合（MoE）模型 Mixtral 8x7B，其平均 token 生成速率如下：

	1 GPU	2 GPU	4 GPU	8 GPU
baseline(bfloat16)	OOM	96.67	155.35	227.82
int8	97.92	155.03	216.87	279.35

请注意，这些基准测试是在一台配备 8 块 A100-80GB 显卡、功率限制为 330W 且采用混合立方体网格拓扑的机器上运行的。所有基准测试均以 batch size=1 运行，因此报告的 tokens/s 数值等同于“tokens/s/user”。此外，它们还使用非常短的提示长度（仅 5 个 token）。

有关 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
meta-llama/Meta-Llama-3.1-8B
meta-llama/Meta-Llama-3.1-70B
meta-llama/Meta-Llama-3.1-405B

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

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

基准测试

基准测试在一台配备 8 块 A100-80GB 显卡、功率限制为 330W 且采用混合立方体网格拓扑的机器上进行。请注意，所有基准测试均以 batch size=1 运行，因此报告的 tokens/s 数值等同于“tokens/s/user”。此外，它们还使用非常短的提示长度（仅 5 个 token）。

模型	技术	Tokens/Second	内存带宽 (GB/s)
Llama-2-7B	Base	104.9	1397.31
	8-bit	155.58	1069.20
	4-bit (G=32)	196.80	862.69
Llama-2-70B	Base	OOM
	8-bit	19.13	1322.58
	4-bit (G=32)	25.25	1097.66
Llama-3.1-8B	Base	93.89	1410.76
	8-bit	137.64	1030.89
Llama-3.1-70B	Base	OOM
	8-bit	18.04	1253.78

推测采样

验证器：Llama-70B (int4)，草稿：Llama-7B (int4)：48.4 tok/s

张量并行

模型	GPU 数量	Tokens/Second	内存带宽 (GB/s)
Llama-2-7B	1	104.9	1397.31
	2	168.84	1181.99
	4	254.02	955.83
	8	328.43	704.10
Llama-2-70B	1	OOM
	2	21.32	1481.87
	4	38.01	1340.76
	8	62.50	1135.29
Llama-3.1-8B	1	93.83	1408.37
	2	149.10	1197.32
	4	217.21	986.32
	8	276.01	772.60
Llama-3.1-70B	1	OOM
	2	16.03	1130.81
	4	37.45	1360.53
	8	58.78	1129.61

张量并行 + 量化

模型	技术	Tokens/Second	内存带宽 (GB/s)
Llama-2-70B	Base	62.50	1135.29
	8-bit	80.44	752.04
	4-bit (G=32)	90.77	548.10
Llama-3.1-70B	Base	58.78	1129.61
	8-bit	75.58	726.57
Llama-3.1-405B	8-bit	15.60	815.87

AMD

基准测试在一块 MI-250x 的 GCD 上进行。

模型	技术	Tokens/Second	内存带宽 (GB/s)
Llama-2-7B	Base	76.33	1028.70
	8-bit	101.86	700.06

生成文本

模型定义位于 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 checkpoint 传递给 generate.py 即可。

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

Int4 仅权重量化

要生成 int4 版本的模型：

# Spits out model at checkpoints/$MODEL_REPO/model_int4.pth
python quantize.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --mode int4

要使用 int4 模型，只需将 int4 checkpoint 传递给 generate.py 即可。

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

Int4 仅权重量化

要生成 int4 版本的模型：

# Spits out model at checkpoints/$MODEL_REPO/model_int4.pth
python quantize.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --mode int4

要使用 int4 模型，只需将 int4 checkpoint 传递给 generate.py 即可。

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

在检查点目录下输出模型：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 即可。
```bash
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 的支持，以及在 Flash Attention、int8 量化和 LoRA 微调方面的工作。
• GGML 推动了 LLM 的快速设备端推理技术。
• Karpathy 领导了简单、可解释且快速的 LLM 实现。
• MLC-LLM 在异构硬件上提升了 4 位量化性能。

gpt-fast 快速上手指南

gpt-fast 是一个基于原生 PyTorch 构建的极简、高效 Transformer 文本生成项目。它无需复杂框架，仅依赖 PyTorch 和 sentencepiece，通过极少的代码量（<1000 行）实现了低延迟推理、int8/int4 量化、投机采样（Speculative Decoding）及张量并行等高级特性。

环境准备

系统要求

• 操作系统: Linux (推荐)
• GPU: 支持 NVIDIA (CUDA) 或 AMD (ROCm) GPU
• Python: 3.8+

前置依赖

本项目核心依赖为 PyTorch Nightly 版本以获得最佳性能支持，以及 sentencepiece。

国内加速建议： 安装 PyTorch 时，建议使用清华大学或阿里云镜像源加速下载。

安装步骤

1. 安装 PyTorch Nightly

访问 PyTorch 官网 获取最新命令，或使用以下通用命令（以 CUDA 12.1 为例，国内用户可替换 pip 源）：

pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
# 国内镜像加速示例 (清华源可能不包含 nightly，建议官方源或寻找特定 nightly 镜像)
# pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121

注：AMD 用户请安装对应的 ROCm 版本 PyTorch。

2. 克隆项目并安装依赖

git clone https://github.com/meta-pytorch/gpt-fast.git
cd gpt-fast
pip install -r requirements.txt

3. 配置 Hugging Face 访问

下载 LLaMA 等模型需要 Hugging Face 权限。

1. 前往 Hugging Face Llama-2 页面 申请访问权限。
2. 在终端登录：

huggingface-cli login

提示：如果网络连接不稳定，建议在终端设置代理或使用国内镜像站下载模型权重后手动放入目录。

基本使用

1. 下载并准备模型权重

以 Llama-2-7b-chat-hf 为例，运行脚本将模型转换为 gpt-fast 格式：

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

支持模型包括：Llama-2/3 系列、CodeLlama、Mistral、Mixtral 等（详见 README Supported Models 列表）。

2. 生成文本

使用默认精度进行文本生成：

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

参数说明：

• --compile: 启用 torch.compile 进行加速（首次运行会有编译耗时）。
• --checkpoint_path: 指向转换后的模型权重文件。
• --prompt: 输入提示词。

3. (可选) 使用量化模型加速

为了进一步降低显存占用并提升速度，推荐使用 int8 或 int4 量化。

生成 int8 模型并运行：

# 1. 量化权重
python quantize.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --mode int8

# 2. 运行生成 (指定 int8 权重)
python generate.py --compile --checkpoint_path checkpoints/$MODEL_REPO/model_int8.pth --device cuda

生成 int4 模型并运行：

# 1. 量化权重 (groupsize=32)
python quantize.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --mode int4 --groupsize 32

# 2. 运行生成
python generate.py --checkpoint_path checkpoints/$MODEL_REPO/model_int4.g32.pth --compile

4. (可选) 多卡张量并行

对于大模型（如 70B），单卡显存不足时可利用多卡并行：

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

将 --nproc_per_node 修改为你拥有的 GPU 数量。