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


A         =   logging.getLogger(__name__)

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





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

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

				lowerCamelCase       :      Dict          =			QuantDescriptor(num_bits=args.aprec							,				calib_method=_A)
				lowerCamelCase       :      str          =			QuantDescriptor(num_bits=args.wprec							,				axis=(None if args.quant_per_tensor else (0,)))
				quant_nn.QuantLinear.set_default_quant_desc_input(_A)
				quant_nn.QuantLinear.set_default_quant_desc_weight(_A)





def  UpperCAmelCase				(	UpperCAmelCase__				:				Any							,				UpperCAmelCase__				:				Any							,				UpperCAmelCase__				:				Any=False							,				UpperCAmelCase__				:				Optional[Any]=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(_A							,				['embeddings']							,				which='weight'							,				_disabled=_A)

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

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

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

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

								if args.recalibrate_weights:
												recalibrate_weights(_A)

								if args.fuse_qkv:
												fuse_qkv(_A							,				_A)

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

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





def  UpperCAmelCase				(	UpperCAmelCase__				:				str):
				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  UpperCAmelCase				(	UpperCAmelCase__				:				Tuple							,				UpperCAmelCase__				:				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(_A)





def  UpperCAmelCase				(	UpperCAmelCase__				:				Dict							,				UpperCAmelCase__				:				Optional[int]):

				def fusea(UpperCAmelCase__				:				Optional[Any]							,				UpperCAmelCase__				:				Optional[Any]							,				UpperCAmelCase__				:				Dict):
								for mod in [qq, qk, qv]:
												if not hasattr(_A							,				'_amax'):
																print('          WARNING: NO AMAX BUFFER')
																return
								lowerCamelCase       :      Dict          =			qq._amax.detach().item()
								lowerCamelCase       :      List[Any]          =			qk._amax.detach().item()
								lowerCamelCase       :      Optional[int]          =			qv._amax.detach().item()

								lowerCamelCase       :      Dict          =			max(_A							,				_A							,				_A)
								qq._amax.fill_(_A)
								qk._amax.fill_(_A)
								qv._amax.fill_(_A)
								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  UpperCAmelCase				(	UpperCAmelCase__				:				Dict							,				UpperCAmelCase__				:				Any):

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

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





def  UpperCAmelCase				(	UpperCAmelCase__				:				List[str]):

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





def  UpperCAmelCase				(	UpperCAmelCase__				:				Dict							,				UpperCAmelCase__				:				Tuple=25							,				UpperCAmelCase__				:				Any=1_80							,				UpperCAmelCase__				:				Optional[int]=None):

				if ignore is None:
								lowerCamelCase       :      Dict          =			[]
				elif not isinstance(_A							,				_A):
								lowerCamelCase       :      int          =			[ignore]

				lowerCamelCase       :      Optional[int]          =			0
				for name, mod in model.named_modules():
								if not hasattr(_A							,				'weight'):
												continue
								lowerCamelCase       :      Dict          =			max(_A							,				len(_A))

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





def  UpperCAmelCase				(	UpperCAmelCase__				:				Optional[Any]):
				lowerCamelCase       :      str          =			0
				for name, mod in model.named_modules():
								if isinstance(_A							,				pytorch_quantization.nn.TensorQuantizer):
												print(F'''{name:80} {mod}''')
												count += 1
				print(F'''{count} TensorQuantizers found in model''')





def  UpperCAmelCase				(	UpperCAmelCase__				:				Dict							,				UpperCAmelCase__				:				Optional[Any]							,				UpperCAmelCase__				:				Union[str, Any]							,				UpperCAmelCase__				:				Union[str, Any]							,				UpperCAmelCase__				:				Dict):
				lowerCamelCase       :      Any          =			getattr(_A							,				_A							,				_A)
				if quantizer_mod is not None:
								assert hasattr(_A							,				_A)
								setattr(_A							,				_A							,				_A)
				else:
								logger.warning(F'''{name} has no {quantizer}''')





def  UpperCAmelCase				(	UpperCAmelCase__				:				str							,				UpperCAmelCase__				:				Any							,				UpperCAmelCase__				:				Optional[int]="both"							,				**UpperCAmelCase__				:				List[str]):
				lowerCamelCase       :      Tuple          =			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(_A							,				_A							,				'_input_quantizer'							,				_A							,				_A)
								if which in ["weight", "both"]:
												set_quantizer(_A							,				_A							,				'_weight_quantizer'							,				_A							,				_A)
				logger.info(_A)





def  UpperCAmelCase				(	UpperCAmelCase__				:				Tuple							,				UpperCAmelCase__				:				List[str]							,				**UpperCAmelCase__				:				Optional[Any]):

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

'''simple docstring'''

import warnings


warnings.warn(
    'memory_utils has been reorganized to utils.memory. Import `find_executable_batchsize` from the main `__init__`: '
    '`from accelerate import find_executable_batch_size` to avoid this warning.',
    FutureWarning,
)




