


'''simple docstring'''


import logging
import re

import pytorch_quantization
import pytorch_quantization.nn as quant_nn
import torch
from pytorch_quantization import calib
from pytorch_quantization.tensor_quant import QuantDescriptor


lowerCamelCase__					=   logging.getLogger(__name__)

lowerCamelCase__					=   50  # max width of layer names
lowerCamelCase__					=   70  # max width of quantizer names






def        __lowerCAmelCase       (__lowerCAmelCase      ):
		_UpperCAmelCase						:				Dict	    =					parser.add_argument_group("quant_trainer arguments"      )
		group.add_argument("--wprec"					,       type=__lowerCAmelCase					,       default=8					,       help="weight precision"      )
		group.add_argument("--aprec"					,       type=__lowerCAmelCase					,       default=8					,       help="activation precision"      )
		group.add_argument("--quant-per-tensor"					,       action="store_true"					,       help="per tensor weight scaling"      )
		group.add_argument("--quant-disable"					,       action="store_true"					,       help="disable all quantizers"      )
		group.add_argument("--quant-disable-embeddings"					,       action="store_true"					,       help="disable all embeddings quantizers"      )
		group.add_argument("--quant-disable-keyword"					,       type=__lowerCAmelCase					,       nargs="+"					,       help="disable quantizers by keyword"      )
		group.add_argument("--quant-disable-layer-module"					,       type=__lowerCAmelCase					,       help="disable quantizers by keyword under layer."      )
		group.add_argument("--quant-enable-layer-module"					,       type=__lowerCAmelCase					,       help="enable quantizers by keyword under layer"      )
		group.add_argument("--calibrator"					,       default="max"					,       help="which quantization range calibrator to use"      )
		group.add_argument("--percentile"					,       default=__lowerCAmelCase					,       type=__lowerCAmelCase					,       help="percentile for PercentileCalibrator"      )
		group.add_argument("--fuse-qkv"					,       action="store_true"					,       help="use the same scale factor for qkv"      )
		group.add_argument("--clip-gelu"					,       metavar="N"					,       type=__lowerCAmelCase					,       help="clip gelu output maximum value to N"      )
		group.add_argument(
		    "--recalibrate-weights"					,       action="store_true"					,       help=(
		        "recalibrate weight amaxes by taking the max of the weights."
		        " amaxes will be computed with the current quantization granularity (axis)."
		    )					,       )






def        __lowerCAmelCase       (__lowerCAmelCase      ):

		if args.calibrator == "max":
				_UpperCAmelCase						:				str	    =					"max"
		elif args.calibrator == "percentile":
				if args.percentile is None:
						raise ValueError("Specify --percentile when using percentile calibrator"      )
				_UpperCAmelCase						:				List[Any]	    =					"histogram"
		elif args.calibrator == "mse":
				_UpperCAmelCase						:				Optional[int]	    =					"histogram"
		else:
				raise ValueError(F"""Invalid calibrator {args.calibrator}"""      )

		_UpperCAmelCase						:				Dict	    =					QuantDescriptor(num_bits=args.aprec					,       calib_method=__lowerCAmelCase      )
		_UpperCAmelCase						:				Optional[int]	    =					QuantDescriptor(num_bits=args.wprec					,       axis=(None if args.quant_per_tensor else (0,))      )
		quant_nn.QuantLinear.set_default_quant_desc_input(__lowerCAmelCase      )
		quant_nn.QuantLinear.set_default_quant_desc_weight(__lowerCAmelCase      )






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase=False					,       __lowerCAmelCase=False      ):
		logger.info("Configuring Model for Quantization"      )
		logger.info(F"""using quantization package {pytorch_quantization.__file__}"""      )

		if not calib:
				if args.quant_disable_embeddings:
						set_quantizer_by_name(__lowerCAmelCase					,       ["embeddings"]					,       which="weight"					,       _disabled=__lowerCAmelCase      )

				if args.quant_disable:
						set_quantizer_by_name(__lowerCAmelCase					,       [""]					,       _disabled=__lowerCAmelCase      )

				if args.quant_disable_keyword:
						set_quantizer_by_name(__lowerCAmelCase					,       args.quant_disable_keyword					,       _disabled=__lowerCAmelCase      )

				if args.quant_disable_layer_module:
						set_quantizer_by_name(__lowerCAmelCase					,       [R"layer.\d+." + args.quant_disable_layer_module]					,       _disabled=__lowerCAmelCase      )

				if args.quant_enable_layer_module:
						set_quantizer_by_name(__lowerCAmelCase					,       [R"layer.\d+." + args.quant_enable_layer_module]					,       _disabled=__lowerCAmelCase      )

				if args.recalibrate_weights:
						recalibrate_weights(__lowerCAmelCase      )

				if args.fuse_qkv:
						fuse_qkv(__lowerCAmelCase					,       __lowerCAmelCase      )

		if args.clip_gelu:
				clip_gelu(__lowerCAmelCase					,       args.clip_gelu      )

		# if args.local_rank in [-1, 0] and not calib:
		print_quant_summary(__lowerCAmelCase      )






