Transformers documentation
compressed-tensors
compressed-tensors
compressed-tensors는 safetensors 파일을 압축된 텐서 데이터 유형으로 확장하여, dense, int-quantized(int8), float-quantized(fp8), pack-quantized(int4 또는 int8 weight-quantized를 int32로 패킹)와 같은 다양한 양자화 형식을 저장하고 로드하기 위한 통합 체크포인트 형식을 제공합니다.
compressed-tensors는 PEFT를 활용한 fine-tuning을 지원하며, 다음 기능도 포함합니다.
- fp8, int4, int8 가중치 및 활성화 정밀도.
- tensor, channel, group, block, token에 대한 양자화 스케일 및 제로 포인트 전략.
- 토큰별 동적 활성화 양자화(또는 모든 정적 전략).
- Quantization of arbitrary modules, not just nn.Linear modules.
- 이름 또는 클래스별로 특정 모듈에 대한 대상 지정 지원.
최신 안정 버전(권장)을 설치하려면 PyPI에서 compressed-tensors를 설치하고, 최신 기능을 사용하려면 소스에서 설치하세요.
pip install compressed-tensors
git clone https://github.com/neuralmagic/compressed-tensors
cd compressed-tensors
pip install -e .Hugging Face Hub에서 tag를 사용하여 압축된 텐서(compressed-tensors)로 검색하여 호환되는 모델을 찾으세요.
사전 양자화 모델은 직접 로드할 수 있습니다. 모델을 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-tensors 체크포인트는 가중치를 압축된 상태(fp8 또는 packed int4/int8)로 저장합니다. 실행 방식은 CompressedTensorsConfig의 두 가지 인자에 의해 결정됩니다.
| 구성 | 로드 후 가중치 | 실행 |
|---|---|---|
| default | compressed-tensors 형식으로 유지됩니다: 직렬화된 모든 텐서(weight_packed, weight_scale, weight_zero_point, …)는 저장된 그대로 로드됩니다 | compressed-tensors가 레이어를 관리하며 첫 번째 forward pass에서 모델을 압축 해제합니다 |
dequantize=True | dtype(예: BF16)로 모델에 로드하는 동안 압축이 해제됩니다 | quantization_scheme가 없는 일반 dense 모델이므로 w8a8 activations도 양자화되지 않습니다. 해당 dtype으로 fine-tuning하거나 저장할 수 있습니다. |
use_optimized_inference=True | 양자화 상태로 유지 | kernel이 실행되는 스킴을 가진 레이어, 현재 W8A8 fp8; 추론 전용 |
기본값으로 save_pretrained()는 압축된 텐서를 변경 없이 다시 기록하므로, 로드된 모델은 동일한 compressed-tensors 형식으로 저장됩니다.
FP8 커널 가속
FP8 compressed-tensors 모델이 FP8로 유지되고, matmul 연산이 하드웨어 가속 FP8 커널(torch.nn.functional.scaled_mm, torch._scaled_mm_v2로 디스패치되며, 이전 torch 버전은 torch._scaled_mm로 폴백)을 통해 수행되도록 use_optimized_inference=True를 전달하세요. 이를 통해 가중치를 BF16으로 역양자화하지 않고, 추론 전반에 걸쳐 가중치를 FP8로 유지하여 메모리 사용량을 줄이고 계산 속도를 높일 수 있습니다. 이 기능은 추론 전용이므로, 파인튜닝 시에는 비활성화해 두세요.
| 장치 | 커널 | The input is a single word "Notes". 참고 |
|---|---|---|
| Intel XPU | torch.nn.functional.scaled_mm | FP8을 지원하는 모든 XPU 장치 |
| NVIDIA CUDA (SM89+) | torch.nn.functional.scaled_mm | Ada Lovelace (L4, L40), Hopper (H100), Blackwell 및 이후 |
| Apple silicon (MPS) | torch.nn.functional.scaled_mm | torch >= 2.15가 필요합니다 |
| CPU / CUDA SM80 (A100) / MPS, torch < 2.15 | 폴백 | use_optimized_inference=True는 무시되며, 모델은 역양자화 상태로 실행됩니다 |
FP8 커널 경로는 이러한 양자화 레이아웃을 지원합니다.
| 전략 | 예시 모델 |
|---|---|
| 채널별 동적 | RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic |
| 텐서별 정적 | RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8 |
FP8 양자화 모델 로드
FP8 커널은 선택적으로 사용되며, use_optimized_inference=True를 지정하여 요청할 수 있습니다. 모델의 config에 FP8 양자화가 지정되어 있고 지원되는 GPU가 사용 가능한 경우 해당 커널이 사용됩니다.
from transformers import AutoModelForCausalLM, AutoTokenizer, CompressedTensorsConfig
model = AutoModelForCausalLM.from_pretrained(
"RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic",
quantization_config=CompressedTensorsConfig(use_optimized_inference=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic")
inputs = tokenizer("Hello, how are you?", return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))로드 시점의 역양자화
use_optimized_inference=True를 사용하지 않으면 모델은 일반적인 compressed-tensors 경로를 따릅니다. 이 경우 가중치는 압축된 상태로 유지되며, compressed-tensors가 첫 번째 forward pass 시점에 이를 압축 해제합니다. 모델의 fine-tuning을 하거나 원래 정밀도(예: BF16)로 저장하려면, 대신 로딩 시점에 압축 해제를 수행하도록 dequantize=True를 전달하세요. 그 결과 일반적인 dense 모델이 생성됩니다. compressed-tensors 스킴이 연결되지 않으므로 activation quantization(예: w8a8)도 함께 제거됩니다.
from transformers import AutoModelForCausalLM, CompressedTensorsConfig
model = AutoModelForCausalLM.from_pretrained(
"RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic",
quantization_config=CompressedTensorsConfig(dequantize=True),
device_map="auto",
)모델 체크포인트
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"
},config 파일은 config 그룹(group_0)의 양자화를 지정하며, 여기에는 정적 per-tensor 전략을 사용한 fp8로의 weight 및 activation 양자화가 포함됩니다. ignore 키에 표시된 대로 lm_head 모듈은 양자화되지 않습니다.
모델 가중에 대해 더 자세히 살펴보기 위해 모델 카드의 safetensors viewer를 사용하면 모든 nn.Linear 모듈의 양자화 가중치, 입력 스케일, 가중치 스케일을 확인할 수 있습니다.
| 텐서 | Shape | 정밀도 |
|---|---|---|
| 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 모델을 로드할 때, 대상 모듈은 compressed-tensors로 전달됩니다. compressed-tensors는 해결된 quantization_scheme를 연결하고, quantization_status를 설정하며, 체크포인트에 저장된 매개변수(weight는 fp8, 그리고 정적 전략의 경우 weight_scale와 input_scale)를 등록하고, 해당 모듈에 자체 forward pass를 설치합니다. 모듈은 nn.Linear 인스턴스로 유지되므로 print 출력에도 그렇게 표시됩니다. 최근 compressed-tensors 버전에서는 더 이상 모듈을 CompressedLinear 하위 클래스로 래핑하지 않습니다. ignore에 나열된 lm_head와 같은 모듈은 그대로 유지됩니다.
dequantize=False(기본값)인 경우, 로딩이 완료된 후에도 가중치는 여전히 압축된 상태로 유지되며, compressed-tensors가 첫 번째 forward pass에서 전체 모델을 압축 해제합니다. dequantize=True인 경우, 각 가중치를 로딩하는 동안 압축을 해제하므로 dense 가중치를 얻기 위해 forward pass가 필요하지 않으며, 모듈은 quantization_scheme이 없는 일반 dense 레이어가 됩니다.
import torch
from transformers import AutoModelForCausalLM, CompressedTensorsConfig
model_id = "nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf"
ct_model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=CompressedTensorsConfig(dequantize=False),
device_map="auto",
)
q_proj = ct_model.model.layers[0].self_attn.q_proj
print(q_proj, q_proj.quantization_status)
# Linear(in_features=4096, out_features=4096, bias=False) QuantizationStatus.COMPRESSED
# ^ compressed-tensors module: fp8 weight, weight_scale, and its own forward
ct_model(input_ids=torch.tensor([[0, 1, 2]], device=ct_model.device))
print(q_proj, q_proj.quantization_status)
# Linear(in_features=4096, out_features=4096, bias=False) QuantizationStatus.DECOMPRESSED
# ^ weight is BF16 now, decompressed by that forward pass
ct_model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=CompressedTensorsConfig(dequantize=True),
device_map="auto",
)
print(ct_model.model.layers[0].self_attn.q_proj)
# Linear(in_features=4096, out_features=4096, bias=False) weight: BF16use_optimized_inference=True를 사용하면 fp8 config 그룹에 해당하는 레이어는 CompressedTensorsFP8Linear로 대체됩니다. 이 클래스는 row-wise matmul 커널이 요구하는 레이아웃으로 fp8 가중치와 해당 scale을 보유합니다. 이러한 가중치는 forward pass를 포함하여 fp8 상태로 유지됩니다.
from transformers import AutoModelForCausalLM, CompressedTensorsConfig
ct_model = AutoModelForCausalLM.from_pretrained(
"nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf",
quantization_config=CompressedTensorsConfig(use_optimized_inference=True),
device_map="auto",
)
print(ct_model.model.layers[0].self_attn.q_proj)
# CompressedTensorsFP8Linear(in_features=4096, out_features=4096, bias=False) weight: F8_E4M3