Transformers 文档
压缩张量
并获得增强的文档体验
开始使用
压缩张量
compressed-tensors 扩展了 safetensors 文件,以支持压缩张量数据类型,从而为存储和加载各种量化和稀疏格式(如稠密、int 量化 (int8)、float 量化 (fp8) 以及打包量化(int4 或 int8 权重量化打包到 int32 中))提供统一的检查点格式。
compressed-tensors 支持使用 PEFT 进行微调,并包含以下功能。
- fp8、int4、int8 权重和激活精度。
- 用于 张量、通道、组、块、token 的量化比例和零点策略。
- 动态的 per-token 激活量化(或任何静态策略)。
- 权重稀疏性(非结构化或半结构化,如 2:4)可以与量化结合使用以实现极高的压缩率。
- 任意模块的量化,不仅仅是 nn.Linear 模块。
- 通过名称或类针对特定模块的支持。
从 PyPI 安装 compressed-tensors 以获取最新的稳定版本(推荐),或者从源代码安装以获取最新的功能。
pip install compressed-tensors
使用 compressed-tensors 标签 搜索以在 Hugging Face Hub 上查找兼容的模型。
目前只能加载已经量化的模型,并且一旦加载模型,就无法保存。要将模型量化为 compressed-tensors 格式,请参阅 llm-compressor。或者,可以独立创建模型并使用 compressed-tensors 配置进行序列化。
from transformers import AutoModelForCausalLM
ct_model = AutoModelForCausalLM.from_pretrained("nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf", device_map="auto")
# measure memory usage
mem_params = sum([param.nelement()*param.element_size() for param in ct_model.parameters()])
print(f"{mem_params/2**30:.4f} GB")
# 8.4575 GB
模型检查点
compressed-tensor 模型通过其配置条目定义。以下示例取自 nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf 的 config.json
文件。
有很多条目允许在压缩期间和压缩后进行灵活的表达,但用于加载和推理的条目可以简化为仅关注几个关键条目。
"quantization_config": {
"config_groups": {
"group_0": {
"input_activations": {
"num_bits": 8,
"strategy": "tensor",
"type": "float"
},
"targets": ["Linear"],
"weights": {
"num_bits": 8,
"strategy": "tensor",
"type": "float"
}
}
},
"format": "naive-quantized",
"ignore": ["lm_head"],
"quant_method": "compressed-tensors",
"quantization_status": "frozen"
},
配置文件指定了配置组 (group_0
) 的量化,其中包括权重和激活量化到 fp8,并采用静态的 per-tensor 策略。 lm_head
模块未量化,如 ignore
键所示。
要更详细地了解模型权重,请使用模型卡上的 safetensors 查看器 查看所有 nn.Linear 模块的量化权重、输入比例和权重比例。
张量 | 形状 | 精度 |
---|---|---|
model.layers.0.input_layernorm.weight | [4 096] | BF16 |
model.layers.0.mlp.down_proj.input_scale | [1] | BF16 |
model.layers.0.mlp.down_proj.weight | [4 096, 14 336] | F8_E4M3 |
model.layers.0.mlp.down_proj.weight_scale | [1] | BF16 |
model.layers.0.mlp.gate_proj.input_scale | [1] | BF16 |
model.layers.0.mlp.gate_proj.weight | [14 336, 4 096] | F8_E4M3 |
model.layers.0.mlp.gate_proj.weight_scale | [1] | BF16 |
model.layers.0.mlp.up_proj.input_scale | [1] | BF16 |
model.layers.0.mlp.up_proj.weight | [14 336, 4 096] | F8_E4M3 |
model.layers.0.mlp.up_proj.weight_scale | [1] | BF16 |
model.layers.0.post_attention_layernorm.weight | [4 096] | BF16 |
model.layers.0.self_attn.k_proj.input_scale | [1] | BF16 |
model.layers.0.self_attn.k_proj.weight | [1 024, 4 096] | F8_E4M3 |
model.layers.0.self_attn.k_proj.weight_scale | [1] | BF16 |
model.layers.0.self_attn.o_proj.input_scale | [1] | BF16 |
model.layers.0.self_attn.o_proj.weight | [4 096, 4 096] | F8_E4M3 |
model.layers.0.self_attn.o_proj.weight_scale | [1] | BF16 |
model.layers.0.self_attn.q_proj.input_scale | [1] | BF16 |
model.layers.0.self_attn.q_proj.weight | [4 096, 4 096] | F8_E4M3 |
model.layers.0.self_attn.q_proj.weight_scale | [1] | BF16 |
model.layers.0.self_attn.v_proj.input_scale | [1] | BF16 |
model.layers.0.self_attn.v_proj.weight | [1 024, 4 096] | F8_E4M3 |
model.layers.0.self_attn.v_proj.weight_scale | [1] | BF16 |
当使用 ~quantizers.HFQuantizer
集成加载 compressed-tensors 模型时,量化配置中指定的所有 nn.Linear 模块都将被 CompressedLinear 模块替换,这些模块管理压缩权重和用于推理的前向传递。 lm_head
模块仍然保留为未量化的 nn.Linear 模块。
from transformers import AutoModelForCausalLM
ct_model = AutoModelForCausalLM.from_pretrained("nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf")
print(ct_model)
"""
LlamaForCausalLM(
(model): LlamaModel(
(embed_tokens): Embedding(128256, 4096)
(layers): ModuleList(
(0-31): 32 x LlamaDecoderLayer(
(self_attn): LlamaSdpaAttention(
(q_proj): CompressedLinear(
in_features=4096, out_features=4096, bias=False
(input_observer): MovingAverageMinMaxObserver()
(weight_observer): MovingAverageMinMaxObserver()
)
(k_proj): CompressedLinear(
in_features=4096, out_features=1024, bias=False
(input_observer): MovingAverageMinMaxObserver()
(weight_observer): MovingAverageMinMaxObserver()
)
(v_proj): CompressedLinear(
in_features=4096, out_features=1024, bias=False
(input_observer): MovingAverageMinMaxObserver()
(weight_observer): MovingAverageMinMaxObserver()
)
(o_proj): CompressedLinear(
in_features=4096, out_features=4096, bias=False
(input_observer): MovingAverageMinMaxObserver()
(weight_observer): MovingAverageMinMaxObserver()
)
(rotary_emb): LlamaRotaryEmbedding()
)
(mlp): LlamaMLP(
(gate_proj): CompressedLinear(
in_features=4096, out_features=14336, bias=False
(input_observer): MovingAverageMinMaxObserver()
(weight_observer): MovingAverageMinMaxObserver()
)
(up_proj): CompressedLinear(
in_features=4096, out_features=14336, bias=False
(input_observer): MovingAverageMinMaxObserver()
(weight_observer): MovingAverageMinMaxObserver()
)
(down_proj): CompressedLinear(
in_features=14336, out_features=4096, bias=False
(input_observer): MovingAverageMinMaxObserver()
(weight_observer): MovingAverageMinMaxObserver()
)
(act_fn): SiLU()
)
(input_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
(post_attention_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
)
)
(norm): LlamaRMSNorm((4096,), eps=1e-05)
(rotary_emb): LlamaRotaryEmbedding()
)
(lm_head): Linear(in_features=4096, out_features=128256, bias=False)
)
"""