AQLM

AQLM的官方PyTorch实现，基于论文《通过加性量化实现大型语言模型的极限压缩》

[2024年5月] AQLM被ICML'2024接收！如果您参加会议，欢迎来这个海报处与我们见面。

[2024年6月] 现在有一种更有效的方法来调优量化模型，称为PV-tuning。我们正在发布经PV-tuning的AQLM模型在这个集合中，代码在pv-tuning分支中。我们将在进行几项技术改进后将pv-tuning代码合并到主分支。

推理

演示

通过以下Google Colab示例学习如何运行预量化模型：

基本AQLM <br> 生成	使用GPU/CPU <br> 进行流式处理	使用CUDA图 <br> 进行推理（3倍加速）	使用PEFT <br> 进行微调	使用vLLM <br> 进行服务
<a target="_blank" href="https://colab.research.google.com/github/Vahe1994/AQLM/blob/main/notebooks/colab_example.ipynb"><img src="https://yellow-cdn.veclightyear.com/835a84d5/3640a2ae-0fbc-42da-a963-8346f4905fc5.svg" alt="AQLM在Colab中"/></a>	<a target="_blank" href="https://colab.research.google.com/github/Vahe1994/AQLM/blob/main/notebooks/streaming_example.ipynb"><img src="https://yellow-cdn.veclightyear.com/835a84d5/3640a2ae-0fbc-42da-a963-8346f4905fc5.svg" alt="AQLM在Colab中"/></a>	<a target="_blank" href="https://colab.research.google.com/github/Vahe1994/AQLM/blob/main/notebooks/aqlm_cuda_graph.ipynb"><img src="https://yellow-cdn.veclightyear.com/835a84d5/3640a2ae-0fbc-42da-a963-8346f4905fc5.svg" alt="在Colab中打开"/></a>	<a target="_blank" href="https://colab.research.google.com/github/Vahe1994/AQLM/blob/main/notebooks/aqlm_2bit_training.ipynb"><img src="https://yellow-cdn.veclightyear.com/835a84d5/3640a2ae-0fbc-42da-a963-8346f4905fc5.svg" alt="在Colab中打开"/></a>	<a target="_blank" href="https://colab.research.google.com/github/Vahe1994/AQLM/blob/main/notebooks/aqlm_vllm.ipynb"><img src="https://yellow-cdn.veclightyear.com/835a84d5/3640a2ae-0fbc-42da-a963-8346f4905fc5.svg" alt="在Colab中打开"/></a>

模型

该存储库目前设计用于LLaMA、Mistral和Mixtral系列模型。 以下报告的模型使用了附录A中描述的全模型微调方法，采用带有教师logits的交叉熵目标函数。

我们提供了一些未经PV-Tuning的预量化AQLM模型（下滑可查看经PV-Tuning的模型）：

模型	AQLM方案	WikiText-2 困惑度	MMLU (5-shot) FP16→AQLM	模型大小, Gb	Hub链接
Llama-3-8b	1x16	-	0.65→0.56	4.1	链接
Llama-3-8b-Instruct	1x16	-	0.66→0.59	4.1	链接
Llama-3-70b	1x16	-	0.79→0.75	21.9	链接
Llama-3-70b-Instruct	1x16	-	0.80→0.76	21.9	链接
Command-R	1x16	-	0.68→0.57	12.7	链接
Command-R+	1x16	-	0.74→0.68	31.9	链接
Mistral-7b	1x16	5.40	-	2.5	链接
Mistral-7B-Instruct-v0.2	2x8	-	0.59→0.44	2.5	链接
Mixtral-8x7b	1x16	3.35	-	12.6	链接
Mixtral-8x7b-Instruct	1x16	-	-	12.6	链接
Llama-2-7b	1x16	5.92	0.46→0.39	2.4	链接
Llama-2-7b	2x8	6.69	-	2.2	链接
Llama-2-7b	8x8	6.61	-	2.2	链接
Llama-2-13b	1x16	5.22	0.55→0.49	4.1	链接
Llama-2-13b	2x8	5.63	-	3.8	链接
Llama-2-70b	1x16	3.83	0.69→0.65	18.8	链接
Llama-2-70b	2x8	4.21	-	18.2	链接
gemma-2b	1x16	-	-	1.7	链接
gemma-2b	2x8	-	-	1.6	链接

您还可以下载通过PV-tuning微调的AQLM模型：

