Metal

Квантование Metal выполняет аффинное квантование на устройствах Apple Silicon (MPS) с использованием Metal-ядер, размещённых на Hugging Face Hub (kernels-community/mlx-quantization-metal-kernels). Эти ядра происходят из фреймворка MLX и компилируются через библиотеку kernels.

Веса упаковываются в тензоры uint32 с масштабами и смещениями для каждой группы, а прямой проход использует объединённое Metal-ядро для деквантования и матричного умножения (affine_qmm_t). Это позволяет удерживать потребление памяти на низком уровне, выполняя инференс полностью на GPU без обращения к CPU.

Поддерживаемые разрядности: 2, 4 и 8. Размер группы настраивается (по умолчанию 64).

Требования

• Mac с Apple Silicon (M1 / M2 / M3 / M4) с поддержкой MPS
• Пакет kernels:

pip install kernels

Ядра Metal автоматически загружаются с Hub при первом использовании — ручная компиляция не требуется.

Квантование на лету

Загрузите любую модель и выполните её квантование при загрузке, передав 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))

Загрузка предварительно квантованной модели

Если контрольная точка уже содержит квантованные веса (weight в виде упакованных uint32, scales, qbiases), они загружаются напрямую — повторная квантизация не требуется.

from transformers import AutoModelForCausalLM, MetalConfig

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

Декуантизация

На машинах без 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",
)

Исключить слои

Некоторые слои (например, lm_head) можно исключить из квантования с помощью modules_to_not_convert:

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

Параметры конфигурации