




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

__a			:	str							     =		logging.get_logger(__name__)




def  a_		(							)	->			Tuple:

			'''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					(				lowerCamelCase_       ):







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

										"""simple docstring"""



										set_hf_deepspeed_config(self				)
										dep_version_check('accelerate'				)
										dep_version_check('deepspeed'				)
										super().__init__(__UpperCAmelCase				)







class     A					(				lowerCamelCase_       ):







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

										"""simple docstring"""



										super().__init__(__UpperCAmelCase				)
										UpperCamelCase_     				= None
										UpperCamelCase_     				= []







							def        lowercase__							(    self   :	Tuple				)		->					List[str]:

										"""simple docstring"""



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







							def        lowercase__							(    self   :	Tuple   ,     __UpperCAmelCase   :	Union[str, Any]				)		->					Optional[Any]:

										"""simple docstring"""



										UpperCamelCase_     				= self.get_value(__UpperCAmelCase				)
										if val is None:
													return False
										else:
													return val == "auto"







							def        lowercase__							(    self   :	int   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	Dict   ,     __UpperCAmelCase   :	Union[str, Any]=None   ,     __UpperCAmelCase   :	Any=True				)		->					Optional[int]:

										"""simple docstring"""



										UpperCamelCase_				,   UpperCamelCase_     				= self.find_config_node(__UpperCAmelCase				)
										if config is None:
													return

										if config.get(__UpperCAmelCase				) == "auto":
													UpperCamelCase_     				= hf_val
													return

										if not must_match:
													return

										UpperCamelCase_     				= 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  :		Any			  =       partialmethod(lowerCamelCase_  ,						must_match=lowerCamelCase_       )







							def        lowercase__							(    self   :	int   ,     __UpperCAmelCase   :	int				)		->					Optional[int]:

										"""simple docstring"""



										UpperCamelCase_     				= 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_     				= 'apex' if args.fpaa_backend == 'apex' else 'amp'
										else:
													UpperCamelCase_     				= 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_     				= self.config.get('checkpoint'   ,     {}				)
													UpperCamelCase_     				= 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_     				= torch.bfloataa
										elif self.is_false('fp16.enabled'				):
													UpperCamelCase_     				= torch.floataa
										else:
													UpperCamelCase_     				= torch.floataa







							def        lowercase__							(    self   :	Dict   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	Union[str, Any]   ,     __UpperCAmelCase   :	str				)		->					str:

										"""simple docstring"""



										UpperCamelCase_     				= [
										    'zero_optimization.reduce_bucket_size',
										    'zero_optimization.stage3_prefetch_bucket_size',
										    'zero_optimization.stage3_param_persistence_threshold',
										]
										UpperCamelCase_     				= [x for x in hidden_size_based_keys if self.is_auto(__UpperCAmelCase				)]

										if len(__UpperCAmelCase				) > 0:
													if hasattr(model.config   ,     'hidden_size'				):
																UpperCamelCase_     				= model.config.hidden_size
													elif hasattr(model.config   ,     'hidden_sizes'				):
																# if there are many hidden sizes pick the largest one
																UpperCamelCase_     				= 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_     				= '\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\'.'''				)







__a			:	Union[str, Any]							     =		None




def  a_		(							__snake_case					)	->			Union[str, Any]:

			'''simple docstring'''


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




def  a_		(							)	->			Union[str, Any]:

			'''simple docstring'''


			global _hf_deepspeed_config_weak_ref
			UpperCamelCase_     				= None




def  a_		(							)	->			Optional[int]:

			'''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  a_		(							)	->			Any:

			'''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  a_		(							__snake_case     ,		__snake_case     ,		__snake_case     ,		__snake_case     ,		__snake_case					)	->			str:

			'''simple docstring'''


			from accelerate.utils import DummyOptim, DummyScheduler

			UpperCamelCase_     				= 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_     				= 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_     				= 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.
						UpperCamelCase_     				= trainer.create_optimizer()
						# To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
						UpperCamelCase_     				= True

			UpperCamelCase_     				= None
			if "scheduler" in config:
						UpperCamelCase_     				= 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.'					)
						UpperCamelCase_     				= trainer.create_scheduler(num_training_steps=__snake_case     ,		optimizer=__snake_case					)

			return optimizer, lr_scheduler




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case=False					)	->			str:

			'''simple docstring'''


			from deepspeed.utils import logger as ds_logger

			UpperCamelCase_     				= trainer.model
			UpperCamelCase_     				= trainer.args

			UpperCamelCase_     				= 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'					)
						UpperCamelCase_				,   UpperCamelCase_     				= None, None
						UpperCamelCase_     				= None
			else:
						UpperCamelCase_     				= None  # important for when deepspeed_init is used as re-init
						UpperCamelCase_     				= list(filter(lambda __snake_case					: p.requires_grad     ,		model.parameters()					)					)
						UpperCamelCase_				,   UpperCamelCase_     				= 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  a_		(							__snake_case     ,		__snake_case					)	->			int:

			'''simple docstring'''


			import glob

			UpperCamelCase_     				= 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
						UpperCamelCase_				,   UpperCamelCase_     				= 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 functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			torch.Tensor:

			'''simple docstring'''


			if not inplace:
						UpperCamelCase_     				= ma + ma
			else:
						ma += ma

			return ma




