



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

A__     				=	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  __lowerCAmelCase			(					a_	):




				def __init__(							self						,    _snake_case		):



									"""simple docstring"""





									set_hf_deepspeed_config(self		)
									dep_version_check("""accelerate"""		)
									dep_version_check("""deepspeed"""		)
									super().__init__(snake_case__		)





class  __lowerCAmelCase			(					a_	):




				def __init__(							self						,    _snake_case		):



									"""simple docstring"""





									super().__init__(snake_case__		)
									_lowerCAmelCase      =   None
									_lowerCAmelCase      =   []




				def 							snake_case       (							self		):



									"""simple docstring"""





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




				def 							snake_case       (							self						,    _snake_case		):



									"""simple docstring"""





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




				def 							snake_case       (							self						,    _snake_case						,    _snake_case						,    _snake_case=None						,    _snake_case=True		):



									"""simple docstring"""





									_lowerCAmelCase     ,     _lowerCAmelCase      =   self.find_config_node(snake_case__		)
									if config is None:
														return

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

									if not must_match:
														return

									_lowerCAmelCase      =   config.get(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(a_	,						must_match=a_	)




				def 							snake_case       (							self						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase      =   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"""						,    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:
														_lowerCAmelCase      =   """apex""" if args.fpaa_backend == """apex""" else """amp"""
									else:
														_lowerCAmelCase      =   None

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




				def 							snake_case       (							self						,    _snake_case						,    _snake_case						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase      =   [
									    """zero_optimization.reduce_bucket_size""",
									    """zero_optimization.stage3_prefetch_bucket_size""",
									    """zero_optimization.stage3_param_persistence_threshold""",
									]
									_lowerCAmelCase      =   [x for x in hidden_size_based_keys if self.is_auto(snake_case__		)]

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

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


def 		_UpperCAmelCase    (				snake_case      ):





					"""simple docstring"""

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


def 		_UpperCAmelCase    (				):





					"""simple docstring"""

					global _hf_deepspeed_config_weak_ref
					_lowerCAmelCase      =   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    (				snake_case     ,						snake_case     ,						snake_case     ,						snake_case     ,						snake_case      ):





					"""simple docstring"""

					from accelerate.utils import DummyOptim, DummyScheduler

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

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

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

					return optimizer, lr_scheduler


def 		_UpperCAmelCase    (				snake_case     ,						snake_case     ,						snake_case=False      ):





					"""simple docstring"""

					from deepspeed.utils import logger as ds_logger

					_lowerCAmelCase      =   trainer.model
					_lowerCAmelCase      =   trainer.args

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

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

					return optimizer, lr_scheduler


def 		_UpperCAmelCase    (				snake_case     ,						snake_case      ):





					"""simple docstring"""

					import glob

					_lowerCAmelCase      =   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
										_lowerCAmelCase     ,     _lowerCAmelCase      =   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}'      )








import unittest
from pathlib import Path

from transformers import is_vision_available, load_tool
from transformers.testing_utils import get_tests_dir

from .test_tools_common import ToolTesterMixin


if is_vision_available():
	from PIL import Image





class  __lowerCAmelCase			(					unittest.TestCase	,						lowerCamelCase__	):




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   load_tool("""image-question-answering"""		)
									self.tool.setup()
									_lowerCAmelCase      =   load_tool("""image-question-answering"""						,    remote=_snake_case		)




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   Image.open(Path(get_tests_dir("""fixtures/tests_samples/COCO"""		)		) / """000000039769.png"""		)
									_lowerCAmelCase      =   self.tool(_snake_case						,    """How many cats are sleeping on the couch?"""		)
									self.assertEqual(_snake_case						,    """2"""		)




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   Image.open(Path(get_tests_dir("""fixtures/tests_samples/COCO"""		)		) / """000000039769.png"""		)
									_lowerCAmelCase      =   self.remote_tool(_snake_case						,    """How many cats are sleeping on the couch?"""		)
									self.assertEqual(_snake_case						,    """2"""		)




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   Image.open(Path(get_tests_dir("""fixtures/tests_samples/COCO"""		)		) / """000000039769.png"""		)
									_lowerCAmelCase      =   self.tool(image=_snake_case						,    question="""How many cats are sleeping on the couch?"""		)
									self.assertEqual(_snake_case						,    """2"""		)




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   Image.open(Path(get_tests_dir("""fixtures/tests_samples/COCO"""		)		) / """000000039769.png"""		)
									_lowerCAmelCase      =   self.remote_tool(image=_snake_case						,    question="""How many cats are sleeping on the couch?"""		)
									self.assertEqual(_snake_case						,    """2"""		)



