文章首发于 slightwind.cn
Mistral-7B’s performance on 5-shot MMLU
Mistral-7B 是一个很强的 7B 开源模型,在 Mistral 官网和论文中声称可以在 5-shot MMLU 上达到 60.1% 的准确率,首先下载官方的模型权重文件(Mistral-7B-v0.1)并直接在原精度(BF16)上进行推理,尝试复现出官方的准确率。
llmtask
这里使用 llmtask 来进行下游任务测试,非常方便快捷,只需要
pip install llmtask==0.0.2
即可完成安装,可以直接测试模型在 C-Eval 和 MMLU 数据集上的表现。
示例代码:
import random
from llmtask import TaskGenerator
choices = ("A", "B", "C", "D")
TG = TaskGenerator("mmlu", max_shot=4)
for task in TG:
TG.feedback(random.choice(choices))
print(TG.summary())
测试 Mistral-7B 原精度推理脚本:
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from llmtask import TaskGenerator
def log(msg):
with open("mmlu_5shot_bf16.log", "a") as f:
f.write(f"{msg}\n")
device = "cuda"
model = AutoModelForCausalLM.from_pretrained("/path/to/Mistral-7B-v0.1", torch_dtype=torch.bfloat16).to(device)
tokenizer = AutoTokenizer.from_pretrained("/path/to/Mistral-7B-v0.1")
cnt = 0
TG = TaskGenerator("mmlu", max_shot=5)
for task in TG:
model_inputs = tokenizer([task], return_tensors="pt").to(device)
input_tokens = len(model_inputs['input_ids'][0])
t0 = time.time()
generated_ids = model.generate(**model_inputs, max_new_tokens=1, pad_token_id=tokenizer.eos_token_id)
ans = tokenizer.batch_decode([generated_ids[0][input_tokens:]])[0]
log(f"[{cnt:5}] [{(time.time() - t0):5.3f} s] => ans:{ans}")
cnt += 1
TG.feedback(ans)
log(TG.summary())
torch.cuda.empty_cache()
测试结果如下(每次只推理一个 Token 作为模型选择的答案,很快就可以测试完成):
Precision | Avg (%) | STEM (%) | Social Science (%) | Humanities (%) | Other (%) | Total Time (s) |
---|---|---|---|---|---|---|
BF16 | 61.00 61.00 61.00 | 50.46 50.46 50.46 | 75.07 75.07 75.07 | 53.47 53.47 53.47 | 68.16 68.16 68.16 | 312.79 312.79 312.79 |
平均每道题耗时 204 204 204ms,最后的测试结果还算比较接近官方的结果,以此作为 baseline 和量化后的模型权重对比推理下游任务准确率的损失情况。
8bit/4bit Quantization
量化使用 transformers 内置的 bitsandbytes 提供的 LLM.int8()
作为 8bit 量化算法(threshold=6.0
),4bit 量化包含两种 4bit 的数据类型 FP4 和 NF4,以及 torch.float32
和 torch.float16
两种计算类型,接下来分别对这些场景进行测试。
8bit
进行 8bit 推理只需要修改加载权重的这一行即可:
虽然官方已经不推荐这样做了,但是这里不需要在
BitsAndBytesConfig
配置额外的参数,可以直接这样使用默认参数。
model = AutoModelForCausalLM.from_pretrained("/path/to/Mistral-7B-v0.1", load_in_8bit=True)
8bit 量化后平均每道题耗时 401 401 401ms,测试结果如下:
Precision | Avg (%) | STEM (%) | Social Science (%) | Humanities (%) | Other (%) | Total Time (s) |
---|---|---|---|---|---|---|
INT8 | 60.87 60.87 60.87 | 51.09 51.09 51.09 | 73.59 73.59 73.59 | 52.89 52.89 52.89 | 69.29 69.29 69.29 | 614.43 614.43 614.43 |
4bit
通过 BitsAndBytesConfig
来配置量化类型(FP4
/NF4
)测试脚本:
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from llmtask import TaskGenerator
def log(msg):
with open("mmlu_5shot_fp4_fp16.log", "a") as f:
f.write(f"{msg}\n")
device = "cuda"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_quant_type="fp4",
bnb_4bit_use_double_quant=False,
bnb_4bit_quant_storage=torch.uint8
)
model = AutoModelForCausalLM.from_pretrained("/path/to/Mistral-7B-v0.1", quantization_config=bnb_config)
tokenizer = AutoTokenizer.from_pretrained("/path/to/Mistral-7B-v0.1")
TG = TaskGenerator("mmlu", max_shot=5)
cnt = 0
for task in TG:
model_inputs = tokenizer([task], return_tensors="pt").to(device)
input_tokens = len(model_inputs['input_ids'][0])
t0 = time.time()
generated_ids = model.generate(**model_inputs, max_new_tokens=1, pad_token_id=tokenizer.eos_token_id)
ans = tokenizer.batch_decode([generated_ids[0][input_tokens:]])[0]
log(f"[{cnt:5}] [{(time.time() - t0):5.3f} s] => ans:{ans}")
cnt += 1
TG.feedback(ans)
log(TG.summary())
torch.cuda.empty_cache()
下面是改变其中某个参数后在 MMLU 数据集上的准确率,可以看出即使是 4bit 对准确率影响都没有很大,首 Token 性能还可以接近原精度,还节省了大量的空间。
Quant Type | Compute Dtype | Double Quant | Avg (%) | Total Time (s) |
---|---|---|---|---|
FP4 | FP16 | False | 59.37 59.37 59.37 | 347.00 347.00 347.00 |
FP4 | FP16 | True | 59.17 59.17 59.17 | 353.22 353.22 353.22 |
FP4 | FP32 | False | 59.50 59.50 59.50 | 1061.27 1061.27 1061.27 |
NF4 | FP16 | False | 59.04 59.04 59.04 | 361.19 361.19 361.19 |
Versions
Python Packages | Version |
---|---|
torch | 2.2.1 |
transformers | 4.39.1 |
bitsandbytes | 0.43.0 |
accelerate | 0.28.0 |
llmtask | 0.0.2 |