Metal

A quantização Metal executa a quantização afim em dispositivos Apple Silicon (MPS) usando kernels Metal hospedados no Hugging Face Hub (kernels-community/mlx-quantization-metal-kernels). Esses kernels originam-se do framework MLX e são compilados por meio da biblioteca kernels.

Os pesos são empacotados em tensores uint32 com escalas e vieses por grupo, e a passagem direta utiliza um kernel Metal de desquantização + matmul fundidos (affine_qmm_t). Isso mantém o uso de memória baixo ao executar a inferência inteiramente na GPU, sem idas e vindas com a CPU.

As larguras de bits suportadas são 2, 4 e 8. O tamanho do grupo é configurável (padrão 64).

Requisitos

• Mac com Apple Silicon (M1 / M2 / M3 / M4) com suporte a MPS
• O pacote kernels:

pip install kernels

Os kernels Metal são baixados automaticamente do Hub no primeiro uso — não é necessária compilação manual.

Quantize em tempo de execução

Carregue qualquer modelo e quantize-o durante o carregamento passando um MetalConfig. Todas as camadas nn.Linear elegíveis são substituídas por versões quantizadas.

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

Carregar um modelo pré-quantizado

Se um checkpoint já contém pesos quantizados (weight como uint32 empacotado, scales, qbiases), eles são carregados diretamente — sem necessidade de re-quantização.

from transformers import AutoModelForCausalLM, MetalConfig

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

Desquantizar

Em máquinas sem MPS, um checkpoint pré-quantizado é automaticamente desquantizado de volta para float, de modo que o modelo continue utilizável em CPU ou CUDA. Você também pode forçar esse comportamento explicitamente:

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 camadas

Certas camadas (por exemplo, lm_head) podem ser excluídas da quantização por meio de modules_to_not_convert:

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

Opções de configuração