

import importlib.util
import weakref
from functools import partialmethod

from .dependency_versions_check import dep_version_check
from .utils import is_accelerate_available, is_torch_available, logging


if is_torch_available():
							import torch

UpperCamelCase__					: Any								=      logging.get_logger(__name__)


def    _UpperCAmelCase    (					):





				"""simple docstring"""


				return importlib.util.find_spec('deepspeed'    ) is not None


if is_accelerate_available() and is_deepspeed_available():
							from accelerate.utils.deepspeed import HfDeepSpeedConfig as DeepSpeedConfig
else:
							# Inherits from a dummy `object` if accelerate is not available, so that python succeeds to import this file.
							# Deepspeed glue code will never inherit this dummy object as it checks if accelerate is available.
							from builtins import object as DeepSpeedConfig



class  __snake_case					(      UpperCamelCase_    ):

		def __init__(   self  ,						_A):
						# set global weakref object
						set_hf_deepspeed_config(self)
						dep_version_check('accelerate')
						dep_version_check('deepspeed')
						super().__init__(__SCREAMING_SNAKE_CASE)






class  __snake_case					(      UpperCamelCase_    ):

		def __init__(   self  ,						_A):
						super().__init__(__SCREAMING_SNAKE_CASE)
						SCREAMING_SNAKE_CASE_													=  None
						SCREAMING_SNAKE_CASE_													=  []

		def   lowerCAmelCase__       (   self):
						if self._dtype is None:
										raise ValueError('trainer_config_process() wasn\'t called yet to tell dtype')
						return self._dtype

		def   lowerCAmelCase__       (   self  ,						_A):
						SCREAMING_SNAKE_CASE_													=  self.get_value(__SCREAMING_SNAKE_CASE)
						if val is None:
										return False
						else:
										return val == "auto"

		def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A=None  ,						_A=True):
						SCREAMING_SNAKE_CASE_							,     SCREAMING_SNAKE_CASE_													=  self.find_config_node(__SCREAMING_SNAKE_CASE)
						if config is None:
										return

						if config.get(__SCREAMING_SNAKE_CASE) == "auto":
										SCREAMING_SNAKE_CASE_													=  hf_val
										return

						if not must_match:
										return

						SCREAMING_SNAKE_CASE_													=  config.get(__SCREAMING_SNAKE_CASE)
						if ds_val is not None and ds_val != hf_val:
										self.mismatches.append(f"""- ds {ds_key_long}={ds_val} vs hf {hf_key}={hf_val}""")

		__lowerCAmelCase       :		int  	=					partialmethod(UpperCamelCase_       ,	must_match=UpperCamelCase_    )

		def   lowerCAmelCase__       (   self  ,						_A):
						SCREAMING_SNAKE_CASE_													=  args.world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps
						self.fill_match(
						    'train_micro_batch_size_per_gpu'  ,						args.per_device_train_batch_size  ,						'per_device_train_batch_size')
						self.fill_match('gradient_accumulation_steps'  ,						args.gradient_accumulation_steps  ,						'gradient_accumulation_steps')
						self.fill_match('train_batch_size'  ,						__SCREAMING_SNAKE_CASE  ,						'train_batch_size (calculated)')
						self.fill_match('gradient_clipping'  ,						args.max_grad_norm  ,						'max_grad_norm')

						self.fill_match('optimizer.params.lr'  ,						args.learning_rate  ,						'learning_rate')
						self.fill_match('optimizer.params.betas'  ,						[args.adam_betaa, args.adam_betaa]  ,						'adam_beta1+adam_beta2')
						self.fill_match('optimizer.params.eps'  ,						args.adam_epsilon  ,						'adam_epsilon')
						self.fill_match('optimizer.params.weight_decay'  ,						args.weight_decay  ,						'weight_decay')

						self.fill_only('scheduler.params.warmup_min_lr'  ,						0)  # not a trainer arg
						self.fill_match('scheduler.params.warmup_max_lr'  ,						args.learning_rate  ,						'learning_rate')
						# total_num_steps - will get set in trainer_config_finalize

						# fp16
						if args.fpaa or args.fpaa_full_eval:
										SCREAMING_SNAKE_CASE_													=  'apex' if args.fpaa_backend == 'apex' else 'amp'
						else:
										SCREAMING_SNAKE_CASE_													=  None

						if args.save_on_each_node:
										# deepspeed uses shared storage by default. Let's override this setting if save_on_each_node == True
										SCREAMING_SNAKE_CASE_													=  self.config.get('checkpoint'  ,						{})
										SCREAMING_SNAKE_CASE_													=  args.save_on_each_node

						# amp: similar to the pytorch native amp - it has a bunch of optional params but we won't set
						# any here unless the user did the work
						self.fill_match(
						    'fp16.enabled'  ,						((args.fpaa or args.fpaa_full_eval) and fpaa_backend == 'amp')  ,						'fp16|fp16_full_eval+fp16_backend(amp)'  ,						)

						# apex: delegates amp work to apex (which needs to be available), but it cannot be used with any
						# ZeRO features
						self.fill_match('amp.enabled'  ,						fpaa_backend == 'apex'  ,						'fp16+fp16_backend(apex)')
						self.fill_match('amp.opt_level'  ,						args.fpaa_opt_level  ,						'fp16_opt_level')

						self.fill_match('bf16.enabled'  ,						(args.bfaa or args.bfaa_full_eval)  ,						'bf16|bf16_full_eval')

						# deepspeed's default mode is fp16 unless there is a config that says differently
						if self.is_true('bf16.enabled'):
										SCREAMING_SNAKE_CASE_													=  torch.bfloataa
						elif self.is_false('fp16.enabled'):
										SCREAMING_SNAKE_CASE_													=  torch.floataa
						else:
										SCREAMING_SNAKE_CASE_													=  torch.floataa


		def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A):
						SCREAMING_SNAKE_CASE_													=  [
						    'zero_optimization.reduce_bucket_size',
						    'zero_optimization.stage3_prefetch_bucket_size',
						    'zero_optimization.stage3_param_persistence_threshold',
						]
						SCREAMING_SNAKE_CASE_													=  [x for x in hidden_size_based_keys if self.is_auto(__SCREAMING_SNAKE_CASE)]

						if len(__SCREAMING_SNAKE_CASE) > 0:
										if hasattr(model.config  ,						'hidden_size'):
														SCREAMING_SNAKE_CASE_													=  model.config.hidden_size
										elif hasattr(model.config  ,						'hidden_sizes'):
														# if there are many hidden sizes pick the largest one
														SCREAMING_SNAKE_CASE_													=  max(model.config.hidden_sizes)
										else:
														raise ValueError(
														    'The model\'s config file has neither `hidden_size` nor `hidden_sizes` entry, '
														    'therefore it\'s not possible to automatically fill out the following `auto` entries '
														    f"""in the DeepSpeed config file: {hidden_size_auto_keys}. You can fix that by replacing """
														    '`auto` values for these keys with an integer value of your choice.')

										self.fill_only('zero_optimization.reduce_bucket_size'  ,						hidden_size * hidden_size)
										if self.is_zeroa():
														# automatically assign the optimal config values based on model config
														self.fill_only('zero_optimization.stage3_prefetch_bucket_size'  ,						0.9 * hidden_size * hidden_size)
														self.fill_only('zero_optimization.stage3_param_persistence_threshold'  ,						10 * hidden_size)

        # scheduler
						self.fill_match('scheduler.params.total_num_steps'  ,						__SCREAMING_SNAKE_CASE  ,						'num_training_steps (calculated)')
						self.fill_match('scheduler.params.warmup_num_steps'  ,						args.get_warmup_steps(__SCREAMING_SNAKE_CASE)  ,						'warmup_steps')

						if len(self.mismatches) > 0:
										SCREAMING_SNAKE_CASE_													=  '\n'.join(self.mismatches)
										raise ValueError(
										    'Please correct the following DeepSpeed config values that mismatch TrainingArguments'
										    f""" values:\n{mismatches}\nThe easiest method is to set these DeepSpeed config values to 'auto'.""")






