SmoothQuant

SmoothQuant is a post-training quantization method that smooths the activation outliers by offline migrating the quantization difficulty from activations to weights with a mathematically equivalent transformation.

SmoothQuant with Qwix

In Qwix, SqCalibrationProvider collects activation scale statistics and SqInferenceProvider runs the model for SmoothQuant inference.

LinenNNX

model = SomeLinenModel()

# Collect activation scale statistics for SmoothQuant calibration.
rules = [
  sq.SqRule(
    module_path='Dense_0',
    weight_qtype=jnp.int4,
    act_qtype=jnp.int4,
    alpha=0.5
  )
]
sq_calibration_provider = sq.SqCalibrationProvider(rules)
cal_model = qwix.quantize_model(model, sq_calibration_provider)
_, new_variables = cal_model.apply(variables, model_input, mutable='quant_stats')
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 SqInferenceProvider for inference and apply SmoothQuant params to the model.
sq_params = sq.quantize_params(
    variables['params'], abs_variables['params'], variables['quant_stats']
)
sq_inference_provider = sq.SqInferenceProvider(rules)
sq_model = qwix.quantize_model(model, sq_inference_provider)
sq_output = sq_model.apply({'params': sq_params}, model_input)