

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

lowerCamelCase				:   str      				=  logging.get_logger(__name__)


def  _lowerCAmelCase      (		)   ->    List[Any]:
     """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   :			str							,      _a   :			Union[str, Any]				)			->					List[Any]:





       '''simple docstring'''






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







class 							A__					( __lowerCamelCase						):





  def __init__(	self   :			str							,      _a   :			List[Any]				)			->					Tuple:





       '''simple docstring'''






       super().__init__(_a				)
       _SCREAMING_SNAKE_CASE												=None
       _SCREAMING_SNAKE_CASE												=[]





  def 		A						(	self   :			Dict				)			->					Optional[Any]:





       '''simple docstring'''






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





  def 		A						(	self   :			List[Any]							,      _a   :			str				)			->					List[Any]:





       '''simple docstring'''






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





  def 		A						(	self   :			int							,      _a   :			List[str]							,      _a   :			str							,      _a   :			Any=None							,      _a   :			Optional[Any]=True				)			->					Any:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE												=self.find_config_node(_a				)
       if config is None:
            return

       if config.get(_a				) == "auto":
            _SCREAMING_SNAKE_CASE												=hf_val
            return

       if not must_match:
            return

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

  A__					=       partialmethod(__lowerCamelCase  ,     must_match=__lowerCamelCase						)





  def 		A						(	self   :			List[str]							,      _a   :			Any				)			->					Optional[int]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE												=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'							,      _a							,      '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:
            _SCREAMING_SNAKE_CASE												='apex' if args.fpaa_backend == 'apex' else 'amp'
       else:
            _SCREAMING_SNAKE_CASE												=None

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





  def 		A						(	self   :			Any							,      _a   :			Any							,      _a   :			str							,      _a   :			Any				)			->					List[Any]:





       '''simple docstring'''






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

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

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

lowerCamelCase				:   str      				=  None


def  _lowerCAmelCase      (		_UpperCamelCase		:  int       )   ->    List[Any]:
     """simple docstring"""




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


def  _lowerCAmelCase      (		)   ->    Optional[int]:
     """simple docstring"""




     global _hf_deepspeed_config_weak_ref
     _SCREAMING_SNAKE_CASE												=None


def  _lowerCAmelCase      (		)   ->    Dict:
     """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      (		)   ->    Union[str, 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  _lowerCAmelCase      (		_UpperCamelCase		:  int							,				_UpperCamelCase		:  Optional[int]							,				_UpperCamelCase		:  Dict							,				_UpperCamelCase		:  Optional[Any]							,				_UpperCamelCase		:  int       )   ->    int:
     """simple docstring"""




     from accelerate.utils import DummyOptim, DummyScheduler

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

     _SCREAMING_SNAKE_CASE												=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.'       )
          _SCREAMING_SNAKE_CASE												=DummyOptim(params=_UpperCamelCase       )
     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.
          _SCREAMING_SNAKE_CASE												=trainer.create_optimizer()
          # To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
          _SCREAMING_SNAKE_CASE												=True

     _SCREAMING_SNAKE_CASE												=None
     if "scheduler" in config:
          _SCREAMING_SNAKE_CASE												=DummyScheduler(_UpperCamelCase       )
     else:
          if isinstance(_UpperCamelCase							,				_UpperCamelCase       ):
               raise ValueError(
                   'Found `optimizer` configured in the DeepSpeed config, but no `scheduler`. '
                   'Please configure a scheduler in the DeepSpeed config.'       )
          _SCREAMING_SNAKE_CASE												=trainer.create_scheduler(num_training_steps=_UpperCamelCase							,				optimizer=_UpperCamelCase       )

     return optimizer, lr_scheduler


def  _lowerCAmelCase      (		_UpperCamelCase		:  int							,				_UpperCamelCase		:  Optional[int]							,				_UpperCamelCase		:  str=False       )   ->    Optional[int]:
     """simple docstring"""




     from deepspeed.utils import logger as ds_logger

     _SCREAMING_SNAKE_CASE												=trainer.model
     _SCREAMING_SNAKE_CASE												=trainer.args

     _SCREAMING_SNAKE_CASE												=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(_UpperCamelCase							,				_UpperCamelCase							,				_UpperCamelCase       )

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

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

     return optimizer, lr_scheduler







def  _lowerCAmelCase      (		_UpperCamelCase		:  int							,				_UpperCamelCase		:  str       )   ->    int:
     """simple docstring"""




     import glob

     _SCREAMING_SNAKE_CASE												=sorted(glob.glob(f"{checkpoint_path}/global_step*"       )       )

     if len(_UpperCamelCase       ) > 0:
          logger.info(f"Attempting to resume from {checkpoint_path}"       )
          # this magically updates self.optimizer and self.lr_scheduler
          _SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE												=deepspeed_engine.load_checkpoint(
              _UpperCamelCase							,				load_optimizer_states=_UpperCamelCase							,				load_lr_scheduler_states=_UpperCamelCase       )
          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'''
from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase				:   Dict      				=  logging.get_logger(__name__)

lowerCamelCase				:   Dict      				=  {
    "uw-madison/nystromformer-512": "https://huggingface.co/uw-madison/nystromformer-512/resolve/main/config.json",
    # See all Nystromformer models at https://huggingface.co/models?filter=nystromformer
}







class 							A__					( A__						):
  A__					=       'nystromformer'





  def __init__(	self   :			List[str]							,      _a   :			Dict=3_0000							,      _a   :			Dict=768							,      _a   :			Union[str, Any]=12							,      _a   :			Optional[int]=12							,      _a   :			Union[str, Any]=3072							,      _a   :			Optional[Any]="gelu_new"							,      _a   :			Dict=0.1							,      _a   :			Dict=0.1							,      _a   :			Union[str, Any]=510							,      _a   :			Any=2							,      _a   :			List[Any]=64							,      _a   :			int=64							,      _a   :			Optional[Any]=65							,      _a   :			str=False							,      _a   :			List[Any]=0.02							,      _a   :			List[str]=1e-5							,      _a   :			Optional[Any]=1							,      _a   :			Optional[Any]=0							,      _a   :			List[Any]=2							,      **_a   :			List[Any]							,      )			->					List[Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE												=vocab_size
       _SCREAMING_SNAKE_CASE												=max_position_embeddings
       _SCREAMING_SNAKE_CASE												=hidden_size
       _SCREAMING_SNAKE_CASE												=num_hidden_layers
       _SCREAMING_SNAKE_CASE												=num_attention_heads
       _SCREAMING_SNAKE_CASE												=intermediate_size
       _SCREAMING_SNAKE_CASE												=hidden_act
       _SCREAMING_SNAKE_CASE												=hidden_dropout_prob
       _SCREAMING_SNAKE_CASE												=attention_probs_dropout_prob
       _SCREAMING_SNAKE_CASE												=initializer_range
       _SCREAMING_SNAKE_CASE												=type_vocab_size
       _SCREAMING_SNAKE_CASE												=segment_means_seq_len
       _SCREAMING_SNAKE_CASE												=num_landmarks
       _SCREAMING_SNAKE_CASE												=conv_kernel_size
       _SCREAMING_SNAKE_CASE												=inv_coeff_init_option
       _SCREAMING_SNAKE_CASE												=layer_norm_eps
       super().__init__(pad_token_id=_a							,      bos_token_id=_a							,      eos_token_id=_a							,      **_a				)
