def clip_gelu(model, maxval):
    """Clip activations generated by GELU to maxval when quantized.
    Implemented by adjusting the amax of the following input_quantizer.
    """

    for name, mod in model.named_modules():
        if name.endswith(".output.dense") and not name.endswith("attention.output.dense"):
            amax_init = mod._input_quantizer._amax.data.detach().item()
            mod._input_quantizer._amax.data.detach().clamp_(max=maxval)
            amax = mod._input_quantizer._amax.data.detach().item()
            logger.info(f"CLIP_GELU: {name:{name_width}} amax: {amax_init:5.2f} -> {amax:5.2f}")
Clip activations generated by GELU to maxval when quantized.
    Implemented by adjusting the amax of the following input_quantizer.
    
dpkd/transformers/examples/research_projects/quantization-qdqbert/quant_trainer.py