Activation-aware Weight Quantization

This is a different type of blog post from what I've done before, but it's part of a new thing I'm trying this summer. I'll pick a paper and try to reproduce its core result through a small experiment. I'd like to do this with no AI coding, at least for the core algorithmic parts.

For this post, I want to recreate the core idea and result from AWQ. The first step is to actually define what outcome we expect to see in our mini-experiment. Let's say we want to replicate a simplified, smaller version of Table 3 from the paper.

Let's focus only on AWQ and RTN (with the FP16 baseline). Additionally, I want to run this on my Mac -- so let's pick a more modest model scaling series: Pythia with 14M, 31M, and 70M. So, we should have a 3x3 grid of 9 perplexity numbers by the end of this post.

For a baseline, we can start with running the model in default FP16 to measure perplexity, which will be our primary metric.

The full WikiText-2 test set is longer than the model context window, so I split it into chunks of 512 tokens and average the loss across chunks. Since FP16, RTN, and AWQ all use the exact same evaluator, this is fine for comparing the methods for the purposes of this post!

Next is the naive quantization baseline: RTN, or round-to-nearest. For every nn.Linear weight matrix, I split the weights into groups of 128 values along the input-channel dimension. Each group gets its own scale and zero point. For 4-bit quantization, the integer range is $0, 1, \dots, 15$, so for a group of weights $\mathbf{w}$, I compute:

Then quantize:

Since I'm slightly lazy to handle int4 packed values + write custom kernels for the same, we can just do "fake quantization" for now. This means we can dequantize the values:

and store these in our FP16 weights. This simulates the quantization error without actually handling any low-precision kernels (and obviously we don't get a speed benefit here).

For the actual AWQ part, I found it easiest to think about the whole thing as a layer-wise function approximation problem. AWQ wants to rescale the input channels individually before quantizing. To choose how to rescale, we want to quantify how well a given rescaling config performs, then work backwards from there:

Let $\mathbf{X}$ be the input activations to some layer, and let $\mathbf{W}$ be the weight for that layer (I’ll write the layer output as $\mathbf{W}\mathbf{X}$). The scale vector $\mathbf{s}$ has one value per input channel. The trick is that we can insert $\mathbf{s}$ without changing the unquantized layer:

Without quantization, the rewrite is useless as the scale gets canceled out. But once the scaled weights are quantized, the choice of $\mathbf{s}$ starts to matter!

Different scale vectors change what the weights look like before quantization, which changes how much information survives. Of course, one activation tensor $\mathbf{X}$ is not enough. I collect a small calibration set through the model and cache the input activations for every linear layer. If those cached activations are $\{\mathbf{X}_k\}_{k=1}^n$, then the objective becomes:

There is one more detail from the paper that makes this search computationaly feasible to carry out. I am not searching over every possible value of $\mathbf{s}$. Instead, I first compute a base activation scale for each input channel:

Then I search over:

where $\alpha$ is a single scalar. In my implementation, I just try a small grid of $\alpha$ values between 0 and 1 (step = 0.05) for each layer. When $\alpha = 0$, every scale is 1, so this collapses back to normal RTN. As $\alpha$ gets larger, channels with larger activations get protected more aggressively.

In code, this maps pretty cleanly to four pieces:

1. A calibration collection function. I add forward pre-hooks to each nn.Linear layer and cache the input activations to that layer.
2. A reconstruction loss function. Given a layer name and a candidate scale vector, compute how close the quantized scaled layer is to the original layer.
3. A base scale function. For each layer, compute mean(abs(X)) per input channel.
4. A grid search function. Try different $\alpha$ values for each layer and keep the best scale.

I follow the pattern from RTN here (not doing real int4 inference). That is, the function $Q$ first quantizes via RTN, then dequantizes the same way, simulating the error. So after I find the appropriate AWQ scale for a layer, I fold it into the weight:

Then I replace the original layer weight with $\mathbf{W}_{\text{AWQ}}$, and the normal model forward pass just works as is:

Here's the final data:

This is roughly the pattern I was hoping to see. FP16 is still much better, which is expected, but AWQ consistently beats RTN across all three model sizes. The gap is especially noisy at this scale because these Pythia models are a couple OOMs smallers compared to the models in the paper. Atleast it beats RTN!

Code is all available here. There's obviously a lot more things that can be done, but this was just a quick exercise to get me warmed up for this series of reproduction posts!