def  a_		(							__snake_case     ,		__snake_case					)	->			torch.Tensor:

			'''simple docstring'''


			UpperCamelCase_     				= -1 * len(__snake_case					)
			UpperCamelCase_     				= list(range(len(tensor.shape[:zero_index]					)					)					)
			return tensor.permute(first_inds + [zero_index + i for i in inds]					)




def  a_		(							__snake_case     ,		__snake_case					)	->			torch.Tensor:

			'''simple docstring'''


			return t.reshape(t.shape[:-no_dims] + (-1,)					)




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case     ,		__snake_case = 1E-4					)	->			torch.Tensor:

			'''simple docstring'''


			UpperCamelCase_     				= mask.expand(*value.shape					)
			return torch.sum(mask * value     ,		dim=__snake_case					) / (eps + torch.sum(__snake_case     ,		dim=__snake_case					))




def  a_		(							__snake_case     ,		__snake_case = 2.31_25     ,		__snake_case = 21.68_75     ,		__snake_case = 6_4					)	->			torch.Tensor:

			'''simple docstring'''


			UpperCamelCase_     				= torch.linspace(__snake_case     ,		__snake_case     ,		no_bins - 1     ,		device=pts.device					)
			UpperCamelCase_     				= torch.sqrt(torch.sum((pts.unsqueeze(-2					) - pts.unsqueeze(-3					)) ** 2     ,		dim=-1					)					)
			return torch.bucketize(__snake_case     ,		__snake_case					)




def  a_		(							__snake_case     ,		__snake_case					)	->			dict:

			'''simple docstring'''


			UpperCamelCase_     				= dicts[0]
			UpperCamelCase_     				= {}
			for k, v in first.items():
						UpperCamelCase_     				= [d[k] for d in dicts]
						if isinstance(__snake_case     ,		__snake_case					):
									UpperCamelCase_     				= dict_multimap(__snake_case     ,		__snake_case					)
						else:
									UpperCamelCase_     				= fn(__snake_case					)

			return new_dict




def  a_		(							__snake_case     ,		__snake_case					)	->			torch.Tensor:

			'''simple docstring'''


			UpperCamelCase_     				= v_bins.view(((1,) * len(x.shape					)) + (len(__snake_case					),)					)
			UpperCamelCase_     				= x[..., None] - reshaped_bins
			UpperCamelCase_     				= torch.argmin(torch.abs(__snake_case					)     ,		dim=-1					)
			return nn.functional.one_hot(__snake_case     ,		num_classes=len(__snake_case					)					).float()




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case = 0     ,		__snake_case = 0					)	->			torch.Tensor:

			'''simple docstring'''


			UpperCamelCase_     				= []
			for i, s in enumerate(data.shape[:no_batch_dims]					):
						UpperCamelCase_     				= torch.arange(__snake_case					)
						UpperCamelCase_     				= r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape					) - i - 1)))					)
						ranges.append(__snake_case					)

			UpperCamelCase_     				= [slice(__snake_case					) for _ in range(len(data.shape					) - no_batch_dims					)]
			UpperCamelCase_     				= inds
			ranges.extend(__snake_case					)
			# Matt note: Editing this to get around the behaviour of using a list as an array index changing
			# in recent Numpy versions
			return data[tuple(__snake_case					)]


__a			:	List[Any]							     =		TypeVar("""T""")




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			Dict[Any, Union[dict, list, tuple, Any]]:

			'''simple docstring'''


			UpperCamelCase_     				= {}
			for k, v in dic.items():
						if isinstance(__snake_case     ,		__snake_case					):
									UpperCamelCase_     				= dict_map(__snake_case     ,		__snake_case     ,		__snake_case					)
						else:
									UpperCamelCase_     				= tree_map(__snake_case     ,		__snake_case     ,		__snake_case					)

			return new_dict




@overload
def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			Any:

			'''simple docstring'''


			...




@overload
def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			dict:

			'''simple docstring'''


			...




@overload
def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			list:

			'''simple docstring'''


			...




@overload
def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			tuple:

			'''simple docstring'''


			...




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			Any:

			'''simple docstring'''


			if isinstance(__snake_case     ,		__snake_case					):
						return dict_map(__snake_case     ,		__snake_case     ,		__snake_case					)
			elif isinstance(__snake_case     ,		__snake_case					):
						return [tree_map(__snake_case     ,		__snake_case     ,		__snake_case					) for x in tree]
			elif isinstance(__snake_case     ,		__snake_case					):
						return tuple(tree_map(__snake_case     ,		__snake_case     ,		__snake_case					) for x in tree					)
			elif isinstance(__snake_case     ,		__snake_case					):
						return fn(__snake_case					)
			else:
						print(type(__snake_case					)					)
						raise ValueError('Not supported'					)





__a			:	Tuple							     =		partial(tree_map, leaf_type=torch.Tensor)
