GPTQ

Contents

GPTQ#

GPTQ is a post-training quantization method based on approximate second-order information.

GPTQ with Qwix#

In Qwix, GptqCalibrationProvider collects the Hessian statistics and PtqProvider runs the model for GPTQ inference.

model = SomeLinenModel()

# Collect the Hessian statistics for GPTQ calibration.
rules = [
  gptq.GptqRule(
    module_path='Dense_0',
    weight_qtype=jnp.int8
  )
]
gptq_calibration_provider = gptq.GptqCalibrationProvider(rules)
cal_model = qwix.quantize_model(model, gptq_calibration_provider)
for batch in calibration_dataset:
    # Pass the batch to the model and capture the mutated quant_stats.
    _, new_variables = cal_model.apply(variables, batch, mutable='quant_stats')
    # Update your variables dictionary so the stats accumulate.
    variables.update(new_variables)

# Use PtqProvider for inference.
ptq_provider = qwix.PtqProvider(rules)
gptq_model = qwix.quantize_model(model, ptq_provider)
abs_variables = jax.eval_shape(gptq_model.init, jax.random.key(0), model_input)

# Apply GPTQ params to the model.
gptq_params = gptq.quantize_params(
    variables['params'], abs_variables['params'], variables['quant_stats']
)
gptq_output = gptq_model.apply({'params': gptq_params}, model_input)
Contents