


'''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  :		int		     =	logging.getLogger(__name__)

_lowerCAmelCase  :		int		     =	50  # max width of layer names
_lowerCAmelCase  :		Any		     =	70  # max width of quantizer names

def 			__UpperCamelCase    (    _A    :				Optional[Any]				)    ->      List[str]:




	"""simple docstring"""





	lowerCAmelCase					:    List[Any]      							= 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    (    _A    :				Union[str, Any]				)    ->      int:




	"""simple docstring"""





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

	lowerCAmelCase					:    List[str]      							= QuantDescriptor(num_bits=args.aprec					,   calib_method=_A				)
	lowerCAmelCase					:    Optional[Any]      							= 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    (    _A    :				int					,   _A    :				Dict					,   _A    :				int=False					,   _A    :				Any=False				)    ->      Optional[int]:




	"""simple docstring"""





	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    (    _A    :				List[str]				)    ->      List[str]:




	"""simple docstring"""





	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    (    _A    :				str					,   _A    :				Optional[Any]				)    ->      str:




	"""simple docstring"""





	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    (    _A    :				Optional[int]					,   _A    :				Optional[Any]				)    ->      Any:




	"""simple docstring"""





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

		lowerCAmelCase					:    Union[str, Any]      							= 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    (    _A    :				Dict					,   _A    :				List[Any]				)    ->      int:




	"""simple docstring"""





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




	"""simple docstring"""





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

def 			__UpperCamelCase    (    _A    :				Union[str, Any]				)    ->      Optional[Any]:




	"""simple docstring"""





	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					:    Optional[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					:    List[Any]      							= 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					:    Optional[Any]      							= amax

def 			__UpperCamelCase    (    _A    :				Optional[Any]					,   _A    :				Any=25					,   _A    :				Optional[Any]=1_80					,   _A    :				str=None				)    ->      List[Any]:




	"""simple docstring"""





	if ignore is None:
		lowerCAmelCase					:    List[Any]      							= []
	elif not isinstance(_A					,   _A				):
		lowerCAmelCase					:    Dict      							= [ignore]

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

	for name, mod in model.named_modules():
		lowerCAmelCase					:    Union[str, Any]      							= getattr(_A					,   '_input_quantizer'					,   _A				)
		lowerCAmelCase					:    Any      							= 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					:    Any      							= F"Act:{input_q.extra_repr()}"
		lowerCAmelCase					:    Optional[Any]      							= F"Wgt:{weight_q.extra_repr()}"
		lowerCAmelCase					:    Tuple      							= 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    (    _A    :				Any				)    ->      List[str]:




	"""simple docstring"""





	lowerCAmelCase					:    int      							= 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    (    _A    :				Any					,   _A    :				Optional[int]					,   _A    :				int					,   _A    :				Optional[int]					,   _A    :				str				)    ->      Dict:




	"""simple docstring"""





	lowerCAmelCase					:    Optional[int]      							= 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    (    _A    :				int					,   _A    :				Tuple					,   _A    :				List[Any]="both"					,   **_A    :				int				)    ->      List[str]:




	"""simple docstring"""





	lowerCAmelCase					:    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(_A					,   _A					,   '_input_quantizer'					,   _A					,   _A				)
		if which in ["weight", "both"]:
			set_quantizer(_A					,   _A					,   '_weight_quantizer'					,   _A					,   _A				)
	logger.info(_A				)






def 			__UpperCamelCase    (    _A    :				int					,   _A    :				Optional[Any]					,   **_A    :				Any				)    ->      Dict:




	"""simple docstring"""





	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					:    Dict      							= 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'''




from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowerCAmelCase  :		str		     =	logging.get_logger(__name__)

_lowerCAmelCase  :		Tuple		     =	{
    's-JoL/Open-Llama-V1': 'https://huggingface.co/s-JoL/Open-Llama-V1/blob/main/config.json',
}



class  lowerCAmelCase				(				a							):
	_lowerCamelCase   :					Union[str, Any]			      = """open-llama"""



	def __init__(					self ,					snake_case__=10_0000 ,					snake_case__=4096 ,					snake_case__=1_1008 ,					snake_case__=32 ,					snake_case__=32 ,					snake_case__="silu" ,					snake_case__=2048 ,					snake_case__=0.0_2 ,					snake_case__=1e-6 ,					snake_case__=True ,					snake_case__=0 ,					snake_case__=1 ,					snake_case__=2 ,					snake_case__=False ,					snake_case__=True ,					snake_case__=0.1 ,					snake_case__=0.1 ,					snake_case__=True ,					snake_case__=True ,					snake_case__=None ,					**snake_case__ ,					):
		lowerCAmelCase					:    Tuple      							= vocab_size
		lowerCAmelCase					:    Optional[Any]      							= max_position_embeddings
		lowerCAmelCase					:    List[Any]      							= hidden_size
		lowerCAmelCase					:    List[Any]      							= intermediate_size
		lowerCAmelCase					:    Tuple      							= num_hidden_layers
		lowerCAmelCase					:    List[Any]      							= num_attention_heads
		lowerCAmelCase					:    List[Any]      							= hidden_act
		lowerCAmelCase					:    Union[str, Any]      							= initializer_range
		lowerCAmelCase					:    str      							= rms_norm_eps
		lowerCAmelCase					:    Optional[int]      							= use_cache
		lowerCAmelCase					:    Dict      							= kwargs.pop(
		    'use_memorry_efficient_attention' ,					snake_case__    )
		lowerCAmelCase					:    Optional[int]      							= hidden_dropout_prob
		lowerCAmelCase					:    Optional[Any]      							= attention_dropout_prob
		lowerCAmelCase					:    Union[str, Any]      							= use_stable_embedding
		lowerCAmelCase					:    Tuple      							= shared_input_output_embedding
		lowerCAmelCase					:    Tuple      							= rope_scaling
		self._rope_scaling_validation()

		super().__init__(
		    pad_token_id=snake_case__ ,					bos_token_id=snake_case__ ,					eos_token_id=snake_case__ ,					tie_word_embeddings=snake_case__ ,					**snake_case__ ,					)



	def 				lowercase     (					self    ):
		if self.rope_scaling is None:
			return

		if not isinstance(self.rope_scaling ,					snake_case__    ) or len(self.rope_scaling    ) != 2:
			raise ValueError(
			    '`rope_scaling` must be a dictionary with with two fields, `name` and `factor`, '
			    f"got {self.rope_scaling}"    )
		lowerCAmelCase					:    List[Any]      							= self.rope_scaling.get('type' ,					snake_case__    )
		lowerCAmelCase					:    List[str]      							= self.rope_scaling.get('factor' ,					snake_case__    )
		if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
			raise ValueError(
			    f"`rope_scaling`'s name field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"    )
		if rope_scaling_factor is None or not isinstance(snake_case__ ,					snake_case__    ) or rope_scaling_factor <= 1.0:
			raise ValueError(f"`rope_scaling`'s factor field must be an float > 1, got {rope_scaling_factor}"    )
