AutoGPTQ大語言模型量化工具包

from transformers import AutoTokenizer, TextGenerationPipeline
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig


pretrained_model_dir = "facebook/opt-125m"
quantized_model_dir = "opt-125m-4bit"


tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True)
examples = [
    tokenizer(
        "auto-gptq is an easy-to-use model quantization library with user-friendly apis, based on GPTQ algorithm."
    )
]

quantize_config = BaseQuantizeConfig(
    bits=4,  # 將模型量化為 4-bit 數(shù)值類型
    group_size=128,  # 一般推薦將此參數(shù)的值設(shè)置為 128
    desc_act=False,  # 設(shè)為 False 可以顯著提升推理速度，但是 ppl 可能會輕微地變差
)

# 加載未量化的模型，默認(rèn)情況下，模型總是會被加載到 CPU 內(nèi)存中
model = AutoGPTQForCausalLM.from_pretrained(pretrained_model_dir, quantize_config)

# 量化模型, 樣本的數(shù)據(jù)類型應(yīng)該為 List[Dict]，其中字典的鍵有且僅有 input_ids 和 attention_mask
model.quantize(examples)

# 保存量化好的模型
model.save_quantized(quantized_model_dir)

# 使用 safetensors 保存量化好的模型
model.save_quantized(quantized_model_dir, use_safetensors=True)

# 將量化好的模型直接上傳至 Hugging Face Hub 
# 當(dāng)使用 use_auth_token=True 時(shí), 確保你已經(jīng)首先使用 huggingface-cli login 進(jìn)行了登錄
# 或者可以使用 use_auth_token="hf_xxxxxxx" 來顯式地添加賬戶認(rèn)證 token
# （取消下面三行代碼的注釋來使用該功能）
# repo_id = f"YourUserName/{quantized_model_dir}"
# commit_message = f"AutoGPTQ model for {pretrained_model_dir}: {quantize_config.bits}bits, gr{quantize_config.group_size}, desc_act={quantize_config.desc_act}"
# model.push_to_hub(repo_id, commit_message=commit_message, use_auth_token=True)

# 或者你也可以同時(shí)將量化好的模型保存到本地并上傳至 Hugging Face Hub
# （取消下面三行代碼的注釋來使用該功能）
# repo_id = f"YourUserName/{quantized_model_dir}"
# commit_message = f"AutoGPTQ model for {pretrained_model_dir}: {quantize_config.bits}bits, gr{quantize_config.group_size}, desc_act={quantize_config.desc_act}"
# model.push_to_hub(repo_id, save_dir=quantized_model_dir, use_safetensors=True, commit_message=commit_message, use_auth_token=True)

# 加載量化好的模型到能被識別到的第一塊顯卡中
model = AutoGPTQForCausalLM.from_quantized(quantized_model_dir, device="cuda:0")

# 從 Hugging Face Hub 下載量化好的模型并加載到能被識別到的第一塊顯卡中
# model = AutoGPTQForCausalLM.from_quantized(repo_id, device="cuda:0", use_safetensors=True, use_triton=False)

# 使用 model.generate 執(zhí)行推理
print(tokenizer.decode(model.generate(**tokenizer("auto_gptq is", return_tensors="pt").to(model.device))[0]))

# 或者使用 TextGenerationPipeline
pipeline = TextGenerationPipeline(model=model, tokenizer=tokenizer)
print(pipeline("auto-gptq is")[0]["generated_text"])

from auto_gptq.modeling import BaseGPTQForCausalLM


class OPTGPTQForCausalLM(BaseGPTQForCausalLM):
    # chained attribute name of transformer layer block
    layers_block_name = "model.decoder.layers"
    # chained attribute names of other nn modules that in the same level as the transformer layer block
    outside_layer_modules = [
        "model.decoder.embed_tokens", "model.decoder.embed_positions", "model.decoder.project_out",
        "model.decoder.project_in", "model.decoder.final_layer_norm"
    ]
    # chained attribute names of linear layers in transformer layer module
    # normally, there are four sub lists, for each one the modules in it can be seen as one operation, 
    # and the order should be the order when they are truly executed, in this case (and usually in most cases), 
    # they are: attention q_k_v projection, attention output projection, MLP project input, MLP project output
    inside_layer_modules = [
        ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"],
        ["self_attn.out_proj"],
        ["fc1"],
        ["fc2"]
    ]