Metal

La quantification Metal effectue une quantification affine sur les appareils Apple Silicon (MPS) à l'aide de noyaux Metal hébergés sur le Hugging Face Hub (kernels-community/mlx-quantization-metal-kernels). Ces noyaux proviennent du framework MLX et sont compilés via la bibliothèque kernels.

Les poids sont regroupés dans des tenseurs uint32 avec des échelles et des biais par groupe, et le passage avant utilise un noyau Metal fusionné de déquantification + matmul (affine_qmm_t). Cela maintient une faible consommation de mémoire tout en exécutant l'inférence entièrement sur le GPU, sans aller-retour avec le CPU.

Les largeurs de bits prises en charge sont 2, 4 et 8. La taille de groupe est configurable (64 par défaut).

Prérequis

• Mac Apple Silicon (M1 / M2 / M3 / M4) avec prise en charge de MPS
• Le package kernels :

pip install kernels

Les noyaux Metal sont téléchargés automatiquement depuis le Hub lors de la première utilisation — aucune compilation manuelle n’est requise.

Quantifier à la volée

Chargez n’importe quel modèle et quantifiez-le au moment du chargement en passant un MetalConfig. Toutes les couches nn.Linear éligibles sont remplacées par des versions quantifiées.

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

Charger un modèle pré-quantifié

Si un point de contrôle contient déjà des poids quantifiés (weight en uint32 empilés, scales, qbiases), ils sont chargés directement — aucune requantification n'est nécessaire.

from transformers import AutoModelForCausalLM, MetalConfig

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

Déquantiser

Sur les machines sans MPS, un point de contrôle pré-quantifié est automatiquement déquantifié en float afin que le modèle reste utilisable sur CPU ou CUDA. Vous pouvez également forcer explicitement ce comportement :

from transformers import AutoModelForCausalLM, MetalConfig

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

Exclure des couches

Certaines couches (par exemple, lm_head) peuvent être exclues de la quantification via modules_to_not_convert :

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

Options de configuration