


import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
       import torch

       from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
       import pyanvml.pyanvml as nvml


__UpperCamelCase					:	Dict          =				logging.get_logger(__name__)





class 			lowerCAmelCase__(     snake_case__    ):






    '''simple docstring'''





    A_	:  PyTorchBenchmarkArguments
    A_	:  PretrainedConfig
    A_	:  str								=			"PyTorch"





    @property
    def        _lowerCamelCase			(						self					: Union[str, Any]						):





       '''simple docstring'''
       return torch.__version__





    def        _lowerCamelCase			(						self					: int							, __snake_case					: str							, __snake_case					: int							, __snake_case					: int						):





       '''simple docstring'''
       UpperCAmelCase_				:   Optional[Any]							= self._prepare_inference_func(__snake_case							, __snake_case							, __snake_case						)
       return self._measure_speed(_inference						)





    def        _lowerCamelCase			(						self					: Optional[int]							, __snake_case					: str							, __snake_case					: int							, __snake_case					: int						):





       '''simple docstring'''
       UpperCAmelCase_				:   List[str]							= self._prepare_inference_func(__snake_case							, __snake_case							, __snake_case						)
       return self._measure_memory(_inference						)





    def        _lowerCamelCase			(						self					: Dict							, __snake_case					: str							, __snake_case					: int							, __snake_case					: int						):





       '''simple docstring'''
       UpperCAmelCase_				:   Union[str, Any]							= self._prepare_train_func(__snake_case							, __snake_case							, __snake_case						)
       return self._measure_speed(_train						)





    def        _lowerCamelCase			(						self					: int							, __snake_case					: str							, __snake_case					: int							, __snake_case					: int						):





       '''simple docstring'''
       UpperCAmelCase_				:   Optional[int]							= self._prepare_train_func(__snake_case							, __snake_case							, __snake_case						)
       return self._measure_memory(_train						)





    def        _lowerCamelCase			(						self					: Optional[int]							, __snake_case					: str							, __snake_case					: int							, __snake_case					: int						):





       '''simple docstring'''
       UpperCAmelCase_				:   int							= self.config_dict[model_name]

       if self.args.torchscript:
          UpperCAmelCase_				:   Any							= True

       UpperCAmelCase_				:   str							= (
           hasattr(__snake_case							, '''architectures'''						)
           and isinstance(config.architectures							, __snake_case						)
           and len(config.architectures						) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
          try:
             UpperCAmelCase_				:   List[Any]							= config.architectures[0]
             UpperCAmelCase_				:   Any							= __import__('''transformers'''							, fromlist=[model_class]						)
             UpperCAmelCase_				:   int							= getattr(__snake_case							, __snake_case						)
             UpperCAmelCase_				:   List[Any]							= model_cls(__snake_case						)
          except ImportError:
             raise ImportError(
                 f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
                 ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''						)
       else:
          UpperCAmelCase_				:   Union[str, Any]							= MODEL_MAPPING[config.__class__](__snake_case						)

       model.eval()
       model.to(self.args.device						)

       # encoder-decoder has vocab size saved differently
       UpperCAmelCase_				:   Tuple							= config.vocab_size if hasattr(__snake_case							, '''vocab_size'''						) else config.encoder.vocab_size
       UpperCAmelCase_				:   Optional[Any]							= torch.randint(__snake_case							, (batch_size, sequence_length)							, dtype=torch.long							, device=self.args.device						)

       if self.args.fpaa:
          logger.info('''Running training in Mixed Precision...'''						)
          if not self.args.is_gpu:
             raise ValueError('''Mixed precision is possible only for GPU.'''						)
          # amp seems to have memory leaks so that memory usage
          # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
          model.half()

       if self.args.torchscript:
          with torch.no_grad():
             UpperCAmelCase_				:   List[Any]							= torch.jit.trace(__snake_case							, __snake_case						)
       else:
          UpperCAmelCase_				:   Dict							= model

       def encoder_decoder_forward():
          with torch.no_grad():
             UpperCAmelCase_				:   Dict							= inference_model(__snake_case							, decoder_input_ids=__snake_case						)
          return outputs

       def encoder_forward():
          with torch.no_grad():
             UpperCAmelCase_				:   Tuple							= inference_model(__snake_case						)
          return outputs

       UpperCAmelCase_				:   Optional[int]							= encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
       return _forward





    def        _lowerCamelCase			(						self					: Union[str, Any]							, __snake_case					: str							, __snake_case					: int							, __snake_case					: int						):





       '''simple docstring'''
       UpperCAmelCase_				:   Optional[int]							= self.config_dict[model_name]

       UpperCAmelCase_				:   str							= (
           hasattr(__snake_case							, '''architectures'''						)
           and isinstance(config.architectures							, __snake_case						)
           and len(config.architectures						) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
          try:
             UpperCAmelCase_				:   List[str]							= config.architectures[0]
             UpperCAmelCase_				:   Optional[int]							= __import__('''transformers'''							, fromlist=[model_class]						)
             UpperCAmelCase_				:   List[str]							= getattr(__snake_case							, __snake_case						)
             UpperCAmelCase_				:   Tuple							= model_cls(__snake_case						)
          except ImportError:
             raise ImportError(
                 f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
                 ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''						)
       else:
          UpperCAmelCase_				:   Union[str, Any]							= MODEL_WITH_LM_HEAD_MAPPING[config.__class__](__snake_case						)

       if self.args.torchscript:
          raise NotImplementedError('''Training for torchscript is currently not implemented'''						)
       else:
          UpperCAmelCase_				:   Optional[Any]							= model

       model.train()
       model.to(self.args.device						)

       # encoder-decoder has vocab size saved differently
       UpperCAmelCase_				:   Dict							= config.vocab_size if hasattr(__snake_case							, '''vocab_size'''						) else config.encoder.vocab_size
       UpperCAmelCase_				:   Dict							= torch.randint(__snake_case							, (batch_size, sequence_length)							, dtype=torch.long							, device=self.args.device						)

       if self.args.fpaa:
          logger.info('''Running training in Mixed Precision...'''						)
          if not self.args.is_gpu:
             raise ValueError('''Mixed precision is possible only for GPU.'''						)

          # amp seems to have memory leaks so that memory usage
          # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
          model.half()

       def compute_loss_and_backprob_encoder():
          UpperCAmelCase_				:   str							= train_model(__snake_case							, labels=__snake_case						)[0]
          loss.backward()
          return loss

       def compute_loss_and_backprob_encoder_decoder():
          UpperCAmelCase_				:   Union[str, Any]							= train_model(__snake_case							, decoder_input_ids=__snake_case							, labels=__snake_case						)[0]
          loss.backward()
          return loss

       UpperCAmelCase_				:   int							= (
           compute_loss_and_backprob_encoder_decoder
           if config.is_encoder_decoder
           else compute_loss_and_backprob_encoder
       )
       return _train





    def        _lowerCamelCase			(						self					: str							, __snake_case					: Any						):





       '''simple docstring'''
       try:
          if self.args.is_tpu or self.args.torchscript:
             # run additional 10 times to stabilize compilation for tpu and torchscript
             logger.info('''Do inference on TPU or torchscript. Running model 5 times to stabilize compilation'''						)
             timeit.repeat(
                 __snake_case							, repeat=1							, number=5							, )

          # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
          UpperCAmelCase_				:   Dict							= timeit.repeat(
              __snake_case							, repeat=self.args.repeat							, number=10							, )

          if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
             import torch_xla.debug.metrics as met

             self.print_fn(met.metrics_report()						)

          return min(__snake_case						) / 10.0
       except RuntimeError as e:
          self.print_fn(f'''Doesn\'t fit on GPU. {e}'''						)
          return "N/A"





    def        _lowerCamelCase			(						self					: int							, __snake_case					: Callable[[], None]						):





       '''simple docstring'''
       try:
          if self.args.trace_memory_line_by_line:
             UpperCAmelCase_				:   Dict							= start_memory_tracing('''transformers'''						)

          if self.args.is_tpu:
             # tpu
             raise NotImplementedError(
                 '''Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with'''
                 ''' `--no-memory` or `args.memory=False`'''						)
          elif self.args.is_gpu:
             if not is_pyanvml_available():
                logger.warning(
                    '''py3nvml not installed, we won\'t log GPU memory usage. '''
                    '''Install py3nvml (pip install py3nvml) to log information about GPU.'''						)
                UpperCAmelCase_				:   Optional[Any]							= '''N/A'''
             else:
                logger.info(
                    '''Measuring total GPU usage on GPU device. Make sure to not have additional processes running'''
                    ''' on the same GPU.'''						)
                # init nvml
                nvml.nvmlInit()
                func()
                UpperCAmelCase_				:   int							= nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx						)
                UpperCAmelCase_				:   Optional[Any]							= nvml.nvmlDeviceGetMemoryInfo(__snake_case						)
                UpperCAmelCase_				:   Dict							= meminfo.used
                UpperCAmelCase_				:   Tuple							= Memory(__snake_case						)
                # shutdown nvml
                nvml.nvmlShutdown()
          else:
             # cpu
             UpperCAmelCase_				:   List[Any]							= measure_peak_memory_cpu(__snake_case						)
             UpperCAmelCase_				:   str							= Memory(__snake_case						) if isinstance(__snake_case							, __snake_case						) else memory_bytes

          if self.args.trace_memory_line_by_line:
             UpperCAmelCase_				:   str							= stop_memory_tracing(__snake_case						)
          else:
             UpperCAmelCase_				:   Optional[Any]							= None

          return memory, summary
       except RuntimeError as e:
          self.print_fn(f'''Doesn\'t fit on GPU. {e}'''						)
          return "N/A", None



import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
       import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


__UpperCamelCase					:	int          =				get_logger(__name__)
def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase = None      ,							):
   # Model states
   for i, state in enumerate(__lowercase							):
      UpperCAmelCase_				:   Dict							= F'''{MODEL_NAME}.bin''' if i == 0 else F'''{MODEL_NAME}_{i}.bin'''
      UpperCAmelCase_				:   List[Any]							= os.path.join(__lowercase      ,							__lowercase							)
      save(__lowercase      ,							__lowercase							)
      logger.info(F'''Model weights saved in {output_model_file}'''							)
   # Optimizer states
   for i, opt in enumerate(__lowercase							):
      UpperCAmelCase_				:   Tuple							= opt.state_dict()
      UpperCAmelCase_				:   Any							= F'''{OPTIMIZER_NAME}.bin''' if i == 0 else F'''{OPTIMIZER_NAME}_{i}.bin'''
      UpperCAmelCase_				:   List[Any]							= os.path.join(__lowercase      ,							__lowercase							)
      save(__lowercase      ,							__lowercase							)
      logger.info(F'''Optimizer state saved in {output_optimizer_file}'''							)
   # Scheduler states
   for i, scheduler in enumerate(__lowercase							):
      UpperCAmelCase_				:   Tuple							= scheduler.state_dict()
      UpperCAmelCase_				:   Tuple							= F'''{SCHEDULER_NAME}.bin''' if i == 0 else F'''{SCHEDULER_NAME}_{i}.bin'''
      UpperCAmelCase_				:   Tuple							= os.path.join(__lowercase      ,							__lowercase							)
      save(__lowercase      ,							__lowercase							)
      logger.info(F'''Scheduler state saved in {output_scheduler_file}'''							)
   # GradScaler state
   if scaler is not None:
      UpperCAmelCase_				:   Union[str, Any]							= scaler.state_dict()
      UpperCAmelCase_				:   Union[str, Any]							= os.path.join(__lowercase      ,							__lowercase							)
      torch.save(__lowercase      ,							__lowercase							)
      logger.info(F'''Gradient scaler state saved in {output_scaler_file}'''							)
   # Random number generator states
   UpperCAmelCase_				:   Union[str, Any]							= {}
   UpperCAmelCase_				:   Tuple							= F'''{RNG_STATE_NAME}_{process_index}.pkl'''
   UpperCAmelCase_				:   List[Any]							= random.getstate()
   UpperCAmelCase_				:   Any							= np.random.get_state()
   UpperCAmelCase_				:   List[Any]							= torch.get_rng_state()
   if is_xpu_available():
      UpperCAmelCase_				:   Optional[int]							= torch.xpu.get_rng_state_all()
   else:
      UpperCAmelCase_				:   Any							= torch.cuda.get_rng_state_all()
   if is_tpu_available():
      UpperCAmelCase_				:   List[str]							= xm.get_rng_state()
   UpperCAmelCase_				:   Any							= os.path.join(__lowercase      ,							__lowercase							)
   torch.save(__lowercase      ,							__lowercase							)
   logger.info(F'''Random states saved in {output_states_file}'''							)
   return output_dir