def        __lowerCAmelCase       (__lowerCAmelCase      ):
		logger.info("Enabling Calibration"      )
		for name, module in model.named_modules():
				if name.endswith("_quantizer"      ):
						if module._calibrator is not None:
								module.disable_quant()
								module.enable_calib()
						else:
								module.disable()
						logger.info(F"""{name:80}: {module}"""      )






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase      ):
		logger.info("Loading calibrated amax"      )
		for name, module in model.named_modules():
				if name.endswith("_quantizer"      ):
						if module._calibrator is not None:
								if isinstance(module._calibrator					,       calib.MaxCalibrator      ):
										module.load_calib_amax()
								else:
										module.load_calib_amax("percentile"					,       percentile=args.percentile      )
								module.enable_quant()
								module.disable_calib()
						else:
								module.enable()
		model.cuda()
		print_quant_summary(__lowerCAmelCase      )






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase      ):

		def fusea(__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
				for mod in [qq, qk, qv]:
						if not hasattr(__lowerCAmelCase					,       "_amax"      ):
								print("          WARNING: NO AMAX BUFFER"      )
								return
				_UpperCAmelCase						:				List[Any]	    =					qq._amax.detach().item()
				_UpperCAmelCase						:				List[Any]	    =					qk._amax.detach().item()
				_UpperCAmelCase						:				List[Any]	    =					qv._amax.detach().item()

				_UpperCAmelCase						:				Optional[Any]	    =					max(__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      )
				qq._amax.fill_(__lowerCAmelCase      )
				qk._amax.fill_(__lowerCAmelCase      )
				qv._amax.fill_(__lowerCAmelCase      )
				logger.info(F"""          q={q:5.2f} k={k:5.2f} v={v:5.2f} -> {amax:5.2f}"""      )

		for name, mod in model.named_modules():
				if name.endswith(".attention.self"      ):
						logger.info(F"""FUSE_QKV: {name:{name_width}}"""      )
						fusea(mod.matmul_q_input_quantizer					,       mod.matmul_k_input_quantizer					,       mod.matmul_v_input_quantizer      )
						if args.quant_per_tensor:
								fusea(mod.query._weight_quantizer					,       mod.key._weight_quantizer					,       mod.value._weight_quantizer      )






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase      ):

		for name, mod in model.named_modules():
				if name.endswith(".output.dense"      ) and not name.endswith("attention.output.dense"      ):
						_UpperCAmelCase						:				Union[str, Any]	    =					mod._input_quantizer._amax.data.detach().item()
						mod._input_quantizer._amax.data.detach().clamp_(max=__lowerCAmelCase      )
						_UpperCAmelCase						:				Tuple	    =					mod._input_quantizer._amax.data.detach().item()
						logger.info(F"""CLIP_GELU: {name:{name_width}} amax: {amax_init:5.2f} -> {amax:5.2f}"""      )






def        __lowerCAmelCase       (__lowerCAmelCase      ):

		for name, mod in model.named_modules():
				if hasattr(__lowerCAmelCase					,       "_weight_quantizer"      ) and mod._weight_quantizer.axis is not None:
						_UpperCAmelCase						:				str	    =					mod.weight.shape[0]
						_UpperCAmelCase						:				List[Any]	    =					mod._weight_quantizer._amax.detach()
						_UpperCAmelCase						:				Tuple	    =					torch.ones(__lowerCAmelCase					,       dtype=amax.dtype					,       device=amax.device      ) * amax
						print(F"""expanding {name} {amax} -> {mod._weight_quantizer._amax}"""      )






