Metal

Metal 양자화는 Apple Silicon(MPS) 장치에서 Hugging Face Hub에 호스팅된 Metal 커널(kernels-community/mlx-quantization-metal-kernels)을 사용하여 아핀 양자화를 수행합니다. 이 커널은 MLX 프레임워크에서 유래했으며 kernels 라이브러리를 통해 컴파일됩니다.

Weights는 그룹별 scale과 bias를 포함하여 uint32 텐서에 패킹되며, forward pass는 fused dequantization + matmul Metal 커널(affine_qmm_t)을 사용합니다. 이를 통해 CPU 왕복 없이 추론을 GPU에서 완전히 수행하면서도 메모리 사용량을 낮게 유지합니다.

지원되는 비트 폭은 2, 4, 8입니다. 그룹 크기는 설정할 수 있으며(기본값 64)

요구 사항

MPS를 지원하는 Apple Silicon Mac (M1 / M2 / M3 / M4)
kernels 패키지:

pip install kernels

Metal 커널은 최초 사용 시 Hub에서 자동으로 다운로드되며, 수동 컴파일이 필요하지 않습니다.

Quantize on-the-fly

MetalConfig를 전달하여 모델을 로드하는 동안 양자화를 수행할 수 있습니다. 대상이 되는 모든 nn.Linear 레이어는 양자화된 버전으로 교체됩니다.

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

사전 양자화된 모델 로드

체크포인트에 이미 양자화된 가중치(weight가 packed uint32로 저장된 경우, scales, qbiases)가 포함되어 있으면 재양자화 없이 직접 로드됩니다.

from transformers import AutoModelForCausalLM, MetalConfig

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

양자화 해제

MPS가 없는 머신에서는 사전 양자화 체크포인트가 자동으로 float으로 역양자화되어 모델이 CPU 또는 CUDA에서도 사용 가능하도록 유지됩니다. 이 동작을 명시적으로 강제할 수도 있습니다:

from transformers import AutoModelForCausalLM, MetalConfig

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

레이어 제외

일부 레이어(예: lm_head)는 modules_to_not_convert를 통해 양자화에서 제외할 수 있습니다:

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

구성 옵션

매개변수	기본값	설명
`bits`	`4`	weight 양자화에 사용할 비트 폭(2, 4 또는 8)
`group_size`	`64`	양자화 그룹당 요소 수
`modules_to_not_convert`	`None`	full 정밀도로 유지할 모듈 이름 목록
`dequantize`	`False`	float로 강제 역양자화(비-MPS 장치용)

Update on GitHub