def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase=None      ,							__lowercase=None      ,							**__lowercase      ,							):
   if map_location not in [None, "cpu", "on_device"]:
      raise TypeError(
          '''Unsupported optimizer map location passed, please choose one of `None`, `\'cpu\'`, or `\'on_device\'`'''							)
   if map_location is None:
      UpperCAmelCase_				:   Optional[Any]							= '''cpu'''
   elif map_location == "on_device":
      UpperCAmelCase_				:   Optional[int]							= PartialState().device
   # Model states
   for i, model in enumerate(__lowercase							):
      UpperCAmelCase_				:   List[Any]							= F'''{MODEL_NAME}.bin''' if i == 0 else F'''{MODEL_NAME}_{i}.bin'''
      UpperCAmelCase_				:   int							= os.path.join(__lowercase      ,							__lowercase							)
      models[i].load_state_dict(torch.load(__lowercase      ,							map_location=__lowercase							)      ,							**__lowercase							)
   logger.info('''All model weights loaded successfully'''							)

   # Optimizer states
   for i, opt in enumerate(__lowercase							):
      UpperCAmelCase_				:   List[str]							= F'''{OPTIMIZER_NAME}.bin''' if i == 0 else F'''{OPTIMIZER_NAME}_{i}.bin'''
      UpperCAmelCase_				:   List[Any]							= os.path.join(__lowercase      ,							__lowercase							)
      UpperCAmelCase_				:   Tuple							= torch.load(__lowercase      ,							map_location=__lowercase							)
      optimizers[i].load_state_dict(__lowercase							)
   logger.info('''All optimizer states loaded successfully'''							)

   # Scheduler states
   for i, scheduler in enumerate(__lowercase							):
      UpperCAmelCase_				:   Dict							= F'''{SCHEDULER_NAME}.bin''' if i == 0 else F'''{SCHEDULER_NAME}_{i}.bin'''
      UpperCAmelCase_				:   int							= os.path.join(__lowercase      ,							__lowercase							)
      scheduler.load_state_dict(torch.load(__lowercase							)							)
   logger.info('''All scheduler states loaded successfully'''							)

   # GradScaler state
   if scaler is not None:
      UpperCAmelCase_				:   List[str]							= os.path.join(__lowercase      ,							__lowercase							)
      scaler.load_state_dict(torch.load(__lowercase							)							)
      logger.info('''GradScaler state loaded successfully'''							)

   # Random states
   try:
      UpperCAmelCase_				:   Optional[Any]							= torch.load(os.path.join(__lowercase      ,							F'''{RNG_STATE_NAME}_{process_index}.pkl'''							)							)
      random.setstate(states['''random_state''']							)
      np.random.set_state(states['''numpy_random_seed''']							)
      torch.set_rng_state(states['''torch_manual_seed''']							)
      if is_xpu_available():
         torch.xpu.set_rng_state_all(states['''torch_xpu_manual_seed''']							)
      else:
         torch.cuda.set_rng_state_all(states['''torch_cuda_manual_seed''']							)
      if is_tpu_available():
         xm.set_rng_state(states['''xm_seed''']							)
      logger.info('''All random states loaded successfully'''							)
   except Exception:
      logger.info('''Could not load random states'''							)
def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase = 0							):
   UpperCAmelCase_				:   Optional[int]							= Path(__lowercase							) / F'''custom_checkpoint_{index}.pkl'''
   logger.info(F'''Saving the state of {get_pretty_name(__lowercase							)} to {save_location}'''							)
   torch.save(obj.state_dict()      ,							__lowercase							)


def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase = 0							):
   UpperCAmelCase_				:   Dict							= F'''{path}/custom_checkpoint_{index}.pkl'''
   logger.info(F'''Loading the state of {get_pretty_name(__lowercase							)} from {load_location}'''							)
   obj.load_state_dict(torch.load(__lowercase      ,							map_location='''cpu'''							)							)