compressed-tensors

compressed-tensors memperluas file safetensors ke tipe data tensor terkompresi untuk menyediakan format checkpoint terpadu dalam menyimpan dan memuat berbagai format kuantisasi, seperti dense, int-quantized (int8), float-quantized (fp8), dan pack-quantized (int4 atau int8 weight-quantized yang dikemas ke dalam int32).

compressed-tensors mendukung fine-tuning dengan PEFT dan juga mencakup fitur-fitur berikut.

• presisi bobot dan aktivasi fp8, int4, int8.
• Strategi skala kuantisasi dan titik nol untuk tensor, channel, group, block, token.
• Kuantisasi aktivasi dinamis per-token (atau strategi statis apa pun).
• Kuantisasi modul apa pun, bukan hanya modul nn.Linear.
• Dukungan yang ditargetkan untuk modul tertentu berdasarkan nama atau kelas.

Instal compressed-tensors dari PyPI untuk mendapatkan rilis stabil terbaru (disarankan) atau instal dari sumber untuk mendapatkan fitur terbaru.

pip install compressed-tensors

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

Gunakan tag compressed-tensors untuk mencari model yang kompatibel di Hugging Face Hub.

Model yang sudah dikuantisasi dapat dimuat secara langsung. Untuk menguantisasi model ke format compressed-tensors, lihat llm-compressor. Sebagai alternatif, model dapat dibuat secara independen dan diserialisasi dengan konfigurasi 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

Mode pemuatan

Checkpoint compressed-tensors menyimpan bobotnya dalam bentuk terkompresi (fp8, atau int4/int8 yang dikemas). Cara eksekusinya ditentukan oleh dua argumen CompressedTensorsConfig.

Dengan pengaturan bawaan, save_pretrained() menuliskan kembali tensor terkompresi tanpa perubahan, sehingga model yang dimuat disimpan dalam format terkompresi yang sama.

Akselerasi kernel FP8

Lewati use_optimized_inference=True untuk mempertahankan model compressed-tensors FP8 dalam FP8 dan menjalankan matmuls-nya melalui kernel FP8 yang dipercepat oleh perangkat keras (torch.nn.functional.scaled_mm, yang mendispatch ke torch._scaled_mm_v2; versi torch yang lebih lama akan kembali ke torch._scaled_mm), alih-alih mendekuantisasi bobot kembali ke BF16. Mempertahankan bobot dalam FP8 sepanjang inferensi menurunkan penggunaan memori dan mempercepat komputasi. Ini hanya untuk inferensi, jadi biarkan nonaktif untuk fine-tuning.

Jalur kernel FP8 mendukung tata letak kuantisasi berikut.

Memuat model FP8 yang sudah dikuantisasi

Kernel FP8 bersifat opt-in: mintanya dengan use_optimized_inference=True, dan kernel ini akan digunakan jika konfigurasi model menentukan kuantisasi FP8 dan GPU yang didukung tersedia.

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

Dequantisasi saat pemuatan

Tanpa use_optimized_inference=True, model mengambil rute compressed-tensors reguler: bobot dibiarkan terkompresi dan compressed-tensors mendekompresinya pada forward pass pertama. Sampaikan dequantize=True untuk mendekompresinya saat loading, yang merupakan hal yang Anda inginkan untuk fine-tune model atau menyimpannya dalam presisi aslinya (mis. BF16). Hasilnya adalah model dense biasa: tidak ada skema compressed-tensors yang terlampir, sehingga activation quantization (mis. w8a8) juga dihapus.

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

Checkpoint model

Model compressed-tensor didefinisikan melalui entri konfigurasinya. Contoh berikut diambil dari file config.json pada nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf.

Ada banyak entri untuk memungkinkan ekspresi yang fleksibel baik selama maupun setelah kompresi, tetapi entri untuk pemuatan dan inferensi dapat disederhanakan untuk berfokus hanya pada beberapa entri kunci.

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

File konfigurasi menentukan kuantisasi dari grup konfigurasi (group_0), yang mencakup kuantisasi bobot dan aktivasi ke fp8 dengan strategi statis per tensor. Modul lm_head tidak dikuantisasi, seperti yang ditunjukkan pada kunci ignore.

Untuk melihat detail lebih lanjut tentang bobot model, gunakan safetensors viewer di kartu model untuk melihat bobot terkuantisasi, skala input, dan skala bobot untuk semua modul nn.Linear.

Saat memuat model compressed-tensors dengan integrasi ~quantizers.HFQuantizer, modul yang ditargetkan diserahkan ke compressed-tensors: integrasi ini menempelkan quantization_scheme yang telah dipecahkan, mengatur quantization_status, mendaftarkan parameter yang disimpan oleh checkpoint (weight dalam fp8, serta weight_scale dan, untuk strategi statis, input_scale), dan memasang jalur forward-nya sendiri di atas parameter tersebut. Modul-modul ini tetap berupa instance nn.Linear, sehingga itulah yang ditampilkan oleh print — versi compressed-tensors terbaru tidak lagi membungkusnya dalam subclass CompressedLinear. Modul yang tercantum di bawah ignore, seperti lm_head, dibiarkan apa adanya.

Dengan dequantize=False (default), bobot masih terkompresi setelah proses pemuatan selesai, dan compressed-tensors akan mendekompresi seluruh model pada langkah forward pertama. dequantize=True mendekompresi setiap bobot saat dimuat, sehingga tidak diperlukan langkah forward untuk mendapatkan bobot padat, dan modulnya menjadi lapisan padat biasa tanpa quantization_scheme.

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

Dengan use_optimized_inference=True, lapisan yang tercakup oleh grup konfigurasi fp8 diganti dengan CompressedTensorsFP8Linear, yang menyimpan bobot fp8 dan skalanya dalam tata letak yang diharapkan oleh kernel matmul baris-nya. Bobot tersebut tetap dalam fp8, termasuk pada jalur maju (forward pass).

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