



"""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__		:				int          =	5_0  # max width of layer names
lowercase__		:				Optional[Any]          =	7_0  # max width of quantizer names

def 				UpperCamelCase_    (    lowerCAmelCase__      :				Union[str, Any]						)						->   Optional[int]:

  """simple docstring"""




  lowerCAmelCase_   :				str								= 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 				UpperCamelCase_    (    lowerCAmelCase__      :				Optional[Any]						)						->   Optional[int]:

  """simple docstring"""




  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_   :				int								= 'histogram'
  elif args.calibrator == "mse":
    lowerCAmelCase_   :				Optional[Any]								= 'histogram'
  else:
    raise ValueError(f"Invalid calibrator {args.calibrator}"						)

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

  """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 				UpperCamelCase_    (    lowerCAmelCase__      :				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_    (    lowerCAmelCase__      :				int    ,   lowerCAmelCase__      :				Tuple						)						->   Union[str, 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 				UpperCamelCase_    (    lowerCAmelCase__      :				Tuple    ,   lowerCAmelCase__      :				Optional[int]						)						->   int:

  """simple docstring"""




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

    lowerCAmelCase_   :				Dict								= 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 				UpperCamelCase_    (    lowerCAmelCase__      :				Optional[Any]    ,   lowerCAmelCase__      :				Optional[int]						)						->   Optional[int]:

  """simple docstring"""




  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=lowerCAmelCase__						)
      lowerCAmelCase_   :				Optional[int]								= 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_    (    lowerCAmelCase__      :				Optional[Any]						)						->   Optional[Any]:

  """simple docstring"""




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

def 				UpperCamelCase_    (    lowerCAmelCase__      :				List[Any]						)						->   Tuple:

  """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)
      lowerCAmelCase_   :				Any								= set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis						)
      lowerCAmelCase_   :				Optional[int]								= set(range(len(mod.weight.size()						)						)						) - axis_set
      lowerCAmelCase_   :				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()}"						)
      lowerCAmelCase_   :				Optional[int]								= amax

def 				UpperCamelCase_    (    lowerCAmelCase__      :				List[str]    ,   lowerCAmelCase__      :				Optional[int]=25    ,   lowerCAmelCase__      :				List[str]=180    ,   lowerCAmelCase__      :				Any=None						)						->   List[Any]:

  """simple docstring"""




  if ignore is None:
    lowerCAmelCase_   :				Any								= []
  elif not isinstance(lowerCAmelCase__    ,   lowerCAmelCase__						):
    lowerCAmelCase_   :				Union[str, Any]								= [ignore]

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

  for name, mod in model.named_modules():
    lowerCAmelCase_   :				int								= getattr(lowerCAmelCase__    ,   '_input_quantizer'    ,   lowerCAmelCase__						)
    lowerCAmelCase_   :				Any								= 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
    lowerCAmelCase_   :				str								= f"Act:{input_q.extra_repr()}"
    lowerCAmelCase_   :				Dict								= f"Wgt:{weight_q.extra_repr()}"
    lowerCAmelCase_   :				Optional[int]								= 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 				UpperCamelCase_    (    lowerCAmelCase__      :				Tuple						)						->   str:

  """simple docstring"""




  lowerCAmelCase_   :				Dict								= 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 				UpperCamelCase_    (    lowerCAmelCase__      :				Optional[Any]    ,   lowerCAmelCase__      :				Union[str, Any]    ,   lowerCAmelCase__      :				List[str]    ,   lowerCAmelCase__      :				List[Any]    ,   lowerCAmelCase__      :				List[Any]						)						->   Tuple:

  """simple docstring"""




  lowerCAmelCase_   :				Tuple								= 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 				UpperCamelCase_    (    lowerCAmelCase__      :				Tuple    ,   lowerCAmelCase__      :				str    ,   lowerCAmelCase__      :				Dict="both"    ,   **lowerCAmelCase__      :				Union[str, Any]						)						->   Tuple:

  """simple docstring"""




  lowerCAmelCase_   :				int								= 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 				UpperCamelCase_    (    lowerCAmelCase__      :				List[str]    ,   lowerCAmelCase__      :				Optional[int]    ,   **lowerCAmelCase__      :				str						)						->   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__						):
          lowerCAmelCase_   :				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__						)










"""simple docstring"""







import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
     import tensorflow as tf
     from tensorflow.python.eager import context
     from tensorflow.python.framework import ops

     from transformers import GradientAccumulator, create_optimizer







@require_tf
class       UpperCamelCase__   (	unittest.TestCase  ):




      """simple docstring"""




      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      str    ,			SCREAMING_SNAKE_CASE_				:      List[Any]    ,			SCREAMING_SNAKE_CASE_				:      List[str]       ):
        self.assertEqual(len(SCREAMING_SNAKE_CASE_       )    ,			len(SCREAMING_SNAKE_CASE_       )       )
        for a, b in zip(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ):
          self.assertAlmostEqual(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			delta=SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      int       ):
        lowerCAmelCase_   :				int								= GradientAccumulator()
        accumulator([tf.constant([1.0, 2.0]       )]       )
        accumulator([tf.constant([-2.0, 1.0]       )]       )
        accumulator([tf.constant([-1.0, 2.0]       )]       )
        with self.assertRaises(SCREAMING_SNAKE_CASE_       ):
          accumulator([tf.constant([1.0, 1.0]       ), tf.constant([2.0, 2.0]       )]       )
        self.assertEqual(accumulator.step    ,			3       )
        self.assertEqual(len(accumulator.gradients       )    ,			1       )
        self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()    ,			[-2.0, 5.0]    ,			tol=1E-2       )
        accumulator.reset()
        self.assertEqual(accumulator.step    ,			0       )
        self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()    ,			[0.0, 0.0]    ,			tol=1E-2       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      str       ):
        lowerCAmelCase_   :				Optional[int]								= None
        ops.enable_eager_execution_internal()
        lowerCAmelCase_   :				str								= tf.config.list_physical_devices('CPU'       )
        if len(SCREAMING_SNAKE_CASE_       ) == 1:
          tf.config.set_logical_device_configuration(
              physical_devices[0]    ,			[tf.config.LogicalDeviceConfiguration(), tf.config.LogicalDeviceConfiguration()]       )
        lowerCAmelCase_   :				Dict								= tf.config.list_logical_devices(device_type='CPU'       )
        lowerCAmelCase_   :				Optional[int]								= tf.distribute.MirroredStrategy(devices=devices[:2]       )

        with strategy.scope():
          lowerCAmelCase_   :				Union[str, Any]								= GradientAccumulator()
          lowerCAmelCase_   :				int								= tf.Variable([4.0, 3.0]       )
          lowerCAmelCase_						,lowerCAmelCase_   :				Optional[int]								= create_optimizer(5E-5    ,			1_0    ,			5       )
          lowerCAmelCase_   :				Union[str, Any]								= tf.Variable([0.0, 0.0]    ,			trainable=SCREAMING_SNAKE_CASE_       )

        def accumulate_on_replica(SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
          accumulator([gradient]       )

        def apply_on_replica():
          optimizer.apply_gradients(list(zip(accumulator.gradients    ,			[variable]       )       )       )

        @tf.function
        def accumulate(SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
          with strategy.scope():
            lowerCAmelCase_   :				Tuple								= strategy.experimental_local_results(SCREAMING_SNAKE_CASE_       )
            local_variables[0].assign(SCREAMING_SNAKE_CASE_       )
            local_variables[1].assign(SCREAMING_SNAKE_CASE_       )
            strategy.run(SCREAMING_SNAKE_CASE_    ,			args=(gradient_placeholder,)       )

        @tf.function
        def apply_grad():
          with strategy.scope():
            strategy.run(SCREAMING_SNAKE_CASE_       )

        def _check_local_values(SCREAMING_SNAKE_CASE_				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
          lowerCAmelCase_   :				List[Any]								= strategy.experimental_local_results(accumulator._gradients[0]       )
          self.assertListAlmostEqual(values[0].value()    ,			SCREAMING_SNAKE_CASE_    ,			tol=1E-2       )
          self.assertListAlmostEqual(values[1].value()    ,			SCREAMING_SNAKE_CASE_    ,			tol=1E-2       )

        accumulate([1.0, 2.0]    ,			[-1.0, 1.0]       )
        accumulate([3.0, -1.0]    ,			[-1.0, -1.0]       )
        accumulate([-2.0, 2.0]    ,			[3.0, -2.0]       )
        self.assertEqual(accumulator.step    ,			3       )
        _check_local_values([2.0, 3.0]    ,			[1.0, -2.0]       )
        apply_grad()
        self.assertListAlmostEqual(variable.value()    ,			[4.0, 3.0]    ,			tol=1E-2       )
        accumulator.reset()
        self.assertEqual(accumulator.step    ,			0       )
        _check_local_values([0.0, 0.0]    ,			[0.0, 0.0]       )





