






from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import DistributedSampler, RandomSampler

from transformers import PreTrainedModel, Trainer, logging
from transformers.integrations import is_fairscale_available
from transformers.models.fsmt.configuration_fsmt import FSMTConfig
from transformers.optimization import (
    Adafactor,
    AdamW,
    get_constant_schedule,
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.trainer_pt_utils import get_tpu_sampler
from transformers.training_args import ParallelMode
from transformers.utils import is_torch_tpu_available


if is_fairscale_available():
					from fairscale.optim import OSS


SCREAMING_SNAKE_CASE__   :      Dict				=		logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__   :      List[str]				=		{
    """linear""": get_linear_schedule_with_warmup,
    """cosine""": get_cosine_schedule_with_warmup,
    """cosine_w_restarts""": get_cosine_with_hard_restarts_schedule_with_warmup,
    """polynomial""": get_polynomial_decay_schedule_with_warmup,
    """constant""": get_constant_schedule,
    """constant_w_warmup""": get_constant_schedule_with_warmup,
}
class     lowerCamelCase_					(   lowerCamelCase			):







				def __init__(    self							,       __lowerCAmelCase=None							,       __lowerCAmelCase=None							,       *__lowerCAmelCase							,       **__lowerCAmelCase  ):
								"""simple docstring"""
								super().__init__(*__lowerCAmelCase							,       **__lowerCAmelCase  )

								if config is None:
												assert isinstance(self.model							,       __lowerCAmelCase  ), (
												    "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
												    F''' {self.model.__class__}'''
												)
												__magic_name__		:Dict   =   self.model.config
								else:
												__magic_name__		:List[Any]   =   config

								__magic_name__		:Tuple   =   data_args
								__magic_name__		:Any   =   self.config.tgt_vocab_size if isinstance(self.config							,       __lowerCAmelCase  ) else self.config.vocab_size

								if self.args.label_smoothing != 0 or (self.data_args is not None and self.data_args.ignore_pad_token_for_loss):
												assert self.config.pad_token_id is not None, (
												    "Make sure that `config.pad_token_id` is correcly defined when ignoring `pad_token` for loss"
												    " calculation or doing label smoothing."
												)

								if self.config.pad_token_id is None and self.config.eos_token_id is not None:
												logger.warning(
												    F'''The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for'''
												    ''' padding..'''  )

								if self.args.label_smoothing == 0:
												__magic_name__		:List[Any]   =   torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id  )
								else:
												# dynamically import label_smoothed_nll_loss
												from utils import label_smoothed_nll_loss

												__magic_name__		:str   =   label_smoothed_nll_loss







				def       A						(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								if self.optimizer is None:
												__magic_name__		:int   =   ['''bias''', '''LayerNorm.weight''']
												__magic_name__		:Union[str, Any]   =   [
												    {
												        '''params''': [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay  )],
												        '''weight_decay''': self.args.weight_decay,
												    },
												    {
												        '''params''': [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay  )],
												        '''weight_decay''': 0.0,
												    },
												]
												__magic_name__		:Tuple   =   Adafactor if self.args.adafactor else AdamW
												if self.args.adafactor:
																__magic_name__		:List[Any]   =   Adafactor
																__magic_name__		:Union[str, Any]   =   {'''scale_parameter''': False, '''relative_step''': False}
												else:
																__magic_name__		:Union[str, Any]   =   AdamW
																__magic_name__		:int   =   {
																    '''betas''': (self.args.adam_betaa, self.args.adam_betaa),
																    '''eps''': self.args.adam_epsilon,
																}
												__magic_name__		:List[str]   =   self.args.learning_rate
												if self.sharded_ddp:
																__magic_name__		:List[str]   =   OSS(
																    params=__lowerCAmelCase							,       optim=__lowerCAmelCase							,       **__lowerCAmelCase							,       )
												else:
																__magic_name__		:List[str]   =   optimizer_cls(__lowerCAmelCase							,       **__lowerCAmelCase  )

								if self.lr_scheduler is None:
												__magic_name__		:Union[str, Any]   =   self._get_lr_scheduler(__lowerCAmelCase  )
								else:  # ignoring --lr_scheduler
												logger.warning('''scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.'''  )







				def       A						(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								__magic_name__		:Optional[int]   =   arg_to_scheduler[self.args.lr_scheduler]
								if self.args.lr_scheduler == "constant":
												__magic_name__		:int   =   schedule_func(self.optimizer  )
								elif self.args.lr_scheduler == "constant_w_warmup":
												__magic_name__		:List[str]   =   schedule_func(self.optimizer							,       num_warmup_steps=self.args.warmup_steps  )
								else:
												__magic_name__		:str   =   schedule_func(
												    self.optimizer							,       num_warmup_steps=self.args.warmup_steps							,       num_training_steps=__lowerCAmelCase  )
								return scheduler







				def       A						(    self  ):
								"""simple docstring"""
								if isinstance(self.train_dataset							,       torch.utils.data.IterableDataset  ):
												return None
								elif is_torch_tpu_available():
												return get_tpu_sampler(self.train_dataset  )
								else:
												if self.args.sortish_sampler:
																self.train_dataset.make_sortish_sampler(
																    self.args.per_device_train_batch_size							,       distributed=(self.args.parallel_mode == ParallelMode.DISTRIBUTED)							,       )

												return (
												    RandomSampler(self.train_dataset  )
												    if self.args.local_rank == -1
												    else DistributedSampler(self.train_dataset  )
												)







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase  ):
								"""simple docstring"""
								if self.args.label_smoothing == 0:
												if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
																# force training to ignore pad token
																__magic_name__		:Tuple   =   model(**__lowerCAmelCase							,       use_cache=__lowerCAmelCase  )[0]
																__magic_name__		:List[Any]   =   self.loss_fn(logits.view(-1							,       logits.shape[-1]  )							,       labels.view(-1  )  )
												else:
																# compute usual loss via models
																__magic_name__  ,					__magic_name__		:Dict   =   model(**__lowerCAmelCase							,       labels=__lowerCAmelCase							,       use_cache=__lowerCAmelCase  )[:2]
								else:
												# compute label smoothed loss
												__magic_name__		:Optional[int]   =   model(**__lowerCAmelCase							,       use_cache=__lowerCAmelCase  )[0]
												__magic_name__		:Dict   =   torch.nn.functional.log_softmax(__lowerCAmelCase							,       dim=-1  )
												__magic_name__  ,					__magic_name__		:int   =   self.loss_fn(__lowerCAmelCase							,       __lowerCAmelCase							,       self.args.label_smoothing							,       ignore_index=self.config.pad_token_id  )
								return loss, logits







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase  ):
								"""simple docstring"""
								__magic_name__		:str   =   inputs.pop('''labels'''  )
								__magic_name__  ,					__magic_name__		:Optional[Any]   =   self._compute_loss(__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase  )
								return loss







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase = None							,       ):
								"""simple docstring"""
								__magic_name__		:Optional[int]   =   self._prepare_inputs(__lowerCAmelCase  )

								__magic_name__		:Optional[int]   =   {
								    '''max_length''': self.data_args.val_max_target_length
								    if self.data_args is not None
								    else self.config.max_length,
								    '''num_beams''': self.data_args.eval_beams if self.data_args is not None else self.config.num_beams,
								}

								if self.args.predict_with_generate and not self.args.prediction_loss_only:
												__magic_name__		:Optional[Any]   =   self.model.generate(
												    inputs['''input_ids''']							,       attention_mask=inputs['''attention_mask''']							,       **__lowerCAmelCase							,       )
												# in case the batch is shorter than max length, the output should be padded
												if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
																__magic_name__		:Optional[Any]   =   self._pad_tensors_to_max_len(__lowerCAmelCase							,       gen_kwargs['''max_length''']  )

								__magic_name__		:int   =   inputs.pop('''labels'''  )
								with torch.no_grad():
												# compute loss on predict data
												__magic_name__  ,					__magic_name__		:Optional[int]   =   self._compute_loss(__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase  )

								__magic_name__		:Tuple   =   loss.mean().detach()
								if self.args.prediction_loss_only:
												return (loss, None, None)

								__magic_name__		:int   =   generated_tokens if self.args.predict_with_generate else logits

								if labels.shape[-1] < gen_kwargs["max_length"]:
												__magic_name__		:Any   =   self._pad_tensors_to_max_len(__lowerCAmelCase							,       gen_kwargs['''max_length''']  )

								return (loss, logits, labels)







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase  ):
								"""simple docstring"""
								# If PAD token is not defined at least EOS token has to be defined
								__magic_name__		:Optional[Any]   =   self.config.pad_token_id if self.config.pad_token_id is not None else self.config.eos_token_id

								if pad_token_id is None:
												raise ValueError(
												    '''Make sure that either `config.pad_token_id` or `config.eos_token_id` is defined if tensor has to be'''
												    F''' padded to `max_length`={max_length}'''  )

								__magic_name__		:str   =   pad_token_id * torch.ones(
								    (tensor.shape[0], max_length)							,       dtype=tensor.dtype							,       device=tensor.device  )
								__magic_name__		:Optional[Any]   =   tensor
								return padded_tensor







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


