



'''simple docstring'''
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											=			logging.get_logger(__name__)


def   lowercase_	(   ):

			"""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 _A  (    __SCREAMING_SNAKE_CASE    ):





		def __init__(     self ,       __UpperCAmelCase				)       ->     List[str]:

					'''simple docstring'''
					# set global weakref object
					set_hf_deepspeed_config(self				)
					dep_version_check("""accelerate"""				)
					dep_version_check("""deepspeed"""				)
					super().__init__(__UpperCAmelCase				)



class _A  (    __SCREAMING_SNAKE_CASE    ):





		def __init__(     self ,       __UpperCAmelCase				)       ->     Dict:

					'''simple docstring'''
					super().__init__(__UpperCAmelCase				)
					__UpperCAmelCase   :      Optional[Any]										=     None
					__UpperCAmelCase   :      Any										=     []





		def 							__A       (     self				)       ->     List[Any]:

					'''simple docstring'''
					if self._dtype is None:
								raise ValueError("""trainer_config_process() wasn't called yet to tell dtype"""				)
					return self._dtype





		def 							__A       (     self ,       __UpperCAmelCase				)       ->     Any:

					'''simple docstring'''
					__UpperCAmelCase   :      int										=     self.get_value(__UpperCAmelCase				)
					if val is None:
								return False
					else:
								return val == "auto"





		def 							__A       (     self ,       __UpperCAmelCase ,       __UpperCAmelCase ,       __UpperCAmelCase=None ,       __UpperCAmelCase=True				)       ->     str:

					'''simple docstring'''
					__UpperCAmelCase				,  __UpperCAmelCase   :      int										=     self.find_config_node(__UpperCAmelCase				)
					if config is None:
								return

					if config.get(__UpperCAmelCase				) == "auto":
								__UpperCAmelCase   :      Union[str, Any]										=     hf_val
								return

					if not must_match:
								return

					__UpperCAmelCase   :      int										=     config.get(__UpperCAmelCase				)
					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}'				)

		_SCREAMING_SNAKE_CASE					:    Optional[int]						      =			partialmethod(__SCREAMING_SNAKE_CASE	,				must_match=__SCREAMING_SNAKE_CASE    )





		def 							__A       (     self ,       __UpperCAmelCase				)       ->     List[Any]:

					'''simple docstring'''
					__UpperCAmelCase   :      Union[str, Any]										=     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""" ,       __UpperCAmelCase ,       """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:
								__UpperCAmelCase   :      Tuple										=     """apex""" if args.fpaa_backend == """apex""" else """amp"""
					else:
								__UpperCAmelCase   :      List[str]										=     None

					if args.save_on_each_node:
								# deepspeed uses shared storage by default. Let's override this setting if save_on_each_node == True
								__UpperCAmelCase   :      Optional[int]										=     self.config.get("""checkpoint""" ,       {}				)
								__UpperCAmelCase   :      List[str]										=     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"""				):
								__UpperCAmelCase   :      Dict										=     torch.bfloataa
					elif self.is_false("""fp16.enabled"""				):
								__UpperCAmelCase   :      Optional[int]										=     torch.floataa
					else:
								__UpperCAmelCase   :      List[Any]										=     torch.floataa





		def 							__A       (     self ,       __UpperCAmelCase ,       __UpperCAmelCase ,       __UpperCAmelCase				)       ->     Any:

					'''simple docstring'''
					__UpperCAmelCase   :      Optional[Any]										=     [
					    """zero_optimization.reduce_bucket_size""",
					    """zero_optimization.stage3_prefetch_bucket_size""",
					    """zero_optimization.stage3_param_persistence_threshold""",
					]
					__UpperCAmelCase   :      Optional[Any]										=     [x for x in hidden_size_based_keys if self.is_auto(__UpperCAmelCase				)]

					if len(__UpperCAmelCase				) > 0:
								if hasattr(model.config ,       """hidden_size"""				):
											__UpperCAmelCase   :      Union[str, Any]										=     model.config.hidden_size
								elif hasattr(model.config ,       """hidden_sizes"""				):
											# if there are many hidden sizes pick the largest one
											__UpperCAmelCase   :      int										=     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""" ,       __UpperCAmelCase ,       """num_training_steps (calculated)"""				)
					self.fill_match("""scheduler.params.warmup_num_steps""" ,       args.get_warmup_steps(__UpperCAmelCase				) ,       """warmup_steps"""				)

					if len(self.mismatches				) > 0:
								__UpperCAmelCase   :      str										=     """\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											=			None


def   lowercase_	(   lowerCAmelCase__	:  Tuple ):

			"""simple docstring"""





			global _hf_deepspeed_config_weak_ref
			# will go away automatically when HfDeepSpeedConfig is destroyed (when TrainingArguments is destroyed)
			__UpperCAmelCase   :      str										=     weakref.ref(lowerCAmelCase__ )


def   lowercase_	(   ):

			"""simple docstring"""





			global _hf_deepspeed_config_weak_ref
			__UpperCAmelCase   :      Any										=     None


def   lowercase_	(   ):

			"""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   lowercase_	(   ):

			"""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   lowercase_	(   lowerCAmelCase__	:  str			,							lowerCAmelCase__	:  Optional[int]			,							lowerCAmelCase__	:  Optional[int]			,							lowerCAmelCase__	:  str			,							lowerCAmelCase__	:  Optional[int] ):

			"""simple docstring"""





			from accelerate.utils import DummyOptim, DummyScheduler

			__UpperCAmelCase   :      Optional[int]										=     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)

			__UpperCAmelCase   :      List[str]										=     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.""" )
						__UpperCAmelCase   :      Optional[int]										=     DummyOptim(params=lowerCAmelCase__ )
			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.
						__UpperCAmelCase   :      Optional[int]										=     trainer.create_optimizer()
						# To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
						__UpperCAmelCase   :      Any										=     True

			__UpperCAmelCase   :      Optional[Any]										=     None
			if "scheduler" in config:
						__UpperCAmelCase   :      int										=     DummyScheduler(lowerCAmelCase__ )
			else:
						if isinstance(lowerCAmelCase__			,							lowerCAmelCase__ ):
									raise ValueError(
									    """Found `optimizer` configured in the DeepSpeed config, but no `scheduler`. """
									    """Please configure a scheduler in the DeepSpeed config.""" )
						__UpperCAmelCase   :      int										=     trainer.create_scheduler(num_training_steps=lowerCAmelCase__			,							optimizer=lowerCAmelCase__ )

			return optimizer, lr_scheduler


