compressed-tensors

compressed-tensors

compressed-tensors는 safetensors 파일을 압축된 텐서 데이터 타입으로 확장해, dense, int 양자화(int8), float 양자화(fp8), pack 양자화(int4 또는 int8 가중치를 int32로 패킹한 것) 같은 다양한 양자화 형식을 저장·로드하기 위한 통합 체크포인트 형식을 제공합니다.

출처: 문서

본문

compressed-tensors는 PEFT로 미세 조정을 지원하며 다음 기능도 포함합니다.

• fp8, int4, int8 가중치 및 활성화 정밀도.
• tensor, channel, group, block, token을 위한 양자화 스케일·제로 포인트 전략.
• 동적 토큰별(dynamic per-token) 활성화 양자화(또는 임의의 정적 전략).
• nn.Linear 모듈뿐만 아니라 임의 모듈의 양자화.
• 이름이나 클래스로 특정 모듈을 타겟팅한 지원.

compressed-tensors를 PyPI에서 설치하면 최신 안정 릴리스(권장)를 얻고, 소스에서 설치하면 최신 기능을 얻을 수 있어요.

pip install compressed-tensors

git clone https://github.com/neuralmagic/compressed-tensors
cd compressed-tensors
pip install -e .

compressed-tensors 태그로 검색해 Hugging Face Hub에서 호환되는 모델을 찾으세요.

사전 양자화된 모델은 바로 로드할 수 있습니다. 모델을 compressed-tensors 형식으로 양자화하려면 llm-compressor를 보세요. 또는 모델을 독립적으로 만들고 compressed-tensors config로 직렬화할 수도 있어요.

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

로딩 모드 (Loading modes)

compressed-tensors 체크포인트는 가중치를 압축된 상태(fp8, 또는 패킹된 int4/int8)로 저장합니다. 어떻게 실행되는지는 두 개의 CompressedTensorsConfig 인자에 달려 있습니다.

FP8 커널 가속 (FP8 kernel acceleration)

use_optimized_inference=True를 넘기면 FP8 compressed-tensors 모델을 FP8로 유지하고, 가중치를 BF16으로 역양자화하는 대신 하드웨어 가속 FP8 커널(torch.nn.functional.scaled_mm, torch._scaled_mm_v2로 디스패치됨; 더 오래된 torch 버전은 torch._scaled_mm으로 대체)을 통해 matmul을 실행합니다. 추론 내내 가중치를 FP8로 유지하면 메모리 사용이 줄고 계산이 빨라져요. 이는 추론 전용이므로 미세 조정할 때는 꺼 두세요.

FP8 커널 경로는 다음 양자화 레이아웃을 지원합니다.

사전 양자화된 FP8 모델 로드 (Loading a pre-quantized FP8 model)

FP8 커널은 옵트인(opt-in)입니다. 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))

로드 시 역양자화 (Dequantizing at load time)

use_optimized_inference=True 없이 실행하면 모델은 일반적인 compressed-tensors 경로를 따릅니다. 가중치는 압축된 상태로 남고 compressed-tensors가 첫 번째 포워드 패스에서 압축을 해제해요. 대신 dequantize=True를 넘기면 로딩 중에 역양자화하는데, 모델을 미세 조정하거나 원래 정밀도(예: BF16)로 저장하고 싶을 때 원하는 것입니다.

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",
)

모델 체크포인트 (Model checkpoint)

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 파일은 가중치와 활성화를 정적 텐서별 전략으로 fp8 양자화하는 config 그룹(group_0)을 지정합니다. lm_head 모듈은 ignore 키에 나와 있듯이 양자화되지 않습니다.

모델 가중치를 더 자세히 보려면 모델 카드의 safetensors 뷰어를 사용해 모든 nn.Linear 모듈의 양자화된 가중치, 입력 스케일, 가중치 스케일을 볼 수 있어요.

~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가 첫 번째 포워드 패스에서 모델 전체를 압축 해제합니다. dequantize=True는 대신 로딩 중에 수행하므로, 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: BF16

use_optimized_inference=True를 쓰면 fp8 config 그룹이 덮는 레이어가 CompressedTensorsFP8Linear로 교체됩니다. 이 클래스는 행별(row-wise) matmul 커널이 기대하는 레이아웃으로 fp8 가중치와 그 스케일을 보유해요. 이 가중치들은 포워드 패스를 포함해 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

더 알아보기 (Learn more)

• compressed-tensors 저장소: 압축된 텐서 형식의 소스 코드
• llm-compressor: 모델을 compressed-tensors 형식으로 양자화하는 방법