UpperCamelCase__  :						Optional[Any]					=      logging.getLogger(__name__)

UpperCamelCase__  :						Dict					=      50  # max width of layer names
UpperCamelCase__  :						Any					=      70  # max width of quantizer names







def __UpperCamelCase( _A       :					Union[str, Any]     ):
   '''simple docstring'''
   UpperCAmelCase__				:					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( _A       :					Tuple     ):
   '''simple docstring'''

   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__				:					Dict 					=	'''histogram'''
   elif args.calibrator == "mse":
      UpperCAmelCase__				:					Any 					=	'''histogram'''
   else:
      raise ValueError(F'''Invalid calibrator {args.calibrator}'''     )

   UpperCAmelCase__				:					Dict 					=	QuantDescriptor(num_bits=args.aprec	,      calib_method=_A     )
   UpperCAmelCase__				:					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( _A       :					Any	,      _A       :					Any	,      _A       :					Any=False	,      _A       :					Optional[Any]=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(_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       :					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       :					Tuple	,      _A       :					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(_A     )







def __UpperCamelCase( _A       :					Dict	,      _A       :					Optional[int]     ):
   '''simple docstring'''

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

      UpperCAmelCase__				:					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( _A       :					Dict	,      _A       :					Any     ):
   '''simple docstring'''

   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=_A     )
         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 __UpperCamelCase( _A       :					str     ):
   '''simple docstring'''

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







def __UpperCamelCase( _A       :					List[str]     ):
   '''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)
         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__				:					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()}'''     )
         UpperCAmelCase__				:					str 					=	amax