模型	AQLM方案	WikiText-2困惑度	模型大小, Gb	Hub链接
Llama-2-7b	1x16g8	5.68	2.4	链接
Llama-2-7b	2x8g8	5.90	2.2	链接
Llama-2-7b	1x16g16	9.21	1.7	链接
Llama-2-13b	1x16g8	5.05	4.1	链接
Llama-2-70b	1x16g8	3.78	18.8	链接
Meta-Llama-3-8B	1x16g8	6.99	4.1	链接
Meta-Llama-3-8B	1x16g16	9.43	3.9	链接
Meta-Llama-3-70B	1x16g8	4.57	21.9	链接
Meta-Llama-3-70B	1x16g16	8.67	13	链接
Mistral-7B-v0.1	1x16g8	5.22	2.51	链接
Phi-3-mini-4k-instruct	1x16g8	6.63	1.4	链接

请注意，方案中带有"g16"的模型需要aqlm推理库v1.1.6或更新版本：

pip install aqlm[gpu,cpu]>=1.1.6

上述困惑度是在Llama 2模型的4k上下文长度和Mistral/Mixtral及Llama 3的8k上下文长度上评估的。 还请注意，令牌级困惑度只能在同一模型系列内进行比较，不应该在使用不同词汇表的模型之间进行比较。 虽然Mistral的困惑度低于Llama 3 8B，但这并不意味着Mistral更好：Llama的困惑度是在更大的字典上计算的，因此每个令牌的困惑度较高。

关于更多评估结果和详细解释，请参阅我们的论文：Egiazarian等人（2024）介绍纯AQLM，Malinovskii等人（2024）介绍PV-Tuned模型。

推理内核

AQLM量化设置主要在使用的码本数量和码本大小（以位为单位）上有所不同。最常用的设置及其支持的推理内核如下：

内核	码本数量	码本大小（位）	方案表示	准确度	加速比	快速GPU推理	快速CPU推理
Triton	K	N	KxN	-	最高 ~0.7x	✅	❌
CUDA	1	16	1x16	最佳	最高 ~1.3x	✅	❌
CUDA	2	8	2x8	可以	最高 ~3.0x	✅	❌
Numba	K	8	Kx8	良好	最高 ~4.0x	❌	✅

安装

要运行这些模型，需要安装推理库：

pip install aqlm[gpu,cpu]

根据推理设置指定gpu、cpu或两者。

然后，可以使用transformers库提供的熟悉的.from_pretrained方法：

from transformers import AutoModelForCausalLM

quantized_model = AutoModelForCausalLM.from_pretrained(
    "ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf",
    trust_remote_code=True, torch_dtype="auto"
).cuda()

注意，在GPU上torch_dtype应设置为torch.float16或"auto"，在CPU上应设置为torch.float32。之后，可以像使用非量化模型一样使用该模型。

量化

依赖项

从requirements.txt安装包：

pip install -r requirements.txt

加载/缓存数据集和分词器

脚本需要下载并在本地缓存相关的分词器和数据集。 除非通过环境变量提供替代位置，否则它们将保存在默认的Huggingface Datasets目录中。 参见相关Datasets文档部分

数据

使用AQLM量化模型时，我们建议使用模型原始训练数据的子集。

对于Llama-2模型，最接近的可用数据集是RedPajama。要加载RedPajama的子集，在--dataset参数中提供"pajama"。 这将处理nsamples数据并使用提供的模型分词器对其进行分词。

此外，我们为LLama和Solar/Mistral模型提供了4096上下文长度的分词RedPajama，存储在Huggingface中。 要加载它，使用：

from huggingface_hub import hf_hub_download

hf_hub_download(repo_id="Vahe1994/AQLM", filename="data/name.pth", repo_type="dataset")

要使用从HF下载的数据，将其放在data文件夹中（可选），并在main.py中的"--dataset"参数中设置正确的路径。 警告： 这些子集已经用相应的模型分词器处理过了。如果你想量化另一个模型(例如 mistral/mixtral)，请使用 src/datautils 中提供的脚本重新对数据进行分词。

WandB 日志记录

可以选择将数据记录到 Weights and Biases 服务(wandb)。 运行 pip install wandb 以启用 W&B 日志记录。 在运行实验之前，指定 $WANDB_ENTITY、$WANDB_PROJECT、$WANDB_NAME 环境变量。使用 --wandb 参数启用日志记录。

GPU 和内存要求

此代码是使用多个 80GB GPU 内存的 A100 GPU 开发和测试的。 你可以使用 --offload activations 选项来减少 VRAM 使用量。 对于 Language Model Evaluation Harness 评估，需要有足够的内存在一个或多个设备上加载整个模型 + 激活张量。

量化时间

AQLM 量化在校准时间上比 GPTQ 等更简单的量化方法要长得多。这只影响量化时间，不影响推理时间。

例如，使用默认配置量化 7B 模型在单个 A100 GPU 上大约需要 1 天时间。同样，在单个 GPU 上量化 70B 模型需要 10-14 天。如果你有多个具有快速互连的 GPU，可以运行 AQLM 多 GPU 以加速比较 - 只需为多个 GPU 设置 CUDA_VISIBLE_DEVICES。在两个 GPU 上量化 7B 模型将量化时间缩短到约 14.5 小时。同样，在 8 个 A100 GPU 上量化 70B 模型需要 3 天 18 小时。

