In [ ]:
pip install --upgrade torch torchao
默认情况下，权重会以全精度（torch.float32）加载，无论实际存储的权重数据类型是什么（例如 torch.float16）。设置 torch_dtype="auto" 可以根据模型的 config.json 文件中定义的数据类型自动加载最节省内存的数据类型。
In [ ]:
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))
序列化和反序列化¶
TorchAO 的量化是通过 张量子类 实现的，因此仅支持 Hugging Face 的非安全张量序列化和反序列化。它依赖于 torch.load(..., weights_only=True) 来避免加载时执行任意用户代码，并使用 add_safe_globals 来白名单一些已知的用户函数。
不支持安全张量序列化的原因是，张量子类允许最大的灵活性，因此我们希望支持新的量化张量格式的工作量较低。而安全张量则优化了最大安全性（不执行用户代码），这意味着我们必须手动支持新的量化格式。
In [ ]:
# 将量化模型保存到本地
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))