compressed-tensors

compressed-tensors mở rộng các tệp safetensors sang các kiểu dữ liệu tensor nén để cung cấp một định dạng checkpoint thống nhất cho việc lưu trữ và nạp các định dạng lượng tử hóa khác nhau như dense, int-quantized (int8), float-quantized (fp8), và pack-quantized (int4 hoặc int8 weight-quantized được đóng gói vào int32).

compressed-tensors hỗ trợ fine-tuning với PEFT và bao gồm các tính năng sau.

• độ chính xác trọng số và kích hoạt fp8, int4, int8.
• Các chiến lược về thang lượng tử hóa và điểm không cho tensor, kênh, nhóm, khối, token.
• Lượng tử hóa kích hoạt động theo từng token (hoặc bất kỳ chiến lược tĩnh nào).
• Lượng tử hóa các module tùy ý, không chỉ giới hạn ở các module nn.Linear.
• Hỗ trợ có mục tiêu cho các module cụ thể theo tên hoặc lớp.

Cài đặt compressed-tensors từ PyPI để có bản phát hành ổn định mới nhất (khuyến nghị) hoặc cài đặt từ mã nguồn để có các tính năng mới nhất.

pip install compressed-tensors

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

Tìm kiếm bằng thẻ compressed-tensors để tìm một mô hình tương thích trên Hugging Face Hub.

Các mô hình đã được lượng tử hóa trước có thể được tải trực tiếp. Để lượng tử hóa một mô hình sang định dạng compressed-tensors, hãy xem llm-compressor. Ngoài ra, các mô hình có thể được tạo độc lập và được serialize với một cấu hình 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

Các chế độ tải

Checkpoint của compressed-tensors lưu trữ trọng số ở dạng nén (fp8, hoặc int4/int8 đã đóng gói). Cách chúng được thực thi phụ thuộc vào hai đối số của CompressedTensorsConfig.

Với cài đặt mặc định, save_pretrained() ghi các tensor nén trở lại mà không thay đổi, do đó một model đã được nạp sẽ được lưu ở định dạng compressed-tensors tương tự.

Tăng tốc kernel FP8

Truyền use_optimized_inference=True để giữ mô hình compressed-tensors nén FP8 ở định dạng FP8 và chạy các phép nhân ma trận của nó thông qua các kernel FP8 được tăng tốc bằng phần cứng (torch.nn.functional.scaled_mm, điều phối đến torch._scaled_mm_v2; các phiên bản torch cũ hơn sẽ dùng phương án dự phòng torch._scaled_mm), thay vì khử lượng hóa trọng số trở lại BF16. Việc giữ trọng số ở FP8 trong suốt quá trình suy luận giúp giảm mức sử dụng bộ nhớ và tăng tốc độ tính toán. Tính năng này chỉ dành cho suy luận, vì vậy hãy tắt nó đi khi tinh chỉnh (fine-tune).

Lộ trình kernel FP8 hỗ trợ các bố cục lượng tử hóa sau.

Tải một mô hình FP8 đã lượng tử hóa

Các kernel FP8 là tùy chọn: yêu cầu chúng bằng use_optimized_inference=True, và chúng sẽ được sử dụng khi cấu hình của mô hình chỉ định lượng hóa FP8 và có sẵn GPU được hỗ trợ.

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

Khử lượng tử hóa khi tải

Nếu không có use_optimized_inference=True, mô hình sẽ đi theo đường dẫn compressed-tensors thông thường: các trọng số được giữ ở dạng nén và compressed-tensors sẽ giải nén chúng trong lần truyền tiến đầu tiên. Truyền dequantize=True để giải nén chúng trong quá trình tải, đây là điều bạn muốn khi tinh chỉnh mô hình hoặc lưu trữ nó ở độ chính xác gốc (ví dụ: BF16). Kết quả là một mô hình dense thông thường: không có lược đồ compressed-tensors nào được gắn kèm, do đó lượng tử hóa kích hoạt (ví dụ: w8a8) cũng bị loại bỏ.

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

Điểm kiểm tra của mô hình

Mô hình compressed-tensor được định nghĩa thông qua mục cấu hình của nó. Ví dụ dưới đây được lấy từ tệp config.json của nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf.

Có rất nhiều mục để cho phép biểu đạt linh hoạt cả trong và sau quá trình nén, nhưng các mục cho việc tải và suy luận có thể được đơn giản hóa để tập trung vào một vài mục chính.

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

Tệp cấu hình xác định lượng tử hóa của một nhóm cấu hình (group_0), bao gồm lượng tử hóa trọng số và kích hoạt sang fp8 với chiến lược tĩnh trên mỗi tensor. Module lm_head không được lượng tử hóa như thể hiện trong khóa ignore.

Để xem chi tiết hơn về trọng số của mô hình, hãy sử dụng safetensors viewer trên thẻ mô hình để xem trọng số đã lượng tử hóa, thang đo đầu vào và thang đo trọng số cho tất cả các module nn.Linear.

Khi nạp một mô hình compressed-tensors với tích hợp ~quantizers.HFQuantizer, các module đích được chuyển giao cho compressed-tensors: nó gắn quantization_scheme đã được giải quyết, đặt quantization_status, đăng ký các tham số mà checkpoint lưu trữ (weight ở fp8, cùng với weight_scale và, đối với chiến lược tĩnh, input_scale) và cài đặt quá trình forward riêng của nó lên chúng. Chúng vẫn là các instance nn.Linear, vì vậy đó là những gì print hiển thị — các phiên bản compressed-tensors gần đây không còn bọc chúng trong lớp con CompressedLinear nữa. Các module được liệt kê dưới ignore, chẳng hạn như lm_head, được giữ nguyên không thay đổi.

Với dequantize=False (mặc định), các trọng số vẫn được nén sau khi quá trình tải hoàn tất, và compressed-tensors sẽ giải nén toàn bộ mô hình ở bước forward đầu tiên. dequantize=True sẽ giải nén từng trọng số trong khi tải, do đó không cần bước forward để có được các trọng số dense, và các module sẽ là các lớp dense thuần túy không có 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

Với use_optimized_inference=True, các lớp được bao phủ bởi một nhóm cấu hình fp8 sẽ được thay thế bằng CompressedTensorsFP8Linear, thành phần này lưu trữ trọng số fp8 cùng với hệ số tỷ lệ của nó theo bố trí mà nhân matmul theo hàng của nó yêu cầu. Các trọng số đó vẫn ở dạng fp8, kể cả trong các lần truyền tiến.

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