⁣Feature-wise quantization-aware training для production inference: как сохранить метрики при внедрении 4-битных моделей на GPU с ограниченной памятью

Квантование до INT4 в продакшене режет память, но часто убивает метрики на чувствительных признаках. Обычный QAT усредняет scale на весь тензор, и редкие, но критичные фичи затираются. Feature-wise QAT решает это на уровне отдельных признаков.

Почему обычный QAT ломает метрики
Стандартное квантование обучает единый scale для всего тензора. В production-моделях, особенно в рекомендательных системах или NLP, некоторые признаки имеют высокую дисперсию или распределены неравномерно. Пример — эмбеддинги редких сущностей или временные ряды с выбросами. Один scale усредняет эти выбросы, и модель теряет важные нюансы, теряя 2-5% на задачах классификации и регрессии.

Как работает Feature-wise QAT
Вместо общего scale ты учишь отдельные параметры для каждого признака: scale и zero-point. Во время fine-tuning модель подстраивает каждый канал под искажения при 4-битном представлении. Псевдокод кастомного слоя:

class QuantLayer(nn.Module):
    def __init__(self, num_features, bits=4):
        super().__init__()
        self.scales = nn.Parameter(torch.ones(num_features))
        self.zero_points = nn.Parameter(torch.zeros(num_features))
        self.max_val = 2**(bits-1) - 1
        
    def forward(self, x):
        x_scaled = x / self.scales
        x_quant = torch.clamp(torch.round(x_scaled), -self.max_val, self.max_val)
        return x_quant * self.scales

Практический совет: используй learnable параметры, инициализированные из предварительной калибровки на репрезентативном датасете. Это ускоряет сходимость и снижает риск переобучения.

Production-метрики и trade-offs
На BERT-подобных моделях FWQAT даёт прирост F1 на 2-3% относительно обычного QAT. ResNet-50 теряет всего 0.5% против FP32, тогда как стандартный QAT режет 2-3%. Типичная ошибка — тонкий fine-tuning без репрезентативной выборки. Для стабильности нужно минимум 1% обучающих данных с сохранением оригинального распределения. На старых GPU без аппаратного INT4 (например, P40) эмуляция дорогая, но гибрид Int8+FP16 через кастомные ops улучшает блочное квантование.

Инженерные нюансы при деплое
В production FWQAT легко интегрируется: кастомные ops для TensorRT или прямой вызов learnable параметров из рантайма. Предупреждение: после fine-tuning обязательно перекалибруй scale и zero-point на валидационном датасете — иначе дрейф данных сломает метрики. Latency растёт на 5-10% из-за per-feature операций, но это окупается при памяти менее 4 ГБ на батч.

Вывод: Feature-wise QAT сохраняет метрики на production-моделях с INT4 за счёт раздельного квантования каждого признака, но требует репрезентативного fine-tuning и обязательной перекалибровки в production-пайплайне.