
'''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 :  List[Any]    							=       logging.getLogger(__name__)

_lowercase :  Optional[int]    							=       5_0  # max width of layer names
_lowercase :  Any    							=       7_0  # max width of quantizer names
def 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Tuple							):



						"""simple docstring"""





						lowercase_	:     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 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Union[str, Any]							):



						"""simple docstring"""





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

						lowercase_	:     Dict													=       QuantDescriptor(num_bits=args.aprec      ,      calib_method=__lowerCAmelCase							)
						lowercase_	:     List[Any]													=       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 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Optional[Any]      ,      __SCREAMING_SNAKE_CASE					: List[Any]      ,      __SCREAMING_SNAKE_CASE					: Union[str, Any]=False      ,      __SCREAMING_SNAKE_CASE					: Dict=False							):



						"""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(__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 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Optional[int]							):



						"""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 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Union[str, Any]      ,      __SCREAMING_SNAKE_CASE					: List[Any]							):



						"""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(__lowerCAmelCase							)
def 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Optional[Any]      ,      __SCREAMING_SNAKE_CASE					: str							):



						"""simple docstring"""





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

												lowercase_	:     Tuple													=       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 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Tuple      ,      __SCREAMING_SNAKE_CASE					: List[Any]							):



						"""simple docstring"""





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



						"""simple docstring"""





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



						"""simple docstring"""





						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)
																		lowercase_	:     str													=       set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis							)
																		lowercase_	:     Tuple													=       set(range(len(mod.weight.size()							)							)							) - axis_set
																		lowercase_	:     Tuple													=       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()}'''							)
																		lowercase_	:     List[Any]													=       amax
def 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Tuple      ,      __SCREAMING_SNAKE_CASE					: Optional[int]=25      ,      __SCREAMING_SNAKE_CASE					: List[Any]=180      ,      __SCREAMING_SNAKE_CASE					: Tuple=None							):



						"""simple docstring"""





						if ignore is None:
												lowercase_	:     Optional[Any]													=       []
						elif not isinstance(__lowerCAmelCase      ,      __lowerCAmelCase							):
												lowercase_	:     Optional[Any]													=       [ignore]

						lowercase_	:     List[str]													=       0
						for name, mod in model.named_modules():
												if not hasattr(__lowerCAmelCase      ,      '''weight'''							):
																		continue
												lowercase_	:     List[Any]													=       max(__lowerCAmelCase      ,      len(__lowerCAmelCase							)							)

						for name, mod in model.named_modules():
												lowercase_	:     Dict													=       getattr(__lowerCAmelCase      ,      '''_input_quantizer'''      ,      __lowerCAmelCase							)
												lowercase_	:     Dict													=       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
												lowercase_	:     Dict													=       F'''Act:{input_q.extra_repr()}'''
												lowercase_	:     Union[str, Any]													=       F'''Wgt:{weight_q.extra_repr()}'''
												lowercase_	:     Tuple													=       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 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Dict							):



						"""simple docstring"""





						lowercase_	:     List[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 					snake_case_		(				__SCREAMING_SNAKE_CASE					: List[Any]      ,      __SCREAMING_SNAKE_CASE					: Optional[int]      ,      __SCREAMING_SNAKE_CASE					: Optional[Any]      ,      __SCREAMING_SNAKE_CASE					: int      ,      __SCREAMING_SNAKE_CASE					: Dict							):



						"""simple docstring"""





						lowercase_	:     List[Any]													=       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 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Optional[Any]      ,      __SCREAMING_SNAKE_CASE					: Union[str, Any]      ,      __SCREAMING_SNAKE_CASE					: Optional[int]="both"      ,      **__SCREAMING_SNAKE_CASE					: str							):



						"""simple docstring"""





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




def 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Any      ,      __SCREAMING_SNAKE_CASE					: Tuple      ,      **__SCREAMING_SNAKE_CASE					: Optional[Any]							):



						"""simple docstring"""





						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							):
																														lowercase_	:     List[Any]													=       F'''Warning: changing {name:{name_width}}'''
																														for k, v in kwargs.items():
																																				s += F''' {k}={v}'''
																																				setattr(__lowerCAmelCase      ,      __lowerCAmelCase      ,      __lowerCAmelCase							)
																														logger.info(__lowerCAmelCase							)







import argparse

import torch

from transformers import GPTaConfig, GPTaModel, load_tf_weights_in_gpta
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()






def 							__lowercase   ( __lowerCAmelCase							:				str					,    __lowerCAmelCase							:				Any					,    __lowerCAmelCase							:				int   ):
							# Construct model
							if gpta_config_file == "":
														a__             =     GPTaConfig()
							else:
														a__             =     GPTaConfig.from_json_file(__lowerCAmelCase   )
							a__             =     GPTaModel(__lowerCAmelCase   )

							# Load weights from numpy
							load_tf_weights_in_gpta(__lowerCAmelCase					,    __lowerCAmelCase					,    __lowerCAmelCase   )

							# Save pytorch-model
							a__             =     pytorch_dump_folder_path + '/' + WEIGHTS_NAME
							a__             =     pytorch_dump_folder_path + '/' + CONFIG_NAME
							print(F'Save PyTorch model to {pytorch_weights_dump_path}'   )
							torch.save(model.state_dict()					,    __lowerCAmelCase   )
							print(F'Save configuration file to {pytorch_config_dump_path}'   )
							with open(__lowerCAmelCase					,    'w'					,    encoding='utf-8'   ) as f:
														f.write(config.to_json_string()   )


if __name__ == "__main__":
						snake_case							:       List[Any]     			=      argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '''--gpt2_checkpoint_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
						)
						parser.add_argument(
						    '''--pytorch_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
						)
						parser.add_argument(
						    '''--gpt2_config_file''',
						    default='''''',
						    type=str,
						    help=(
						        '''An optional config json file corresponding to the pre-trained OpenAI model. \n'''
						        '''This specifies the model architecture.'''
						    ),
						)
						snake_case							:       Any     			=      parser.parse_args()
						convert_gpta_checkpoint_to_pytorch(args.gpta_checkpoint_path, args.gpta_config_file, args.pytorch_dump_folder_path)