UpperCamelCase__					: List[str]								=      None


def    _UpperCAmelCase    (					_SCREAMING_SNAKE_CASE   :						Optional[Any]    ):





				"""simple docstring"""


				global _hf_deepspeed_config_weak_ref
				# will go away automatically when HfDeepSpeedConfig is destroyed (when TrainingArguments is destroyed)
				SCREAMING_SNAKE_CASE_													=  weakref.ref(snake_case__    )


def    _UpperCAmelCase    (					):





				"""simple docstring"""


				global _hf_deepspeed_config_weak_ref
				SCREAMING_SNAKE_CASE_													=  None


def    _UpperCAmelCase    (					):





				"""simple docstring"""


				if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
								return _hf_deepspeed_config_weak_ref().is_zeroa()
				else:
								return False


def    _UpperCAmelCase    (					):





				"""simple docstring"""


				if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
								return _hf_deepspeed_config_weak_ref().config
				else:
								return None


def    _UpperCAmelCase    (					_SCREAMING_SNAKE_CASE   :						Optional[int]	,				_SCREAMING_SNAKE_CASE   :						List[Any]	,				_SCREAMING_SNAKE_CASE   :						Dict	,				_SCREAMING_SNAKE_CASE   :						Optional[Any]	,				_SCREAMING_SNAKE_CASE   :						Dict    ):





				"""simple docstring"""


				from accelerate.utils import DummyOptim, DummyScheduler

				SCREAMING_SNAKE_CASE_													=  hf_deepspeed_config.config

				# Optimizer + Scheduler
				# Currently supported combos:
				# 1. DS scheduler + DS optimizer: Yes
				# 2. HF scheduler + HF optimizer: Yes
				# 3. DS scheduler + HF optimizer: Yes
				# 4. HF scheduler + DS optimizer: No
				#
				# Unless Offload is enabled in which case it's:
				# 1. DS scheduler + DS optimizer: Yes
				# 2. HF scheduler + HF optimizer: Mostly*
				# 3. DS scheduler + HF optimizer: Mostly*
				# 4. HF scheduler + DS optimizer: No
				#
				# Mostly*: All non-native DeepSpeed optimizers that have both CPU and GPU implementation should work (except LAMB)

				SCREAMING_SNAKE_CASE_													=  None
				if "optimizer" in config:
								if args.adafactor:
												raise ValueError(
												    '--adafactor was passed, but also found `optimizer` configured in the DeepSpeed config. '
												    'Only one optimizer can be configured.'    )
								SCREAMING_SNAKE_CASE_													=  DummyOptim(params=snake_case__    )
				else:
								if hf_deepspeed_config.is_offload():
												logger.info(
												    'Detected ZeRO Offload and non-DeepSpeed optimizers: This combination should work as long as the'
												    ' custom optimizer has both CPU and GPU implementation (except LAMB)'    )

								# ds supports Adam, OneBitAdam, and Lamb optimizers and can import other optimizers from torch.
								# But trainer uses AdamW by default.
								SCREAMING_SNAKE_CASE_													=  trainer.create_optimizer()
								# To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
								SCREAMING_SNAKE_CASE_													=  True

				SCREAMING_SNAKE_CASE_													=  None
				if "scheduler" in config:
								SCREAMING_SNAKE_CASE_													=  DummyScheduler(snake_case__    )
				else:
								if isinstance(snake_case__	,				snake_case__    ):
												raise ValueError(
												    'Found `optimizer` configured in the DeepSpeed config, but no `scheduler`. '
												    'Please configure a scheduler in the DeepSpeed config.'    )
								SCREAMING_SNAKE_CASE_													=  trainer.create_scheduler(num_training_steps=snake_case__	,				optimizer=snake_case__    )

				return optimizer, lr_scheduler