def   lowercase_	(   lowerCAmelCase__	:  int			,							lowerCAmelCase__	:  Optional[int]			,							lowerCAmelCase__	:  Optional[Any]=False ):

			"""simple docstring"""





			from deepspeed.utils import logger as ds_logger

			__UpperCAmelCase   :      str										=     trainer.model
			__UpperCAmelCase   :      List[Any]										=     trainer.args

			__UpperCAmelCase   :      Dict										=     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(lowerCAmelCase__			,							lowerCAmelCase__			,							lowerCAmelCase__ )

			# 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""" )
						__UpperCAmelCase				,  __UpperCAmelCase   :      Optional[Any]										=     None, None
						__UpperCAmelCase   :      Tuple										=     None
			else:
						__UpperCAmelCase   :      List[Any]										=     None  # important for when deepspeed_init is used as re-init
						__UpperCAmelCase   :      List[str]										=     list(filter(lambda lowerCAmelCase__ : p.requires_grad			,							model.parameters() ) )
						__UpperCAmelCase				,  __UpperCAmelCase   :      List[Any]										=     deepspeed_optim_sched(
						    lowerCAmelCase__			,							lowerCAmelCase__			,							lowerCAmelCase__			,							lowerCAmelCase__			,							lowerCAmelCase__ )

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

			return optimizer, lr_scheduler






def   lowercase_	(   lowerCAmelCase__	:  Dict			,							lowerCAmelCase__	:  Any ):

			"""simple docstring"""





			import glob

			__UpperCAmelCase   :      Union[str, Any]										=     sorted(glob.glob(f'{checkpoint_path}/global_step*' ) )

			if len(lowerCAmelCase__ ) > 0:
						logger.info(f'Attempting to resume from {checkpoint_path}' )
						# this magically updates self.optimizer and self.lr_scheduler
						__UpperCAmelCase				,  __UpperCAmelCase   :      Optional[int]										=     deepspeed_engine.load_checkpoint(
						    lowerCAmelCase__			,							load_optimizer_states=lowerCAmelCase__			,							load_lr_scheduler_states=lowerCAmelCase__ )
						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}' )









'''simple docstring'''
from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_UpperCamelCase											=			logging.get_logger(__name__)