def        __lowerCAmelCase       (__lowerCAmelCase      ):

		for name, mod in model.named_modules():
				if hasattr(__lowerCAmelCase					,       "_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)
						_UpperCAmelCase						:				Any	    =					set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis      )
						_UpperCAmelCase						:				Optional[Any]	    =					set(range(len(mod.weight.size()      )      )      ) - axis_set
						_UpperCAmelCase						:				List[str]	    =					pytorch_quantization.utils.reduce_amax(mod.weight					,       axis=__lowerCAmelCase					,       keepdims=__lowerCAmelCase      ).detach()
						logger.info(F"""RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}"""      )
						_UpperCAmelCase						:				Optional[Any]	    =					amax






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase=25					,       __lowerCAmelCase=180					,       __lowerCAmelCase=None      ):

		if ignore is None:
				_UpperCAmelCase						:				Union[str, Any]	    =					[]
		elif not isinstance(__lowerCAmelCase					,       __lowerCAmelCase      ):
				_UpperCAmelCase						:				Union[str, Any]	    =					[ignore]

		_UpperCAmelCase						:				Dict	    =					0
		for name, mod in model.named_modules():
				if not hasattr(__lowerCAmelCase					,       "weight"      ):
						continue
				_UpperCAmelCase						:				Any	    =					max(__lowerCAmelCase					,       len(__lowerCAmelCase      )      )

		for name, mod in model.named_modules():
				_UpperCAmelCase						:				Optional[Any]	    =					getattr(__lowerCAmelCase					,       "_input_quantizer"					,       __lowerCAmelCase      )
				_UpperCAmelCase						:				Union[str, Any]	    =					getattr(__lowerCAmelCase					,       "_weight_quantizer"					,       __lowerCAmelCase      )
				if not hasattr(__lowerCAmelCase					,       "weight"      ):
						continue
				if type(__lowerCAmelCase      ) in ignore:
						continue
				if [True for s in ignore if type(__lowerCAmelCase      ) is str and s in name]:
						continue
				_UpperCAmelCase						:				int	    =					F"""Act:{input_q.extra_repr()}"""
				_UpperCAmelCase						:				Optional[int]	    =					F"""Wgt:{weight_q.extra_repr()}"""
				_UpperCAmelCase						:				Dict	    =					F"""{name:{name_width}} {act_str} {wgt_str}"""
				if len(__lowerCAmelCase      ) <= line_width:
						logger.info(__lowerCAmelCase      )
				else:
						logger.info(F"""{name:{name_width}} {act_str}"""      )
						logger.info(F"""{'  ':{name_width}} {wgt_str}"""      )






def        __lowerCAmelCase       (__lowerCAmelCase      ):
		_UpperCAmelCase						:				Optional[Any]	    =					0
		for name, mod in model.named_modules():
				if isinstance(__lowerCAmelCase					,       pytorch_quantization.nn.TensorQuantizer      ):
						print(F"""{name:80} {mod}"""      )
						count += 1
		print(F"""{count} TensorQuantizers found in model"""      )






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				Tuple	    =					getattr(__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      )
		if quantizer_mod is not None:
				assert hasattr(__lowerCAmelCase					,       __lowerCAmelCase      )
				setattr(__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      )
		else:
				logger.warning(F"""{name} has no {quantizer}"""      )






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase="both"					,       **__lowerCAmelCase      ):
		_UpperCAmelCase						:				Optional[Any]	    =					F"""Warning: changing {which} quantizers of {name:{qname_width}}"""
		for k, v in kwargs.items():
				s += F""" {k}={v}"""
				if which in ["input", "both"]:
						set_quantizer(__lowerCAmelCase					,       __lowerCAmelCase					,       "_input_quantizer"					,       __lowerCAmelCase					,       __lowerCAmelCase      )
				if which in ["weight", "both"]:
						set_quantizer(__lowerCAmelCase					,       __lowerCAmelCase					,       "_weight_quantizer"					,       __lowerCAmelCase					,       __lowerCAmelCase      )
		logger.info(__lowerCAmelCase      )






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       **__lowerCAmelCase      ):

		for name, mod in model.named_modules():
				if hasattr(__lowerCAmelCase					,       "_input_quantizer"      ) or hasattr(__lowerCAmelCase					,       "_weight_quantizer"      ):
						for n in names:
								if re.search(__lowerCAmelCase					,       __lowerCAmelCase      ):
										set_quantizers(__lowerCAmelCase					,       __lowerCAmelCase					,       **__lowerCAmelCase      )
				elif name.endswith("_quantizer"      ):
						for n in names:
								if re.search(__lowerCAmelCase					,       __lowerCAmelCase      ):
										_UpperCAmelCase						:				Tuple	    =					F"""Warning: changing {name:{name_width}}"""
										for k, v in kwargs.items():
												s += F""" {k}={v}"""
												setattr(__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      )
										logger.info(__lowerCAmelCase      )







'''simple docstring'''


def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase      ):

		if density <= 0:
				raise ValueError("Impossible fluid density"      )
		if bulk_modulus <= 0:
				raise ValueError("Impossible bulk modulus"      )

		return (bulk_modulus / density) ** 0.5


if __name__ == "__main__":
						import doctest

						doctest.testmod()