如果你需要在不增加更多 GPU 的情况下加速量化，你也可以增加 --relative_mse_tolerance 或设置 --init_max_points_per_centroid 或限制 --finetune_max_epochs。 然而，这通常会以降低模型精度为代价。

模型下载

该代码要求以 Huggingface 格式下载 LLaMA 模型并本地保存。以下脚本假设 $TRANSFORMERS_CACHE 变量指向 Huggingface Transformers 缓存文件夹。 要下载和缓存模型，请在相同环境中运行以下代码：

from transformers import AutoTokenizer, AutoModelForCausalLM
model_name = "meta-llama/Llama-2-7b-hf"  # 或其他你想下载的模型
tokenizer = AutoTokenizer.from_pretrained(model_name, torch_dtype="auto")
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")

如何使用 AQLM 量化模型

此脚本压缩模型，然后使用 WikiText2、C4 和 Penn Treebank 数据集测试其困惑度性能。

启动脚本的命令应如下所示：

export CUDA_VISIBLE_DEVICES=0   # 或者如 0,1,2,3
export MODEL_PATH=<HUB上模型的路径>
export DATASET_PATH=<插入数据集名称或自定义数据路径>
export SAVE_PATH=/保存量化模型的路径/
export WANDB_PROJECT=MY_AQ_EXPS
export WANDB_NAME=COOL_EXP_NAME

python main.py $MODEL_PATH $DATASET_PATH \
 --nsamples=1024 \
 --val_size=128 \
 --num_codebooks=1 \
 --nbits_per_codebook=16 \
 --in_group_size=8 \
 --relative_mse_tolerance=0.01 \
 --finetune_batch_size=32 \
 --finetune_max_epochs=10 \
 --finetune_early_stop=3 \
 --finetune_keep_best \
 --local_batch_size=1 \
 --offload_activations \
 --wandb \
 --resume \
 --save $SAVE_PATH

主要 CLI 参数：

• CUDA_VISIBLE_DEVICES - 默认情况下，代码将使用所有可用的 GPU。如果你想使用特定的 GPU(或一个 GPU)，请使用此变量。
• MODEL_PATH - 指向 Hugging Face hub (如 meta-llama/Llama-2-7b-hf) 或包含 transformers 模型和分词器的本地文件夹的路径。
• DATASET_PATH - 校准数据的路径(见上文)或标准数据集 [c4, ptb, wikitext2]
  • 对于 llama-2 模型，你可以使用 DATASET_PATH=./data/red_pajama_n=1024_4096_context_length.pth 获取 RedPajama 的一个切片(最多 1024 个样本)
• --nsamples - 校准数据序列的数量(训练 + 验证)。如果未设置此参数，则使用所有可用的校准数据。
• --val_size - 用于块微调早停的验证序列数量。默认为 0。必须小于 --nsamples。
• --num_codebooks - 每层的码本数量
• --nbits_per_codebook - 每个码本将包含 2 ** nbits_per_codebook 个向量
• --in_group_size - 一起量化多少个权重(即 arXiv 论文中的 "g")
• --finetune_batch_size - (仅用于微调) 每个优化步骤使用的序列总数
• --local_batch_size - 在累积 finetune_batch_size 时，每个 GPU 每次前向传播处理的样本数(影响 GPU 内存使用)
• --relative_mse_tolerance - (用于初始校准) - 当 (current_epoch_mse / previous_epoch_mse) > (1 - relative_mse_tolerance) 时停止训练
• --finetune_max_epochs - 块调优时通过校准数据的最大轮数。
• --finetune_early_stop - 在验证集上没有改善的情况下，通过校准数据的最大轮数。
• --offload_activations - 在校准期间，将激活从 GPU 内存移到 RAM。这减少了 VRAM 使用，同时使校准速度降低约 10%(取决于你的硬件)。
• --save - 保存/加载量化模型的路径。(另见：--load)
• --wandb - 如果设置此参数，代码将结果记录到 wandb
• --attn_implementation - 指定注意力实现(适用于 transformers >= 4.38)。Sdpa 注意力有时会引起问题，建议使用 eager 实现。

还有其他可用的超参数。运行 python main.py --help 以获取有关命令行参数的更多详细信息，包括压缩参数。

微调

注意 此代码仅会微调连续参数。要同时微调连续和离散参数，请切换到 pv-tuning 分支并按照其 readme 中的说明操作。

可以通过块微调进一步提高量化模型的准确性。首先，float16/bfloat16 的 logits 被缓存在 RAM 中。然后优化量化模型的可微参数，以最小化与教师 logits 的 KL 散度。通常，我们使用与模型量化相同的校准数据。

