AWQ

Contents

AWQ#

AWQ (Activation-aware Weight Quantization) is a post-training quantization method that identifies salient weight channels based on activation magnitudes and applies per-channel scaling to improve quantization accuracy.

AWQ with Qwix#

In Qwix, AwqCalibrationProvider collects the activation statistics and AwqInferenceProvider runs the model for AWQ inference.

model = SomeLinenModel()

# Collect the activation statistics for AWQ calibration.
rules = [
  awq.AwqRule(
    module_path='Dense_0',
    weight_qtype=jnp.int4,
    tile_size=64
  )
]
awq_calibration_provider = awq.AwqCalibrationProvider(rules)
cal_model = qwix.quantize_model(model, awq_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 to get the abstract quantized params tree.
ptq_provider = qwix.PtqProvider(rules)
ptq_model = qwix.quantize_model(model, ptq_provider)
abs_variables = jax.eval_shape(ptq_model.init, jax.random.key(0), model_input)

# Use AwqInferenceProvider for inference and apply AWQ params to the model.
awq_params = awq.quantize_params(
    variables['params'], abs_variables['params'], variables['quant_stats']
)
awq_inference_provider = awq.AwqInferenceProvider(rules)
awq_model = qwix.quantize_model(model, awq_inference_provider)
awq_output = awq_model.apply({'params': awq_params}, model_input)
Contents