Почему ваше квантование тормозит модель? (И как это исправить)
Вы качаете квантованную модель, запускаете инференс и ждете чуда скорости. А получаете ускорение в 1.3 раза вместо обещанных 2-3х. Знакомо? Проблема не в вас. Проблема в том, что стандартное квантование - тупое. Оно обращается со всеми весами матрицы как с равными. А они не равны.
Представьте, что у вас есть матрица весов на 10 миллиардов параметров. 99% этих весов - шум, мелкие значения около нуля. Но 1% - критически важные веса, которые определяют смысл. При равномерном квантовании в INT4 вы давите и те, и другие одинаково. Результат? Модель теряет способность к сложным рассуждениям, а скорость растет незначительно из-за overhead на деквантование.
Традиционные методы вроде GPTQ или AWQ, о которых мы писали в гайде по квантованию в vLLM, работают на уровне блоков или каналов. Per-weight идет глубже - до каждого отдельного числа.
Per-weight mixed precision: зачем делить веса на «важных» и «обычных»
Идея проста до гениальности. Вместо того чтобы квантовать всю матрицу в INT4, мы анализируем каждый вес индивидуально. Если вес важный (его абсолютное значение выше определенного порога), мы оставляем его в FP16 или BF16. Если вес неважный - переводим в INT4. В итоге 90-95% весов становятся INT4, а 5-10% критических весов остаются в полной точности.
Почему это работает быстрее? Потому что modern GPU (NVIDIA с Ampere и новее, AMD MI300, Apple Silicon) имеют отдельные tensor cores для INT4 вычислений. Когда вы мешаете INT4 и FP16 в одной операции, драйвер может параллелить загрузку. Но главное - вы сокращаете объем памяти для весов в 2.5-3 раза, а не в 4, как при полном INT4. Это значит, что модель помещается в кэш, а не торчит в медленной VRAM.
Под капотом: как найти «важные» веса
Ключевой вопрос - как определить порог. Самый простой способ - использовать статистику распределения. Берем абсолютные значения весов в слое, сортируем, берем 95-й перцентиль. Все, что выше - FP16, ниже - INT4.
Но это слишком примитивно. На практике важность веса определяется не только его величиной, но и контекстом - градиентом во время калибровки, влиянием на выходную ошибку. Метод, который мы реализуем, использует калибровочный датасет (100-200 примеров) для оценки чувствительности каждого веса.
1 Подготовка модели и калибровочных данных
Возьмем модель Llama 3.1 8B (самая свежая на 11.04.2026 в своем классе) и подготовим датасет из 128 случайных промптов. Важно: промпты должны отражать реальное использование модели.
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np
model_name = "meta-llama/Llama-3.1-8B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="auto"
)
# Генерируем калибровочные данные
calibration_prompts = [
"Explain quantum computing in simple terms.",
"Write a Python function to merge two sorted lists.",
# ... 126 других промптов
] * 2 # Удваиваем для большего разнообразия
calibration_inputs = tokenizer(
calibration_prompts,
return_tensors="pt",
padding=True,
truncation=True,
max_length=512
).to(model.device)2 Вычисление весовой важности через градиент
Включаем режим обучения, прогоняем данные, считаем градиенты для весов. Веса с большим средним градиентом - более важные.
def compute_weight_importance(model, inputs):
importance = {}
model.train()
# Прямой проход с сохранением скрытых состояний
outputs = model(**inputs, labels=inputs["input_ids"])
loss = outputs.loss
loss.backward()
# Для каждого линейного слоя собираем средний абсолютный градиент
for name, param in model.named_parameters():
if "weight" in name and param.grad is not None:
# Усредняем по всем измерениям
grad_importance = param.grad.abs().mean().item()
importance[name] = grad_importance
model.zero_grad()
return importance
importance_scores = compute_weight_importance(model, calibration_inputs)Не делайте эту ошибку: не используйте один промпт для калибровки. Модель адаптируется под конкретный контекст, и вы получите смещенную важность. 128 промптов - минимальный порог для Llama 3.1.
3 Применение per-weight mixed precision
Теперь самое интересное - создаем функцию, которая преобразует веса слоя в смешанный формат. Мы будем хранить два тензора: один для INT4 весов, другой - маску для FP16 весов.
def apply_per_weight_mixed_precision(layer_weight, importance, fp16_ratio=0.1):
"""
Применяет per-weight mixed precision к весам слоя.
Args:
layer_weight: тензор весов слоя (FP16)
importance: важность весов (тензор той же формы)
fp16_ratio: доля весов, которые останутся в FP16
Returns:
quantized_weight: квантованный тензор в формате смешанной точности
fp16_mask: маска для весов в FP16
"""
# Определяем порог для сохранения в FP16
flat_importance = importance.flatten()
threshold = torch.quantile(flat_importance, 1 - fp16_ratio)
# Создаем маску для FP16 весов
fp16_mask = importance >= threshold
# Квантуем остальные веса в INT4
# Для простоты используем симметричное квантование
weight_to_quantize = layer_weight[~fp16_mask]
# Масштаб и zero point для INT4
max_val = weight_to_quantize.abs().max()
scale = max_val / 7 # INT4 диапазон: [-7, 7]
# Квантование
quantized = torch.clamp(torch.round(weight_to_quantize / scale), -8, 7).to(torch.int8)
# Упаковываем 2 INT4 значения в один INT8 байт
# (реальная реализация сложнее, здесь упрощенно)
packed = pack_int4(quantized) # Предполагаем функцию упаковки
return {
"packed_int4": packed,
"fp16_weights": layer_weight[fp16_mask],
"fp16_mask": fp16_mask,
"scale": scale,
"original_shape": layer_weight.shape
}
# Применяем ко всем линейным слоям модели
quantized_layers = {}
for name, param in model.named_parameters():
if "weight" in name and name in importance_scores:
# Создаем тензор важности той же формы
imp_tensor = torch.full_like(param, importance_scores[name])
quantized_layers[name] = apply_per_weight_mixed_precision(
param.data,
imp_tensor
)Реальная реализация функции `pack_int4` требует битовых операций. Вот ее код:
def pack_int4(tensor_int8):
"""Упаковывает тензор INT8 (фактически INT4) в байты."""
# Сдвигаем значения из диапазона [-8,7] в [0,15]
tensor_uint4 = tensor_int8.to(torch.uint8) + 8
# Разрежаем форму
flat = tensor_uint4.flatten()
# Упаковываем два 4-битных значения в один байт
packed = torch.zeros((flat.shape[0] + 1) // 2, dtype=torch.uint8)
packed[::2] = flat[::2] # Младшие 4 бита
packed[1::2] = flat[1::2] << 4 # Старшие 4 бита
return packed