def    _UpperCAmelCase    (					_SCREAMING_SNAKE_CASE   :						Any	,				_SCREAMING_SNAKE_CASE   :						Optional[int]	,				_SCREAMING_SNAKE_CASE   :						Optional[int]=False    ):





				"""simple docstring"""


				from deepspeed.utils import logger as ds_logger

				SCREAMING_SNAKE_CASE_													=  trainer.model
				SCREAMING_SNAKE_CASE_													=  trainer.args

				SCREAMING_SNAKE_CASE_													=  trainer.accelerator.state.deepspeed_plugin.hf_ds_config

				# resume config update - some bits like `model` and `num_training_steps` only become available during train
				hf_deepspeed_config.trainer_config_finalize(snake_case__	,				snake_case__	,				snake_case__    )

				# set the Deepspeed log level consistent with the Trainer
				ds_logger.setLevel(args.get_process_log_level()    )

				if inference:
								# only Z3 makes sense for the inference
								if not hf_deepspeed_config.is_zeroa():
												raise ValueError('ZeRO inference only makes sense with ZeRO Stage 3 - please adjust your config'    )

								# in case the training config is re-used for inference
								hf_deepspeed_config.del_config_sub_tree('optimizer'    )
								hf_deepspeed_config.del_config_sub_tree('lr_scheduler'    )
								SCREAMING_SNAKE_CASE_							,     SCREAMING_SNAKE_CASE_													=  None, None
								SCREAMING_SNAKE_CASE_													=  None
				else:
								SCREAMING_SNAKE_CASE_													=  None  # important for when deepspeed_init is used as re-init
								SCREAMING_SNAKE_CASE_													=  list(filter(lambda _SCREAMING_SNAKE_CASE    : p.requires_grad	,				model.parameters()    )    )
								SCREAMING_SNAKE_CASE_							,     SCREAMING_SNAKE_CASE_													=  deepspeed_optim_sched(
								    snake_case__	,				snake_case__	,				snake_case__	,				snake_case__	,				snake_case__    )

				# keep for quick debug:
				# from pprint import pprint; pprint(config)

				return optimizer, lr_scheduler





