def recalibrate_weights(model):
    """Performs max calibration on the weights and updates amax."""

    for name, mod in model.named_modules():
        if hasattr(mod, "_weight_quantizer"):
            if not hasattr(mod.weight_quantizer, "_amax"):
                print("RECALIB: {name:{name_width}} WARNING: NO AMAX BUFFER")
                continue

            # determine which axes to reduce across
            # e.g. a 4D tensor quantized per axis 0 should reduce over (1,2,3)
            axis_set = set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis)
            reduce_axis = set(range(len(mod.weight.size()))) - axis_set
            amax = pytorch_quantization.utils.reduce_amax(mod.weight, axis=reduce_axis, keepdims=True).detach()
            logger.info(f"RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}")
            mod._weight_quantizer._amax = amax
Performs max calibration on the weights and updates amax.
dpkd/transformers/examples/research_projects/quantization-qdqbert/quant_trainer.py