Metal

Metal

Metal 양자화는 Hugging Face Hub에서 호스팅되는 Metal 커널(kernels-community/mlx-quantization-metal-kernels)을 사용해 Apple Silicon(MPS) 기기에서 아핀(affine) 양자화를 수행합니다.

출처: 문서

본문

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

가중치는 그룹별 스케일과 바이어스와 함께 uint32 텐서에 패킹되며, forward 패스는 융합된 역양자화 + matmul Metal 커널(affine_qmm_t)을 사용합니다. 이렇게 하면 메모리 사용량을 낮게 유지하면서 CPU 왕복 없이 추론을 완전히 GPU에서 실행할 수 있습니다.

지원되는 비트 폭은 2, 4, 8입니다. 그룹 크기는 설정 가능합니다(기본 64).

Requirements

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

Load a pre-quantized model

체크포인트에 이미 양자화된 가중치(weight가 패킹된 uint32, scales, qbiases)가 들어 있으면 재양자화 없이 직접 로드됩니다.

from transformers import AutoModelForCausalLM, MetalConfig

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

Dequantize

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

Exclude layers

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

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

Configuration options

Parameter Default Description
bits 4 가중치 양자화 비트 폭(2, 4, 8)
group_size 64 양자화 그룹당 요소 수
modules_to_not_convert None 전체 정밀도로 유지할 모듈 이름 목록
dequantize False float로의 역양자화 강제(비-MPS 기기용)

더 알아보기 (Learn more)