Skip to content

Proposal on QAT LLM quantization

Jump to bottom
Roman V Shaposhnik edited this page Aug 26, 2024 · 2 revisions

Introduction

From Information Theory LLM is a discrete object. So considering quantization noise as channel error as long as LLM capacity as channel is higher than data rate it can be quantized almost without degradation. The problem is the weights optimization for the discrete object has exponential complexity while SGD training methods which consider LLM as continual have polynomial complexity. But continuity means an infinite accuracy and redundant bits for LLM FP weights and activations. One approach is just adding some quantization noise (PTQ), but this always distorts activation distributions and reduces quality. PTQ directly operates on LLM weights and activations without or with very small amount of training data, so it has low complexity.

Discrete-continuous contradiction can be solved by considering discrete (quantized) LLM as continuous object with noise. Similarly to Shannon’s theorem on AWGN channel capacity in this case SNR defines the channel bit-width. SGD can still be applied to continuous part of LLM giving reasoning to QAT methods.

QAT can be fixed-width and mixed-width. In later case the optimal bit-width is selected for each channel based on quality. Mixed-width QAT also covers pruning because it corresponds to zero bit quantization of pruned channels.

Unfortunately, T in QAT means training. LLM training is prohibitive both due to computations and amount and availability of training data. So for practical purposes it is necessary to develop data-free QAT with similar complexity to PTQ.

Problem statement

What we have

  1. LLM PTQ
    1. Pro: PTQ does not need training and training data.
    2. Contra: PTQ reduces NN quality as it distorts activation distributions due to quantization error.
  2. Bit-width differentiable CNN QAT with random noise injection
    1. Pro: QAT can recover quality if NN capacity is enough. Mixed-precision quantization.
    2. Contra: QAT is the same speed or slower than FP training which is prohibitable for LLM. Training data is needed.

What we do not have

Data-free fast QAT.

Proposal

Stage 0. Prerequisites

  1. Identify small LLM which can be trained locally
    1. Several transformer-based models can be chosen
      1. MobileBERT (25.3M)
      2. MobileBERT_tiny (15.1M)
    2. Although most of the papers provide benchmarks for fairly large models, such as LLAMA 7B, etc.
  2. Setup inference and training pipeline
  3. Repeat SOTA training quality (define SOTA)

Stage 1. Adapt QAT for LLM

  1. Integrate per-channel mixed-precision QAT into LLM training
    1. Different bit-widths may be considered for different parts of the transformer
      1. Embedding layers
      2. FFN (perceptron) layers
      3. Multi-head Attention layers
    2. Quantize LLM with various bit-width targets
    3. Compare quality with PTQ

Stage 2. Data-free QAT

  1. Implement random text generation using pretrained FP LLM by random sampling from generator input distribution with empty prompt
  2. Implement LLM data-free distillation using quantized network as student.
  3. Integrate distillation and QAT
  4. Compare results with stage 1 QAT

Stage 3. Forward tiled QAT

  1. Quantize LLM by 3 layer tiles one-by-one starting from input. Use FP LLM tile output as reference output with L1/L2 distance (Hinton-like forward learning).
  2. Integrate forward tiled QAT with stage 2 distillation
  3. Compere results with stage 1,2 QAT

Stage 4. Evaluation

  1. Evaluate on larger LLMs

Questions

What is LLM channel?

For CNN each channel has individual scale and per channel weights are of the same magnitude. It is similar to blocks in block quantization, as each block has individual scale. It seems natural to identify LLM channels with blocks used for quantization.

There are several types of layers with weights that transformer consist of. Most of the layers are Linear aka Fully-connected layers that composed in a different ways.

  • Embedding
  • FFN, Multi-Head attention
    • For these types of layers several techniques are exist.
      • per-row quantization - basically each row has it’s own scale factor
      • group-wise quantization - approach used in LUT-GEMM, varying amount of scale-factors per-group in specific layer channel.
      • block-wise quantization - approach used in GGUF/GGML, which is similar to group-wise quantization, but no respect to channels.
Clone this wiki locally