Lesson 10 of 10

GPTQ's error-compensation loop

Round-to-nearest quantizes each weight independently. GPTQ notices that what we care about is the layer's output on real inputs X, and that the error from rounding one column can be partly cancelled by nudging the columns that haven't been quantized yet.

With H = 2 X^T X / N and U the upper Cholesky factor of H^-1 (so H^-1 = U^T U), process columns j = 0, 1, ... of W (shape out x in):

q = round_to_grid(W[:, j])
err = (W[:, j] - q) / U[j, j]
W[:, j+1:] -= outer(err, U[j, j+1:])
store q as column j of the result.

Step 3 is the Optimal Brain Surgeon update: the change to the remaining weights that best compensates the error, given the input statistics.

You get quantize_column(col, scale, qmax) and the Cholesky factor. Implement gptq(W, X, bits) with one scale per output row computed from the original W. The test checks that your output error on X is well below plain rounding's.

Your code

import numpy as np


def hessian_inverse_chol(X, damp=0.01):
    H = 2 * X.T @ X / len(X)
    H += damp * np.mean(np.diag(H)) * np.eye(len(H))
    return np.linalg.cholesky(np.linalg.inv(H)).T


def quantize_column(col, scale, qmax):
    return np.clip(np.round(col / scale), -qmax, qmax) * scale


def gptq(W, X, bits=3):
    """W: (out, in) weights. X: (N, in) calibration inputs. Return quantized W."""
    qmax = 2 ** (bits - 1) - 1
    U = hessian_inverse_chol(X)
    scale = np.maximum(np.abs(W).max(axis=1), 1e-12) / qmax
    raise NotImplementedError

Checking your code and saving a draft need JavaScript, which is turned off in this browser. Turn it on and reload the page to check your code. You can still edit it here, and the hints and the reference solution are below.

Pigeon's hints (2)

Hint 1 of 2

Work on W = W.copy() because step 3 rewrites the columns you haven't reached.

Hint 2 of 2

U[j, j + 1:] is a row vector; np.outer(err, U[j, j + 1:]) has the shape of W[:, j + 1:].

The reference solution

import numpy as np


def hessian_inverse_chol(X, damp=0.01):
    H = 2 * X.T @ X / len(X)
    H += damp * np.mean(np.diag(H)) * np.eye(len(H))
    return np.linalg.cholesky(np.linalg.inv(H)).T


def quantize_column(col, scale, qmax):
    return np.clip(np.round(col / scale), -qmax, qmax) * scale


def gptq(W, X, bits=3):
    qmax = 2 ** (bits - 1) - 1
    U = hessian_inverse_chol(X)
    scale = np.maximum(np.abs(W).max(axis=1), 1e-12) / qmax
    W = W.copy()
    Q = np.zeros_like(W)
    for j in range(W.shape[1]):
        q = quantize_column(W[:, j], scale, qmax)
        Q[:, j] = q
        err = (W[:, j] - q) / U[j, j]
        W[:, j + 1:] -= np.outer(err, U[j, j + 1:])
    return Q

Your code

Where this is used