





"""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 	__lowercase  (     )     ->List[str]:



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





 """simple docstring"""


 def __init__(					self      ,		__lowerCamelCase						):



     '''simple docstring'''




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






class 	__snake_case      (	SCREAMING_SNAKE_CASE__  ):





 """simple docstring"""


 def __init__(					self      ,		__lowerCamelCase						):



     '''simple docstring'''




     super().__init__(__lowerCamelCase						)
     __A					:		List[Any]						=  None
     __A					:		Tuple						=  []

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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

 def       UpperCamelCase__(					self      ,		__lowerCamelCase						):



     '''simple docstring'''




     __A					:		List[str]						=  self.get_value(__lowerCamelCase						)
     if val is None:
         return False
     else:
         return val == "auto"

 def       UpperCamelCase__(					self      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase=None      ,		__lowerCamelCase=True						):



     '''simple docstring'''




     __A				, __A					:		List[Any]						=  self.find_config_node(__lowerCamelCase						)
     if config is None:
         return

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

     if not must_match:
         return

     __A					:		Optional[int]						=  config.get(__lowerCamelCase						)
     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(SCREAMING_SNAKE_CASE__							,							must_match=SCREAMING_SNAKE_CASE__  )

 def       UpperCamelCase__(					self      ,		__lowerCamelCase						):



     '''simple docstring'''




     __A					:		List[str]						=  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'''      ,		__lowerCamelCase      ,		'''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					:		Any						=  '''apex''' if args.fpaa_backend == '''apex''' else '''amp'''
     else:
         __A					:		Tuple						=  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					:		Any						=  self.config.get('''checkpoint'''      ,		{}						)
         __A					:		int						=  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					:		Dict						=  torch.bfloataa
     elif self.is_false('''fp16.enabled'''						):
         __A					:		Optional[int]						=  torch.floataa
     else:
         __A					:		List[str]						=  torch.floataa







 def       UpperCamelCase__(					self      ,		__lowerCamelCase      ,		__lowerCamelCase      ,		__lowerCamelCase						):



     '''simple docstring'''




     __A					:		Tuple						=  [
         '''zero_optimization.reduce_bucket_size''',
         '''zero_optimization.stage3_prefetch_bucket_size''',
         '''zero_optimization.stage3_param_persistence_threshold''',
     ]
     __A					:		int						=  [x for x in hidden_size_based_keys if self.is_auto(__lowerCamelCase						)]

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

     if len(self.mismatches						) > 0:
         __A					:		List[Any]						=  '''\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 	__lowercase  (     snake_case_						:						str		)     ->List[str]:



    '''simple docstring'''







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







def 	__lowercase  (     )     ->Optional[Any]:



    '''simple docstring'''







    global _hf_deepspeed_config_weak_ref
    __A					:		Optional[Any]						=  None







def 	__lowercase  (     )     ->Optional[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().is_zeroa()
    else:
        return False







def 	__lowercase  (     )     ->Tuple:



    '''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  (     snake_case_						:						List[Any]       ,snake_case_						:						Union[str, Any]       ,snake_case_						:						str       ,snake_case_						:						List[Any]       ,snake_case_						:						str		)     ->int:



    '''simple docstring'''







    from accelerate.utils import DummyOptim, DummyScheduler

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

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

    __A					:		Optional[Any]						=  None
    if "scheduler" in config:
        __A					:		List[str]						=  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.'''		)
        __A					:		Any						=  trainer.create_scheduler(num_training_steps=snake_case_       ,optimizer=snake_case_		)

    return optimizer, lr_scheduler







def 	__lowercase  (     snake_case_						:						List[str]       ,snake_case_						:						List[str]       ,snake_case_						:						int=False		)     ->int:



    '''simple docstring'''







    from deepspeed.utils import logger as ds_logger

    __A					:		int						=  trainer.model
    __A					:		int						=  trainer.args

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



    '''simple docstring'''







    import glob

    __A					:		List[Any]						=  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
        __A				, __A					:		Dict						=  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}"""		)











"""simple docstring"""



import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


a_				=    """
Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.

In March 2021, Hugging Face raised $40 million in a Series B funding round.[3]

On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]
"""






class 	__snake_case      (	unittest.TestCase							,							SCREAMING_SNAKE_CASE__  ):





 """simple docstring"""


 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		Optional[Any]						=  load_tool('''summarization'''						)
     self.tool.setup()
     __A					:		Tuple						=  load_tool('''summarization'''      ,		remote=__lowerCamelCase						)

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		List[str]						=  self.tool(__lowerCamelCase						)
     self.assertEqual(
         __lowerCamelCase      ,		'''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''      ,		)

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		int						=  self.remote_tool(__lowerCamelCase						)
     self.assertEqual(
         __lowerCamelCase      ,		'''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''      ,		)

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		List[Any]						=  self.tool(text=__lowerCamelCase						)
     self.assertEqual(
         __lowerCamelCase      ,		'''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''      ,		)







 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		Optional[int]						=  self.remote_tool(text=__lowerCamelCase						)
     self.assertEqual(
         __lowerCamelCase      ,		'''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''      ,		)




