


'''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


__lowercase     :						Tuple	=     logging.getLogger(__name__)

__lowercase     :						List[Any]	=     50  # max width of layer names
__lowercase     :						List[Any]	=     70  # max width of quantizer names






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       Optional[Any]       ):
					__a      :	List[str]    =  parser.add_argument_group('quant_trainer arguments'       )
					group.add_argument('--wprec'						,				type=_SCREAMING_SNAKE_CASE						,				default=8						,				help='weight precision'       )
					group.add_argument('--aprec'						,				type=_SCREAMING_SNAKE_CASE						,				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=_SCREAMING_SNAKE_CASE						,				nargs='+'						,				help='disable quantizers by keyword'       )
					group.add_argument('--quant-disable-layer-module'						,				type=_SCREAMING_SNAKE_CASE						,				help='disable quantizers by keyword under layer.'       )
					group.add_argument('--quant-enable-layer-module'						,				type=_SCREAMING_SNAKE_CASE						,				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=_SCREAMING_SNAKE_CASE						,				type=_SCREAMING_SNAKE_CASE						,				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=_SCREAMING_SNAKE_CASE						,				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							(_SCREAMING_SNAKE_CASE				:       Dict       ):

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

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






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       Any						,				_SCREAMING_SNAKE_CASE				:       str						,				_SCREAMING_SNAKE_CASE				:       Tuple=False						,				_SCREAMING_SNAKE_CASE				:       Optional[int]=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(_SCREAMING_SNAKE_CASE						,				['embeddings']						,				which='weight'						,				_disabled=_SCREAMING_SNAKE_CASE       )

										if args.quant_disable:
															set_quantizer_by_name(_SCREAMING_SNAKE_CASE						,				['']						,				_disabled=_SCREAMING_SNAKE_CASE       )

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

										if args.quant_disable_layer_module:
															set_quantizer_by_name(_SCREAMING_SNAKE_CASE						,				[r'layer.\d+.' + args.quant_disable_layer_module]						,				_disabled=_SCREAMING_SNAKE_CASE       )

										if args.quant_enable_layer_module:
															set_quantizer_by_name(_SCREAMING_SNAKE_CASE						,				[r'layer.\d+.' + args.quant_enable_layer_module]						,				_disabled=_SCREAMING_SNAKE_CASE       )

										if args.recalibrate_weights:
															recalibrate_weights(_SCREAMING_SNAKE_CASE       )

										if args.fuse_qkv:
															fuse_qkv(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       )

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

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






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       int       ):
					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							(_SCREAMING_SNAKE_CASE				:       Optional[Any]						,				_SCREAMING_SNAKE_CASE				:       Union[str, Any]       ):
					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(_SCREAMING_SNAKE_CASE       )






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       Union[str, Any]						,				_SCREAMING_SNAKE_CASE				:       int       ):

					def fusea(_SCREAMING_SNAKE_CASE				:       Union[str, Any]						,				_SCREAMING_SNAKE_CASE				:       Union[str, Any]						,				_SCREAMING_SNAKE_CASE				:       str       ):
										for mod in [qq, qk, qv]:
															if not hasattr(_SCREAMING_SNAKE_CASE						,				'_amax'       ):
																				print('          WARNING: NO AMAX BUFFER'       )
																				return
										__a      :	int    =  qq._amax.detach().item()
										__a      :	Tuple    =  qk._amax.detach().item()
										__a      :	Union[str, Any]    =  qv._amax.detach().item()

										__a      :	Optional[int]    =  max(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       )
										qq._amax.fill_(_SCREAMING_SNAKE_CASE       )
										qk._amax.fill_(_SCREAMING_SNAKE_CASE       )
										qv._amax.fill_(_SCREAMING_SNAKE_CASE       )
										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							(_SCREAMING_SNAKE_CASE				:       List[Any]						,				_SCREAMING_SNAKE_CASE				:       Union[str, Any]       ):

					for name, mod in model.named_modules():
										if name.endswith('.output.dense'       ) and not name.endswith('attention.output.dense'       ):
															__a      :	List[str]    =  mod._input_quantizer._amax.data.detach().item()
															mod._input_quantizer._amax.data.detach().clamp_(max=_SCREAMING_SNAKE_CASE       )
															__a      :	Any    =  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							(_SCREAMING_SNAKE_CASE				:       Union[str, Any]       ):

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






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       int       ):

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






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       Tuple						,				_SCREAMING_SNAKE_CASE				:       List[Any]=25						,				_SCREAMING_SNAKE_CASE				:       List[str]=180						,				_SCREAMING_SNAKE_CASE				:       Optional[int]=None       ):

					if ignore is None:
										__a      :	Optional[int]    =  []
					elif not isinstance(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       ):
										__a      :	Any    =  [ignore]

					__a      :	Optional[Any]    =  0
					for name, mod in model.named_modules():
										if not hasattr(_SCREAMING_SNAKE_CASE						,				'weight'       ):
															continue
										__a      :	Union[str, Any]    =  max(_SCREAMING_SNAKE_CASE						,				len(_SCREAMING_SNAKE_CASE       )       )

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






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       Tuple       ):
					__a      :	List[str]    =  0
					for name, mod in model.named_modules():
										if isinstance(_SCREAMING_SNAKE_CASE						,				pytorch_quantization.nn.TensorQuantizer       ):
															print(F"""{name:80} {mod}"""       )
															count += 1
					print(F"""{count} TensorQuantizers found in model"""       )






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       int						,				_SCREAMING_SNAKE_CASE				:       Any						,				_SCREAMING_SNAKE_CASE				:       Optional[Any]						,				_SCREAMING_SNAKE_CASE				:       List[Any]						,				_SCREAMING_SNAKE_CASE				:       Dict       ):
					__a      :	Tuple    =  getattr(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       )
					if quantizer_mod is not None:
										assert hasattr(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       )
										setattr(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       )
					else:
										logger.warning(F"""{name} has no {quantizer}"""       )






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       Union[str, Any]						,				_SCREAMING_SNAKE_CASE				:       Union[str, Any]						,				_SCREAMING_SNAKE_CASE				:       Any="both"						,				**_SCREAMING_SNAKE_CASE				:       int       ):
					__a      :	Union[str, 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(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				'_input_quantizer'						,				_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       )
										if which in ["weight", "both"]:
															set_quantizer(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				'_weight_quantizer'						,				_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       )
					logger.info(_SCREAMING_SNAKE_CASE       )






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       str						,				_SCREAMING_SNAKE_CASE				:       int						,				**_SCREAMING_SNAKE_CASE				:       List[Any]       ):

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







'''simple docstring'''



def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       int       ):
					return number & 1 == 0


if __name__ == "__main__":
	import doctest

	doctest.testmod()