def    _UpperCAmelCase    (					_SCREAMING_SNAKE_CASE   :						Tuple	,				_SCREAMING_SNAKE_CASE   :						Tuple    ):





				"""simple docstring"""


				import glob

				SCREAMING_SNAKE_CASE_													=  sorted(glob.glob(f"""{checkpoint_path}/global_step*"""    )    )

				if len(snake_case__    ) > 0:
								logger.info(f"""Attempting to resume from {checkpoint_path}"""    )
								# this magically updates self.optimizer and self.lr_scheduler
								SCREAMING_SNAKE_CASE_							,     SCREAMING_SNAKE_CASE_													=  deepspeed_engine.load_checkpoint(
								    snake_case__	,				load_optimizer_states=snake_case__	,				load_lr_scheduler_states=snake_case__    )
								if load_path is None:
												raise ValueError(f"""[deepspeed] failed to resume from checkpoint {checkpoint_path}"""    )
				else:
								raise ValueError(f"""Can't find a valid checkpoint at {checkpoint_path}"""    )





from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__					: Optional[int]								=      logging.get_logger(__name__)

UpperCamelCase__					: int								=      {
    "google/pegasus-x-base": "https://huggingface.co/google/pegasus-x-base/resolve/main/config.json",
    "google/pegasus-x-large": "https://huggingface.co/google/pegasus-x-large/resolve/main/config.json",
    # See all PEGASUS-X models at https://huggingface.co/models?filter=pegasus-x
}



class  __snake_case					(      lowerCAmelCase__    ):
  __lowerCAmelCase       :		List[Any]  	=					'pegasus_x'
  __lowerCAmelCase       :		Optional[int]  	=					['past_key_values']
  __lowerCAmelCase       :		List[str]  	=					{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}

  def __init__(   self  ,						_A=96103  ,						_A=16384  ,						_A=16  ,						_A=4096  ,						_A=16  ,						_A=16  ,						_A=4096  ,						_A=16  ,						_A=0.0  ,						_A=0.0  ,						_A=True  ,						_A=True  ,						_A="gelu"  ,						_A=1024  ,						_A=0.1  ,						_A=0.0  ,						_A=0.0  ,						_A=0.0_2  ,						_A=0  ,						_A=True  ,						_A=0  ,						_A=1  ,						_A=1  ,						_A=32  ,						_A=512  ,						_A=True  ,						**_A  ,						):
      SCREAMING_SNAKE_CASE_													=  vocab_size
      SCREAMING_SNAKE_CASE_													=  max_position_embeddings
      SCREAMING_SNAKE_CASE_													=  d_model
      SCREAMING_SNAKE_CASE_													=  encoder_ffn_dim
      SCREAMING_SNAKE_CASE_													=  encoder_layers
      SCREAMING_SNAKE_CASE_													=  encoder_attention_heads
      SCREAMING_SNAKE_CASE_													=  decoder_ffn_dim
      SCREAMING_SNAKE_CASE_													=  decoder_layers
      SCREAMING_SNAKE_CASE_													=  decoder_attention_heads
      SCREAMING_SNAKE_CASE_													=  dropout
      SCREAMING_SNAKE_CASE_													=  attention_dropout
      SCREAMING_SNAKE_CASE_													=  activation_dropout
      SCREAMING_SNAKE_CASE_													=  activation_function
      SCREAMING_SNAKE_CASE_													=  init_std
      SCREAMING_SNAKE_CASE_													=  encoder_layerdrop
      SCREAMING_SNAKE_CASE_													=  decoder_layerdrop
      SCREAMING_SNAKE_CASE_													=  use_cache
      SCREAMING_SNAKE_CASE_													=  encoder_layers
      SCREAMING_SNAKE_CASE_													=  scale_embedding  # scale factor will be sqrt(d_model) if True

      SCREAMING_SNAKE_CASE_													=  num_global_tokens
      SCREAMING_SNAKE_CASE_													=  block_size
      SCREAMING_SNAKE_CASE_													=  stagger_local_blocks

      super().__init__(
          pad_token_id=_A  ,						eos_token_id=_A  ,						is_encoder_decoder=_A  ,						decoder_start_token_id=_A  ,						forced_eos_token_id=_A  ,						**_A  ,						)

  @property
  def   lowerCAmelCase__       (   self):
      return self.encoder_attention_heads


  @property
  def   lowerCAmelCase__       (   self):
      return self.d_model