_UpperCamelCase											=			{
    '''nghuyong/ernie-1.0-base-zh''': '''https://huggingface.co/nghuyong/ernie-1.0-base-zh/resolve/main/config.json''',
    '''nghuyong/ernie-2.0-base-en''': '''https://huggingface.co/nghuyong/ernie-2.0-base-en/resolve/main/config.json''',
    '''nghuyong/ernie-2.0-large-en''': '''https://huggingface.co/nghuyong/ernie-2.0-large-en/resolve/main/config.json''',
    '''nghuyong/ernie-3.0-base-zh''': '''https://huggingface.co/nghuyong/ernie-3.0-base-zh/resolve/main/config.json''',
    '''nghuyong/ernie-3.0-medium-zh''': '''https://huggingface.co/nghuyong/ernie-3.0-medium-zh/resolve/main/config.json''',
    '''nghuyong/ernie-3.0-mini-zh''': '''https://huggingface.co/nghuyong/ernie-3.0-mini-zh/resolve/main/config.json''',
    '''nghuyong/ernie-3.0-micro-zh''': '''https://huggingface.co/nghuyong/ernie-3.0-micro-zh/resolve/main/config.json''',
    '''nghuyong/ernie-3.0-nano-zh''': '''https://huggingface.co/nghuyong/ernie-3.0-nano-zh/resolve/main/config.json''',
    '''nghuyong/ernie-gram-zh''': '''https://huggingface.co/nghuyong/ernie-gram-zh/resolve/main/config.json''',
    '''nghuyong/ernie-health-zh''': '''https://huggingface.co/nghuyong/ernie-health-zh/resolve/main/config.json''',
}



class _A  (    __SCREAMING_SNAKE_CASE    ):
		_SCREAMING_SNAKE_CASE					:    Optional[int]						      =			"ernie"





		def __init__(     self ,       __UpperCAmelCase=30_522 ,       __UpperCAmelCase=768 ,       __UpperCAmelCase=12 ,       __UpperCAmelCase=12 ,       __UpperCAmelCase=3_072 ,       __UpperCAmelCase="gelu" ,       __UpperCAmelCase=0.1 ,       __UpperCAmelCase=0.1 ,       __UpperCAmelCase=512 ,       __UpperCAmelCase=2 ,       __UpperCAmelCase=3 ,       __UpperCAmelCase=False ,       __UpperCAmelCase=0.02 ,       __UpperCAmelCase=1E-12 ,       __UpperCAmelCase=0 ,       __UpperCAmelCase="absolute" ,       __UpperCAmelCase=True ,       __UpperCAmelCase=None ,       **__UpperCAmelCase ,       )       ->     List[Any]:

					'''simple docstring'''
					super().__init__(pad_token_id=__UpperCAmelCase ,       **__UpperCAmelCase				)

					__UpperCAmelCase   :      List[str]										=     vocab_size
					__UpperCAmelCase   :      Optional[int]										=     hidden_size
					__UpperCAmelCase   :      Optional[Any]										=     num_hidden_layers
					__UpperCAmelCase   :      int										=     num_attention_heads
					__UpperCAmelCase   :      int										=     hidden_act
					__UpperCAmelCase   :      List[Any]										=     intermediate_size
					__UpperCAmelCase   :      Optional[Any]										=     hidden_dropout_prob
					__UpperCAmelCase   :      Dict										=     attention_probs_dropout_prob
					__UpperCAmelCase   :      int										=     max_position_embeddings
					__UpperCAmelCase   :      str										=     type_vocab_size
					__UpperCAmelCase   :      int										=     task_type_vocab_size
					__UpperCAmelCase   :      List[str]										=     use_task_id
					__UpperCAmelCase   :      Union[str, Any]										=     initializer_range
					__UpperCAmelCase   :      Tuple										=     layer_norm_eps
					__UpperCAmelCase   :      List[str]										=     position_embedding_type
					__UpperCAmelCase   :      Union[str, Any]										=     use_cache
					__UpperCAmelCase   :      Optional[int]										=     classifier_dropout



class _A  (    __SCREAMING_SNAKE_CASE    ):





		@property
		def 							__A       (     self				)       ->     Mapping[str, Mapping[int, str]]:

					'''simple docstring'''
					if self.task == "multiple-choice":
								__UpperCAmelCase   :      Union[str, Any]										=     {0: """batch""", 1: """choice""", 2: """sequence"""}
					else:
								__UpperCAmelCase   :      Tuple										=     {0: """batch""", 1: """sequence"""}
					return OrderedDict(
					    [
					        ("""input_ids""", dynamic_axis),
					        ("""attention_mask""", dynamic_axis),
					        ("""token_type_ids""", dynamic_axis),
					        ("""task_type_ids""", dynamic_axis),
					    ]				)




