



"""simple docstring"""
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


__a           =							logging.get_logger(__name__)
class 							lowerCamelCase    (     _lowerCAmelCase							):






					'''simple docstring'''





					_A						:     PyTorchBenchmarkArguments
					_A						:     PretrainedConfig
					_A						:     str		=	"PyTorch"






					@property
					def   lowerCAmelCase_							(			self:				str		)  ->      Optional[Any]:
											return torch.__version__






					def   lowerCAmelCase_							(			self:				Optional[int]		,	snake_case:				str		,	snake_case:				int		,	snake_case:				int		)  ->      float:
											snake_case_					:Dict		=     self._prepare_inference_func(snake_case		,	snake_case		,	snake_case		)
											return self._measure_speed(_inference		)






					def   lowerCAmelCase_							(			self:				Tuple		,	snake_case:				str		,	snake_case:				int		,	snake_case:				int		)  ->      [Memory, Optional[MemorySummary]]:
											snake_case_					:str		=     self._prepare_inference_func(snake_case		,	snake_case		,	snake_case		)
											return self._measure_memory(_inference		)






					def   lowerCAmelCase_							(			self:				Any		,	snake_case:				str		,	snake_case:				int		,	snake_case:				int		)  ->      float:
											snake_case_					:int		=     self._prepare_train_func(snake_case		,	snake_case		,	snake_case		)
											return self._measure_speed(_train		)






					def   lowerCAmelCase_							(			self:				Any		,	snake_case:				str		,	snake_case:				int		,	snake_case:				int		)  ->      [Memory, Optional[MemorySummary]]:
											snake_case_					:List[str]		=     self._prepare_train_func(snake_case		,	snake_case		,	snake_case		)
											return self._measure_memory(_train		)






					def   lowerCAmelCase_							(			self:				List[Any]		,	snake_case:				str		,	snake_case:				int		,	snake_case:				int		)  ->      Callable[[], None]:
											snake_case_					:Union[str, Any]		=     self.config_dict[model_name]

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

											snake_case_					:List[Any]		=     (
											    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:
																							snake_case_					:Optional[int]		=     config.architectures[0]
																							snake_case_					:Union[str, Any]		=     __import__("""transformers"""		,	fromlist=[model_class]		)
																							snake_case_					:Dict		=     getattr(snake_case		,	snake_case		)
																							snake_case_					:Union[str, 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:
																	snake_case_					:Any		=     MODEL_MAPPING[config.__class__](snake_case		)

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

											# encoder-decoder has vocab size saved differently
											snake_case_					:List[str]		=     config.vocab_size if hasattr(snake_case		,	"""vocab_size"""		) else config.encoder.vocab_size
											snake_case_					: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():
																							snake_case_					:Optional[Any]		=     torch.jit.trace(snake_case		,	snake_case		)
											else:
																	snake_case_					:Any		=     model

											def encoder_decoder_forward():
																	with torch.no_grad():
																							snake_case_					:List[str]		=     inference_model(snake_case		,	decoder_input_ids=snake_case		)
																	return outputs

											def encoder_forward():
																	with torch.no_grad():
																							snake_case_					:Optional[int]		=     inference_model(snake_case		)
																	return outputs

											snake_case_					:Dict		=     encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
											return _forward






					def   lowerCAmelCase_							(			self:				Tuple		,	snake_case:				str		,	snake_case:				int		,	snake_case:				int		)  ->      Callable[[], None]:
											snake_case_					:Tuple		=     self.config_dict[model_name]

											snake_case_					:Dict		=     (
											    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:
																							snake_case_					:Optional[int]		=     config.architectures[0]
																							snake_case_					:List[str]		=     __import__("""transformers"""		,	fromlist=[model_class]		)
																							snake_case_					:Tuple		=     getattr(snake_case		,	snake_case		)
																							snake_case_					:str		=     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:
																	snake_case_					:Dict		=     MODEL_WITH_LM_HEAD_MAPPING[config.__class__](snake_case		)

											if self.args.torchscript:
																	raise NotImplementedError("""Training for torchscript is currently not implemented"""		)
											else:
																	snake_case_					:int		=     model

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

											# encoder-decoder has vocab size saved differently
											snake_case_					:int		=     config.vocab_size if hasattr(snake_case		,	"""vocab_size"""		) else config.encoder.vocab_size
											snake_case_					:int		=     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():
																	snake_case_					:Optional[int]		=     train_model(snake_case		,	labels=snake_case		)[0]
																	loss.backward()
																	return loss

											def compute_loss_and_backprob_encoder_decoder():
																	snake_case_					:Any		=     train_model(snake_case		,	decoder_input_ids=snake_case		,	labels=snake_case		)[0]
																	loss.backward()
																	return loss

											snake_case_					:Optional[Any]		=     (
											    compute_loss_and_backprob_encoder_decoder
											    if config.is_encoder_decoder
											    else compute_loss_and_backprob_encoder
											)
											return _train






					def   lowerCAmelCase_							(			self:				List[Any]		,	snake_case:				Optional[Any]		)  ->      float:
											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
																	snake_case_					:Optional[int]		=     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		) / 1_0.0
											except RuntimeError as e:
																	self.print_fn(f"""Doesn't fit on GPU. {e}"""		)
																	return "N/A"






					def   lowerCAmelCase_							(			self:				Tuple		,	snake_case:				Callable[[], None]		)  ->      [Memory, MemorySummary]:
											try:
																	if self.args.trace_memory_line_by_line:
																							snake_case_					:Tuple		=     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."""		)
																													snake_case_					: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()
																													snake_case_					:Dict		=     nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx		)
																													snake_case_					:str		=     nvml.nvmlDeviceGetMemoryInfo(snake_case		)
																													snake_case_					:Tuple		=     meminfo.used
																													snake_case_					:Optional[int]		=     Memory(snake_case		)
																													# shutdown nvml
																													nvml.nvmlShutdown()
																	else:
																							# cpu
																							snake_case_					:int		=     measure_peak_memory_cpu(snake_case		)
																							snake_case_					:List[str]		=     Memory(snake_case		) if isinstance(snake_case		,	snake_case		) else memory_bytes

																	if self.args.trace_memory_line_by_line:
																							snake_case_					:List[str]		=     stop_memory_tracing(snake_case		)
																	else:
																							snake_case_					:Tuple		=     None

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

import importlib
import importlib.metadata
import os
import warnings
from distutils.util import strtobool
from functools import lru_cache

import torch
from packaging import version
from packaging.version import parse

from .environment import parse_flag_from_env
from .versions import compare_versions, is_torch_version


try:
    import torch_xla.core.xla_model as xm  # noqa: F401

    _lowerCAmelCase        =       True
except ImportError:
    _lowerCAmelCase        =       False


# Cache this result has it's a C FFI call which can be pretty time-consuming
_lowerCAmelCase        =       torch.distributed.is_available()


def       a__      (					a							)		->	Optional[Any]:
       # Check we're not importing a "pkg_name" directory somewhere but the actual library by trying to grab the version
       A_							:       str	      =				importlib.util.find_spec(a							) is not None
       if package_exists:
              try:
                     A_							:       str	      =				importlib.metadata.metadata(a							)
                     return True
              except importlib.metadata.PackageNotFoundError:
                     return False


def       a__      (					)		->	bool:
       return _torch_distributed_available


def       a__      (					)		->	int:
       try:
              pass
       except ImportError:
              print(
                  '''Intel(R) oneCCL Bindings for PyTorch* is required to run DDP on Intel(R) GPUs, but it is not'''
                  ''' detected. If you see \"ValueError: Invalid backend: \'ccl\'\" error, please install Intel(R) oneCCL'''
                  ''' Bindings for PyTorch*.'''							)
       return (
           importlib.util.find_spec('''torch_ccl'''							) is not None
           or importlib.util.find_spec('''oneccl_bindings_for_pytorch'''							) is not None
       )


def       a__      (					)		->	List[str]:
       return importlib.metadata.version('''oneccl_bind_pt'''							)


def       a__      (					)		->	List[str]:
       return _is_package_available('''transformer_engine'''							)


@lru_cache
def       a__      (					a=True							)		->	Union[str, Any]:
       # Due to bugs on the amp series GPUs, we disable torch-xla on them
       if torch.cuda.is_available():
              return False
       if _tpu_available and check_device:
              try:
                     # Will raise a RuntimeError if no XLA configuration is found
                     A_							:       Tuple	      =				xm.xla_device()
                     return True
              except RuntimeError:
                     return False
       return _tpu_available


def       a__      (					)		->	Dict:
       return _is_package_available('''deepspeed'''							)


def       a__      (					a=False							)		->	Optional[int]:
       if is_tpu_available():
              return not ignore_tpu
       if torch.cuda.is_available():
              return torch.cuda.is_bfaa_supported()
       if is_npu_available():
              return False
       return True


def       a__      (					)		->	Optional[Any]:
       A_							:       Optional[int]	      =				_is_package_available('''bitsandbytes'''							)
       if package_exists:
              A_							:       Any	      =				version.parse(importlib.metadata.version('''bitsandbytes'''							)							)
              return compare_versions(a  ,   '''>='''  ,   '''0.39.0'''							)
       return False


def       a__      (					)		->	Any:
       A_							:       List[str]	      =				_is_package_available('''bitsandbytes'''							)
       if package_exists:
              A_							:       Tuple	      =				version.parse(importlib.metadata.version('''bitsandbytes'''							)							)
              return compare_versions(a  ,   '''>='''  ,   '''0.37.2'''							)
       return False


def       a__      (					)		->	List[str]:
       return _is_package_available('''bitsandbytes'''							)


def       a__      (					)		->	List[str]:
       if strtobool(os.environ.get('''ACCELERATE_USE_MEGATRON_LM'''  ,   '''False'''							)							) == 1:
              A_							:       str	      =				importlib.util.find_spec('''megatron'''							) is not None
              if package_exists:
                     try:
                            A_							:       Dict	      =				parse(importlib.metadata.version('''megatron-lm'''							)							)
                            return compare_versions(a  ,   '''>='''  ,   '''2.2.0'''							)
                     except Exception as e:
                            warnings.warn(f"""Parse Megatron version failed. Exception:{e}"""							)
                            return False


def       a__      (					)		->	Tuple:
       return _is_package_available('''safetensors'''							)


def       a__      (					)		->	Tuple:
       return _is_package_available('''transformers'''							)


def       a__      (					)		->	int:
       return _is_package_available('''datasets'''							)


def       a__      (					)		->	Optional[int]:
       return _is_package_available('''aim'''							)


def       a__      (					)		->	Optional[Any]:
       return _is_package_available('''tensorboard'''							) or _is_package_available('''tensorboardX'''							)


def       a__      (					)		->	Optional[int]:
       return _is_package_available('''wandb'''							)


def       a__      (					)		->	Tuple:
       return _is_package_available('''comet_ml'''							)


def       a__      (					)		->	Tuple:
       return _is_package_available('''boto3'''							)


def       a__      (					)		->	Tuple:
       if _is_package_available('''rich'''							):
              if "ACCELERATE_DISABLE_RICH" in os.environ:
                     warnings.warn(
                         '''`ACCELERATE_DISABLE_RICH` is deprecated and will be removed in v0.22.0 and deactivated by default. Please use `ACCELERATE_ENABLE_RICH` if you wish to use `rich`.'''							)
                     return not parse_flag_from_env('''ACCELERATE_DISABLE_RICH'''  ,   a							)
              return parse_flag_from_env('''ACCELERATE_ENABLE_RICH'''  ,   a							)
       return False


def       a__      (					)		->	Any:
       return _is_package_available('''sagemaker'''							)


def       a__      (					)		->	Union[str, Any]:
       return _is_package_available('''tqdm'''							)


def       a__      (					)		->	List[str]:
       return _is_package_available('''mlflow'''							)


def       a__      (					)		->	Tuple:
       return is_torch_version('''>='''  ,   '''1.12'''							) and torch.backends.mps.is_available() and torch.backends.mps.is_built()


def       a__      (					)		->	Tuple:
       def get_major_and_minor_from_version(a							):
              return str(version.parse(a							).major							) + "." + str(version.parse(a							).minor							)

       A_							:       Tuple	      =				importlib.metadata.version('''torch'''							)
       if importlib.util.find_spec('''intel_extension_for_pytorch'''							) is None:
              return False
       A_							:       Optional[int]	      =				'''N/A'''
       try:
              A_							:       List[Any]	      =				importlib.metadata.version('''intel_extension_for_pytorch'''							)
       except importlib.metadata.PackageNotFoundError:
              return False
       A_							:       Tuple	      =				get_major_and_minor_from_version(_torch_version							)
       A_							:       int	      =				get_major_and_minor_from_version(_ipex_version							)
       if torch_major_and_minor != ipex_major_and_minor:
              warnings.warn(
                  f"""Intel Extension for PyTorch {ipex_major_and_minor} needs to work with PyTorch {ipex_major_and_minor}.*,"""
                  f""" but PyTorch {_torch_version} is found. Please switch to the matching version and run again."""							)
              return False
       return True


@lru_cache
def       a__      (					a=False							)		->	Optional[int]:
       if importlib.util.find_spec('''torch'''							) is None or importlib.util.find_spec('''torch_npu'''							) is None:
              return False

       import torch
       import torch_npu  # noqa: F401

       if check_device:
              try:
                     # Will raise a RuntimeError if no NPU is found
                     A_							:       Dict	      =				torch.npu.device_count()
                     return torch.npu.is_available()
              except RuntimeError:
                     return False
       return hasattr(a  ,   '''npu'''							) and torch.npu.is_available()


@lru_cache
def       a__      (					a=False							)		->	Optional[Any]:
       if not parse_flag_from_env('''ACCELERATE_USE_XPU'''  ,   default=a							):
              return False
       if is_ipex_available():
              import torch

              if is_torch_version('''<='''  ,   '''1.12'''							):
                     return False
       else:
              return False

       import intel_extension_for_pytorch  # noqa: F401

       if check_device:
              try:
                     # Will raise a RuntimeError if no XPU  is found
                     A_							:       Tuple	      =				torch.xpu.device_count()
                     return torch.xpu.is_available()
              except RuntimeError:
                     return False
       return hasattr(a  ,   '''xpu'''							) and torch.xpu.is_available()


