压缩张量存储库使用指南

1. 项目介绍

compressed-tensors 是一个开源库，它扩展了 safetensors 格式，提供了一种灵活且高效的方式来存储和管理压缩的张量数据。该库支持多种量化方法和稀疏模式，可以统一处理不同的模型优化技术，如 GPTQ、AWQ、SmoothQuant、INT8、FP8、SparseGPT 等。它的目标是简化模型部署流程，降低支持多种压缩格式在推理引擎中的开销。

2. 项目快速启动

首先，您需要安装 compressed-tensors 库。可以通过以下命令从 PyPI 安装稳定版本：

pip install compressed-tensors

或者，如果您希望安装最新版本，可以使用以下命令：

pip install --pre compressed-tensors

接下来，以下是一个简单的示例，展示如何使用 BitmaskConfig 来压缩和加载张量：

from compressed_tensors import save_compressed, load_compressed, BitmaskConfig
from torch import Tensor
from typing import Dict

# 创建压缩配置
compression_config = BitmaskConfig()

# 定义一些张量
tensors: Dict[str, Tensor] = {
    "tensor_1": Tensor([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])
}

# 压缩张量并保存到磁盘
save_compressed(tensors, "model.safetensors", compression_format=compression_config.format)

# 从磁盘加载压缩的张量
decompressed_tensors = {}
for tensor_name, tensor in load_compressed("model.safetensors", compression_config=compression_config):
    decompressed_tensors[tensor_name] = tensor

3. 应用案例和最佳实践

以下是一个使用 compressed-tensors 库压缩整个模型的示例：

from compressed_tensors import save_compressed_model, load_compressed
from transformers import AutoModelForCausalLM
from compressed_tensors import BitmaskConfig

# 加载预训练模型
model_name = "neuralmagic/llama2.c-stories110M-pruned50"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")
original_state_dict = model.state_dict()

# 创建压缩配置
compression_config = BitmaskConfig()

# 保存压缩的模型权重
save_compressed_model(model, "compressed_model.safetensors", compression_format=compression_config.format)

# 加载压缩的模型权重
state_dict = dict(load_compressed("compressed_model.safetensors", compression_config=compression_config))

4. 典型生态项目

目前，compressed-tensors 已被设计为与 Hugging Face 模型和 PyTorch 无缝集成，这意味着开发者可以轻松地尝试组合不同的量化方法，简化模型部署流程。此外，由于它的开源特性，社区中可能有更多的集成和扩展项目出现，以进一步丰富其生态系统。

以上就是关于 compressed-tensors 的基本介绍、快速启动指南、应用案例和生态项目概述。希望这些信息能够帮助您更好地了解和使用这个项目。