Metal

La cuantización Metal realiza una cuantización afín en dispositivos Apple Silicon (MPS) mediante kernels de Metal alojados en el Hugging Face Hub (kernels-community/mlx-quantization-metal-kernels). Estos kernels provienen del framework MLX y se compilan a través de la biblioteca kernels.

Los pesos se empaquetan en tensores uint32 con escalas y sesgos por grupo, y la pasada hacia adelante utiliza un kernel Metal fusionado de descuantización + matmul (affine_qmm_t). Esto mantiene bajo el uso de memoria mientras se ejecuta la inferencia por completo en la GPU, sin intercambios con la CPU.

Las anchuras de bits admitidas son 2, 4 y 8. El tamaño de grupo es configurable (predeterminado: 64).

Requisitos

• Mac con Apple Silicon (M1 / M2 / M3 / M4) con soporte para MPS
• El paquete kernels:

pip install kernels

Los kernels de Metal se descargan automáticamente desde el Hub en el primer uso, sin necesidad de compilación manual.

Cuantizar sobre la marcha

Cargue cualquier modelo y cuantízelo durante la carga pasando una MetalConfig. Todas las capas nn.Linear elegibles se reemplazan por versiones cuantizadas.

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

Cargar un modelo pre-cuantizado

Si un checkpoint ya contiene pesos cuantizados (weight como uint32 empaquetados, scales, qbiases), se cargan directamente: no es necesario volver a cuantizarlos.

from transformers import AutoModelForCausalLM, MetalConfig

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

Descomprimir

En máquinas sin MPS, un punto de control pre-cuantizado se descuantiza automáticamente a float para que el modelo siga siendo utilizable en CPU o CUDA. También puedes forzar este comportamiento de forma explícita:

from transformers import AutoModelForCausalLM, MetalConfig

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

Excluir capas

Ciertas capas (p. ej., lm_head) se pueden excluir de la cuantización mediante modules_to_not_convert:

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

Opciones de configuración