启动脚本的命令应如下所示：

python finetune.py \
  --base_model $MODEL_PATH \
  --quant_model $INPUT_PATH \
  --dataset $DATASET_PATH \
  --nsamples=<总数量> \
  --val_size=<验证集大小> \
  --lr=1e-5 \
  --adam_beta1=0.90 \
  --adam_beta2=0.999 \
  --epochs=5 \
  --early_stop=3 \
  --batch_size=8 \
  --microbatch_size=4 \
  --save $DATA_PATH \
  --gradient_checkpointing

主要的命令行参数：

• --base_model - 原始浮点模型的路径或名称
• --quant_model - 量化模型权重的路径
• --dataset - 校准数据集的路径或名称
• --nsamples - 校准数据序列的数量（训练+验证）。如果未设置此参数，则使用所有可用的校准数据。
• --val_size - 用于端到端微调早停的验证序列数量。默认为0。必须小于 --nsamples。
• --gradient_checkpointing - 是否使用梯度检查点。以较长的运行时间为代价减少峰值内存使用。
• --finetune_dtype - 微调时应使用的数据类型。默认为 float32。
• --amp - 是否在微调时使用自动混合精度。需要 --finetune_dtype=float32。

对于较大的模型，需要多GPU训练。目前尚未实现FSDP训练，模型在单个进程上进行微调，参数分布在可用设备上。

通过LM评估工具进行零样本基准测试

为了进行零样本评估，我们采用了语言模型评估工具框架。我们的代码适用于标准 transformers 格式的模型，并可以通过 --aqlm_checkpoint_path 参数（可选）加载量化模型的权重。

PV-Tuning中的评估结果是使用 lm-eval=0.4.0 生成的。

要运行评估，请确保安装了正确版本或通过以下命令安装： pip install lm-eval==0.4.0。

启动评估程序的主脚本是 lmeval.py。

export CUDA_VISIBLE_DEVICES=0,1,2,3  # 可选：选择GPU
export QUANTIZED_MODEL=<从MAIN.py保存的量化模型路径>
export MODEL_PATH=<插入HUB上原始模型的路径>
export DATASET=<插入数据集名称或自定义数据路径>
export WANDB_PROJECT=MY_AQLM_EVAL
export WANDB_NAME=COOL_EVAL_NAME

# 用于0样本评估
python lmeval.py \
    --model hf \
    --model_args pretrained=$MODEL_PATH,dtype=float16,parallelize=True \
    --tasks winogrande,piqa,hellaswag,arc_easy,arc_challenge \
    --batch_size <评估批次大小> \
    --aqlm_checkpoint_path QUANTIZED_MODEL # 如果评估量化模型

# 用于5样本MMLU
python lmeval.py \
    --model hf \
    --model_args pretrained=$MODEL_PATH,dtype=float16,parallelize=True \
    --tasks mmlu \
    --batch_size <评估批次大小> \
    --num_fewshot 5 \
    --aqlm_checkpoint_path QUANTIZED_MODEL # 如果评估量化模型

准备模型进行推理

要将模型转换为兼容Hugging Face的格式，请使用带有相应参数的 convert_to_hf.py model in_path out_path：

• model - 原始预训练模型（对应 main.py 的 MODEL_PATH，例如 meta-llama/Llama-2-7b-hf）。
• in_path - 包含初始量化模型的文件夹（对应 main.py 的 --save）。
• out_path - 保存 transformers 模型的文件夹。

您还可以指定诸如 --save_safetensors 之类的标志来控制保存的模型格式（详见 --help）。

示例命令：python convert_to_hf.py meta-llama/Llama-2-7b-hf ./path/to/saved/quantization ./converted-llama2-7b-hf --save_safetensors

贡献

如果您想贡献一些实质性的内容（不仅仅是修正错别字），请先开一个问题。 我们对所有拉取请求使用black和isort。在提交代码之前，请运行 black . && isort .

引用

如果您觉得这项工作有用，请考虑引用：

@misc{egiazarian2024extreme,
      title={Extreme Compression of Large Language Models via Additive Quantization}, 
      author={Vage Egiazarian and Andrei Panferov and Denis Kuznedelev and Elias Frantar and Artem Babenko and Dan Alistarh},
      year={2024},
      eprint={2401.06118},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}
@misc{malinovskii2024pvtuning,
      title={PV-Tuning: Beyond Straight-Through Estimation for Extreme LLM Compression}, 
      author={Vladimir Malinovskii and Denis Mazur and Ivan Ilin and Denis Kuznedelev and Konstantin Burlachenko and Kai Yi and Dan Alistarh and Peter Richtarik},
      year={2024},
      eprint={2405.14852},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}