




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

_lowercase						   =     logging.get_logger(__name__)

def 		lowerCAmelCase				(			):
						'''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    __lowercase							(      snake_case_						):


				'''simple docstring'''


				def __init__( self :  Optional[int]   ,SCREAMING_SNAKE_CASE_ :  Optional[Any]   ):
										# set global weakref object
										set_hf_deepspeed_config(self   )
										dep_version_check('''accelerate'''   )
										dep_version_check('''deepspeed'''   )
										super().__init__(SCREAMING_SNAKE_CASE_   )







class    __lowercase							(      snake_case_						):


				'''simple docstring'''


				def __init__( self :  Optional[int]   ,SCREAMING_SNAKE_CASE_ :  int   ):
										super().__init__(SCREAMING_SNAKE_CASE_   )
										snake_case__							:       Union[str, Any]				      =       None
										snake_case__							:       List[Any]				      =       []

				def    _lowerCamelCase       ( self :  List[str]   ):
										if self._dtype is None:
																raise ValueError('''trainer_config_process() wasn\'t called yet to tell dtype'''   )
										return self._dtype

				def    _lowerCamelCase       ( self :  List[str]   ,SCREAMING_SNAKE_CASE_ :  Dict   ):
										snake_case__							:       List[Any]				      =       self.get_value(SCREAMING_SNAKE_CASE_   )
										if val is None:
																return False
										else:
																return val == "auto"

				def    _lowerCamelCase       ( self :  int   ,SCREAMING_SNAKE_CASE_ :  Optional[Any]   ,SCREAMING_SNAKE_CASE_ :  Dict   ,SCREAMING_SNAKE_CASE_ :  Tuple=None   ,SCREAMING_SNAKE_CASE_ :  Union[str, Any]=True   ):
										snake_case__      ,						snake_case__							:       int				      =       self.find_config_node(SCREAMING_SNAKE_CASE_   )
										if config is None:
																return

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

										if not must_match:
																return

										snake_case__							:       List[Any]				      =       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						      =partialmethod(snake_case_    , must_match=snake_case_						)

				def    _lowerCamelCase       ( self :  Tuple   ,SCREAMING_SNAKE_CASE_ :  int   ):
										snake_case__							:       Dict				      =       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:
																snake_case__							:       Tuple				      =       '''apex''' if args.fpaa_backend == '''apex''' else '''amp'''
										else:
																snake_case__							:       Union[str, Any]				      =       None

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





				def    _lowerCamelCase       ( self :  Tuple   ,SCREAMING_SNAKE_CASE_ :  str   ,SCREAMING_SNAKE_CASE_ :  int   ,SCREAMING_SNAKE_CASE_ :  int   ):
										snake_case__							:       Union[str, Any]				      =       [
										    '''zero_optimization.reduce_bucket_size''',
										    '''zero_optimization.stage3_prefetch_bucket_size''',
										    '''zero_optimization.stage3_param_persistence_threshold''',
										]
										snake_case__							:       List[str]				      =       [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'''   ):
																						snake_case__							:       List[str]				      =       model.config.hidden_size
																elif hasattr(model.config   ,'''hidden_sizes'''   ):
																						# if there are many hidden sizes pick the largest one
																						snake_case__							:       List[str]				      =       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'''   ,1_0 * 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:
																snake_case__							:       Optional[int]				      =       '''\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\'.'''   )







_lowercase						   =     None

def 		lowerCAmelCase				(			lowerCamelCase		):
						'''simple docstring'''


						global _hf_deepspeed_config_weak_ref
						# will go away automatically when HfDeepSpeedConfig is destroyed (when TrainingArguments is destroyed)
						snake_case__							:       Optional[Any]				      =       weakref.ref(lowerCamelCase		)

def 		lowerCAmelCase				(			):
						'''simple docstring'''


						global _hf_deepspeed_config_weak_ref
						snake_case__							:       str				      =       None

def 		lowerCAmelCase				(			):
						'''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 		lowerCAmelCase				(			):
						'''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 		lowerCAmelCase				(			lowerCamelCase   ,      lowerCamelCase   ,      lowerCamelCase   ,      lowerCamelCase   ,      lowerCamelCase		):
						'''simple docstring'''


						from accelerate.utils import DummyOptim, DummyScheduler

						snake_case__							:       List[Any]				      =       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)

						snake_case__							:       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.'''		)
												snake_case__							:       Union[str, Any]				      =       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.
												snake_case__							:       int				      =       trainer.create_optimizer()
												# To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
												snake_case__							:       Tuple				      =       True

						snake_case__							:       Tuple				      =       None
						if "scheduler" in config:
												snake_case__							:       Any				      =       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.'''		)
												snake_case__							:       Optional[int]				      =       trainer.create_scheduler(num_training_steps=lowerCamelCase   ,      optimizer=lowerCamelCase		)

						return optimizer, lr_scheduler

def 		lowerCAmelCase				(			lowerCamelCase   ,      lowerCamelCase   ,      lowerCamelCase=False		):
						'''simple docstring'''


						from deepspeed.utils import logger as ds_logger

						snake_case__							:       List[Any]				      =       trainer.model
						snake_case__							:       Optional[Any]				      =       trainer.args

						snake_case__							:       Union[str, Any]				      =       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'''		)
												snake_case__      ,						snake_case__							:       List[Any]				      =       None, None
												snake_case__							:       List[str]				      =       None
						else:
												snake_case__							:       List[str]				      =       None  # important for when deepspeed_init is used as re-init
												snake_case__							:       int				      =       list(filter(lambda lowerCamelCase		: p.requires_grad   ,      model.parameters()		)		)
												snake_case__      ,						snake_case__							:       int				      =       deepspeed_optim_sched(
												    lowerCamelCase   ,      lowerCamelCase   ,      lowerCamelCase   ,      lowerCamelCase   ,      lowerCamelCase		)

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

						return optimizer, lr_scheduler

def 		lowerCAmelCase				(			lowerCamelCase   ,      lowerCamelCase		):
						'''simple docstring'''


						import glob

						snake_case__							:       int				      =       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
												snake_case__      ,						snake_case__							:       Union[str, Any]				      =       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"""





_lowercase						   =     [0.4_85, 0.4_56, 0.4_06]
_lowercase						   =     [0.2_29, 0.2_24, 0.2_25]
_lowercase						   =     [0.5, 0.5, 0.5]
_lowercase						   =     [0.5, 0.5, 0.5]
_lowercase						   =     [0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]
_lowercase						   =     [0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]