def __UpperCamelCase( _A       :					Dict	,      _A       :					Tuple=25	,      _A       :					Any=1_80	,      _A       :					Optional[int]=None     ):
   '''simple docstring'''

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

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

   for name, mod in model.named_modules():
      UpperCAmelCase__				:					str 					=	getattr(_A	,      '''_input_quantizer'''	,      _A     )
      UpperCAmelCase__				:					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
      UpperCAmelCase__				:					Dict 					=	F'''Act:{input_q.extra_repr()}'''
      UpperCAmelCase__				:					int 					=	F'''Wgt:{weight_q.extra_repr()}'''
      UpperCAmelCase__				:					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( _A       :					Optional[Any]     ):
   '''simple docstring'''
   UpperCAmelCase__				:					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( _A       :					Dict	,      _A       :					Optional[Any]	,      _A       :					Union[str, Any]	,      _A       :					Union[str, Any]	,      _A       :					Dict     ):
   '''simple docstring'''
   UpperCAmelCase__				:					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( _A       :					str	,      _A       :					Any	,      _A       :					Optional[int]="both"	,      **_A       :					List[str]     ):
   '''simple docstring'''
   UpperCAmelCase__				:					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( _A       :					Tuple	,      _A       :					List[str]	,      **_A       :					Optional[Any]     ):
   '''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     ):
               UpperCAmelCase__				:					str 					=	F'''Warning: changing {name:{name_width}}'''
               for k, v in kwargs.items():
                  s += F''' {k}={v}'''
                  setattr(_A	,      _A	,      _A     )
               logger.info(_A     )




