pip install --upgrade torch torchao

import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Meta-Llama-3-8B"

# 支持的量化配置有 int4_weight_only、int8_weight_only 和 int8_dynamic_activation_int8_weight
# 更多参数示例和文档可以在这里找到：https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques
quantization_config = TorchAoConfig("int4_weight_only", group_size=128)

# 加载量化模型
quantized_model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    torch_dtype="auto", 
    device_map="auto", 
    quantization_config=quantization_config
)

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_name)
input_text = "我们晚餐吃什么？"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

# 编译量化模型以加速
import torchao
torchao.quantization.utils.recommended_inductor_config_setter()
quantized_model = torch.compile(quantized_model, mode="max-autotune")

# 生成输出
output = quantized_model.generate(**input_ids, max_new_tokens=10)
print(tokenizer.decode(output[0], skip_special_tokens=True))

# 测试性能
import torch.utils.benchmark as benchmark

def benchmark_fn(f, *args, **kwargs):
    # 预热
    for _ in range(5):
        f(*args, **kwargs)
        
    t0 = benchmark.Timer(
        stmt="f(*args, **kwargs)",
        globals={"args": args, "kwargs": kwargs, "f": f},
        num_threads=torch.get_num_threads(),
    )
    return f"{(t0.blocked_autorange().mean):.3f}"

MAX_NEW_TOKENS = 1000
print("int4wo-128 模型:", benchmark_fn(quantized_model.generate, **input_ids, max_new_tokens=MAX_NEW_TOKENS))

# 使用 bfloat16 模型进行对比
bf16_model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cuda", torch_dtype=torch.bfloat16)
bf16_model = torch.compile(bf16_model, mode="max-autotune")
print("bf16 模型:", benchmark_fn(bf16_model.generate, **input_ids, max_new_tokens=MAX_NEW_TOKENS))

# 将量化模型保存到本地
output_dir = "llama3-8b-int4wo-128"
quantized_model.save_pretrained(output_dir, safe_serialization=False)

# 推送到 Hugging Face Hub
# save_to = "{user_id}/llama3-8b-int4wo-128"
# quantized_model.push_to_hub(save_to, safe_serialization=False)

# 加载量化模型
ckpt_id = "llama3-8b-int4wo-128"  # 或者 Hugging Face Hub 模型 ID
loaded_quantized_model = AutoModelForCausalLM.from_pretrained(ckpt_id, device_map="cuda")

# 确认加速效果
loaded_quantized_model = torch.compile(loaded_quantized_model, mode="max-autotune")
print("加载的 int4wo-128 模型:", benchmark_fn(loaded_quantized_model.generate, **input_ids, max_new_tokens=MAX_NEW_TOKENS))

TorchAO —— 用于 PyTorch 的架构优化库¶

序列化和反序列化¶