






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

snake_case_   :	Optional[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        a__      (	lowerCAmelCase__		):






		def __init__(      self							, lowercase__						)				->				Tuple:
				# set global weakref object
				set_hf_deepspeed_config(self						)
				dep_version_check("accelerate"						)
				dep_version_check("deepspeed"						)
				super().__init__(lowercase__						)





class        a__      (	lowerCAmelCase__		):






		def __init__(      self							, lowercase__						)				->				Union[str, Any]:
				super().__init__(lowercase__						)
				__A    =							None
				__A    =							[]






		def 					_lowerCamelCase					(      self						)				->				Union[str, Any]:
				if self._dtype is None:
						raise ValueError("trainer_config_process() wasn't called yet to tell dtype"						)
				return self._dtype






		def 					_lowerCamelCase					(      self							, lowercase__						)				->				str:
				__A    =							self.get_value(lowercase__						)
				if val is None:
						return False
				else:
						return val == "auto"






		def 					_lowerCamelCase					(      self							, lowercase__							, lowercase__							, lowercase__=None							, lowercase__=True						)				->				Dict:
				__A ,   __A    =							self.find_config_node(lowercase__						)
				if config is None:
						return

				if config.get(lowercase__						) == "auto":
						__A    =							hf_val
						return

				if not must_match:
						return

				__A    =							config.get(lowercase__						)
				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}"""						)

		UpperCAmelCase_							:       int 					= partialmethod(lowerCAmelCase__	,       must_match=lowerCAmelCase__		)






		def 					_lowerCamelCase					(      self							, lowercase__						)				->				int:
				__A    =							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"							, lowercase__							, "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:
						__A    =							"apex" if args.fpaa_backend == "apex" else "amp"
				else:
						__A    =							None

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






		def 					_lowerCamelCase					(      self							, lowercase__							, lowercase__							, lowercase__						)				->				Dict:
				__A    =							[
				    "zero_optimization.reduce_bucket_size",
				    "zero_optimization.stage3_prefetch_bucket_size",
				    "zero_optimization.stage3_param_persistence_threshold",
				]
				__A    =							[x for x in hidden_size_based_keys if self.is_auto(lowercase__						)]

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

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



snake_case_   :	Tuple							      =None



def      UpperCAmelCase						( lowerCAmelCase__	):


		'''simple docstring'''

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



def      UpperCAmelCase						( ):


		'''simple docstring'''

		global _hf_deepspeed_config_weak_ref
		__A    =							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						( lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__	):


		'''simple docstring'''

		from accelerate.utils import DummyOptim, DummyScheduler

		__A    =							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)

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

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

		return optimizer, lr_scheduler



def      UpperCAmelCase						( lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__=False	):


		'''simple docstring'''

		from deepspeed.utils import logger as ds_logger

		__A    =							trainer.model
		__A    =							trainer.args

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

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

		return optimizer, lr_scheduler





def      UpperCAmelCase						( lowerCAmelCase__					,     lowerCAmelCase__	):


		'''simple docstring'''

		import glob

		__A    =							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
				__A ,   __A    =							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'''







_UpperCamelCase  						=			{
    "kilogram": 1,
    "gram": pow(10, 3),
    "milligram": pow(10, 6),
    "metric-ton": pow(10, -3),
    "long-ton": 0.0009842073,
    "short-ton": 0.0011023122,
    "pound": 2.2046244202,
    "stone": 0.1574731728,
    "ounce": 35.273990723,
    "carrat": 5000,
    "atomic-mass-unit": 6.0_2213_6652e26,
}

_UpperCamelCase  						=			{
    "kilogram": 1,
    "gram": pow(10, -3),
    "milligram": pow(10, -6),
    "metric-ton": pow(10, 3),
    "long-ton": 1016.04608,
    "short-ton": 907.184,
    "pound": 0.453592,
    "stone": 6.35029,
    "ounce": 0.0283495,
    "carrat": 0.0002,
    "atomic-mass-unit": 1.6_6054_0199e-27,
}
def 			_lowerCAmelCase(      UpperCAmelCase_      :	str					, UpperCAmelCase_      :	str					, UpperCAmelCase_      :	float						) ->							float:
							if to_type not in KILOGRAM_CHART or from_type not in WEIGHT_TYPE_CHART:
														lowerCAmelCase__  =   (
														    F'''Invalid \'from_type\' or \'to_type\' value: {from_type!r}, {to_type!r}\n'''
														    F'''Supported values are: {", ".join(UpperCAmelCase_						)}'''
														)
														raise ValueError(UpperCAmelCase_						)
							return value * KILOGRAM_CHART[to_type] * WEIGHT_TYPE_CHART[from_type]


if __name__ == "__main__":
		import doctest

		doctest.testmod()





