Metal

Lượng tử hóa Metal thực hiện lượng tử hóa affine trên các thiết bị Apple Silicon (MPS) bằng cách sử dụng các kernel Metal được lưu trữ trên Hugging Face Hub (kernels-community/mlx-quantization-metal-kernels). Các kernel này có nguồn gốc từ framework MLX và được biên dịch thông qua thư viện kernels.

Trọng số được đóng gói vào các tensor uint32 với các thang đo và độ lệch theo nhóm, và bước truyền thuận sử dụng kernel Metal kết hợp khử lượng tử hóa + nhân ma trận (affine_qmm_t). Điều này giúp giữ mức sử dụng bộ nhớ thấp trong khi chạy suy luận hoàn toàn trên GPU mà không có các vòng lặp qua CPU.

Các độ rộng bit được hỗ trợ là 2, 4 và 8. Kích thước nhóm có thể cấu hình (mặc định là 64).

Yêu cầu

• Mac Apple Silicon (M1 / M2 / M3 / M4) hỗ trợ MPS
• Gói kernels:

pip install kernels

Các kernel Metal được tải tự động từ Hub khi sử dụng lần đầu — không cần biên dịch thủ công.

Thực hiện lượng tử hóa ngay khi chạy

Tải bất kỳ model nào và lượng tử hóa nó trong quá trình tải bằng cách truyền vào một MetalConfig. Tất cả các lớp nn.Linear đủ điều kiện sẽ được thay thế bằng các phiên bản đã lượng tử hóa.

from transformers import AutoModelForCausalLM, AutoTokenizer, MetalConfig

quantization_config = MetalConfig(bits=4, group_size=64)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B",
    device_map="mps",
    quantization_config=quantization_config,
)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
inputs = tokenizer("Apple Silicon is", return_tensors="pt").to("mps")
output = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(output[0], skip_special_tokens=True))

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

Nếu checkpoint đã chứa các trọng số đã lượng tử hóa (weight dưới dạng uint32 đã đóng gói, scales, qbiases), chúng sẽ được tải trực tiếp — không cần lượng tử hóa lại.

from transformers import AutoModelForCausalLM, MetalConfig

model = AutoModelForCausalLM.from_pretrained(
    "your-org/model-metal-4bit",
    device_map="mps",
)

Khử lượng tử

Trên các máy không có MPS, checkpoint đã lượng tử hóa trước sẽ được tự động khử lượng tử hóa trở lại float để mô hình vẫn có thể sử dụng trên CPU hoặc CUDA. Bạn cũng có thể ép buộc hành vi này một cách rõ ràng:

from transformers import AutoModelForCausalLM, MetalConfig

config = MetalConfig(dequantize=True)
model = AutoModelForCausalLM.from_pretrained(
    "your-org/model-metal-4bit",
    quantization_config=config,
    device_map="cpu",
)

Loại trừ các lớp

Một số lớp (ví dụ: lm_head) có thể được loại trừ khỏi quá trình lượng tử hóa thông qua modules_to_not_convert:

config = MetalConfig(bits=4, group_size=64, modules_to_not_convert=["lm_head"])

Các tùy chọn cấu hình