Buckets:

rtrm's picture
|
download
raw
3.23 kB
# Metal
Metal quantization performs affine quantization on Apple Silicon (MPS) devices using Metal kernels hosted on the Hugging Face Hub ([kernels-community/mlx-quantization-metal-kernels](https://huggingface.co/kernels-community/mlx-quantization-metal-kernels)). These kernels originate from the [MLX](https://github.com/ml-explore/mlx) framework and are compiled via the [`kernels`](https://github.com/huggingface/kernels) library.
Weights are packed into `uint32` tensors with per-group scales and biases, and the forward pass uses a fused dequantization + matmul Metal kernel (`affine_qmm_t`). This keeps memory usage low while running inference entirely on the GPU with no CPU round-trips.
Supported bit-widths are **2, 4, and 8**. Group size is configurable (default 64).
## Requirements
- Apple Silicon Mac (M1 / M2 / M3 / M4) with MPS support
- The `kernels` package:
```bash
pip install kernels
```
The Metal kernels are downloaded from the Hub automatically on first use — no manual compilation required.
## Quantize on-the-fly
Load any model and quantize it during loading by passing a [MetalConfig](/docs/transformers/pr_26617/en/main_classes/quantization#transformers.MetalConfig). All eligible `nn.Linear` layers are replaced with quantized versions.
```python
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
If a checkpoint already contains quantized weights (`weight` as packed uint32, `scales`, `qbiases`), they are loaded directly — no re-quantization needed.
```python
from transformers import AutoModelForCausalLM, MetalConfig
model = AutoModelForCausalLM.from_pretrained(
"your-org/model-metal-4bit",
device_map="mps",
)
```
## Dequantize
On machines without MPS, a pre-quantized checkpoint is automatically dequantized back to float so the model remains usable on CPU or CUDA. You can also force this behavior explicitly:
```python
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
Certain layers (e.g., `lm_head`) can be excluded from quantization via `modules_to_not_convert`:
```python
config = MetalConfig(bits=4, group_size=64, modules_to_not_convert=["lm_head"])
```
## Configuration options
| Parameter | Default | Description |
|---|---|---|
| `bits` | `4` | Bit-width for weight quantization (2, 4, or 8) |
| `group_size` | `64` | Number of elements per quantization group |
| `modules_to_not_convert` | `None` | List of module names to keep in full precision |
| `dequantize` | `False` | Force dequantization to float (for non-MPS devices) |

Xet Storage Details

Size:
3.23 kB
·
Xet hash:
ecaa0a1d5cdd3c4d28532a23adaa790e8f49e8784a29ba073818ea5f93160ccf

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.