




# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings("ignore", category=UserWarning, module="torch.optim.lr_scheduler")




class 		__A   :

						def __init__(		self :Any     ,			__snake_case :Optional[Any]     ,			__snake_case :Optional[int]     ,			__snake_case :bool = True     ,			__snake_case :bool = False			):
										'''simple docstring'''

										__magic_name__						:						Union[str, Any]										=scheduler
										__magic_name__						:						Optional[Any]										=optimizers if isinstance(__snake_case     ,			(list, tuple)			) else [optimizers]
										__magic_name__						:						Any										=split_batches
										__magic_name__						:						List[str]										=step_with_optimizer
										__magic_name__						:						Dict										=GradientState()

						def      A__    (		self :List[str]     ,			*__snake_case :Optional[Any]     ,			**__snake_case :Any			):
										'''simple docstring'''

										if not self.step_with_optimizer:
														# No link between scheduler and optimizer -> just step
														self.scheduler.step(*__snake_case     ,			**__snake_case			)
														return

										# Otherwise, first make sure the optimizer was stepped.
										if not self.gradient_state.sync_gradients:
														if self.gradient_state.adjust_scheduler:
																		self.scheduler._step_count += 1
														return

										for opt in self.optimizers:
														if opt.step_was_skipped:
																		return
										if self.split_batches:
														# Split batches -> the training dataloader batch size is not changed so one step per training step
														self.scheduler.step(*__snake_case     ,			**__snake_case			)
										else:
														# Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
														# num_processes steps per training step
														__magic_name__						:						str										=AcceleratorState().num_processes
														for _ in range(__snake_case			):
																		# Special case when using OneCycle and `drop_last` was not used
																		if hasattr(self.scheduler     ,			"""total_steps"""			):
																						if self.scheduler._step_count <= self.scheduler.total_steps:
																										self.scheduler.step(*__snake_case     ,			**__snake_case			)
																		else:
																						self.scheduler.step(*__snake_case     ,			**__snake_case			)

						def      A__    (		self :Union[str, Any]			):
										'''simple docstring'''

										return self.scheduler.get_last_lr()

						def      A__    (		self :Tuple			):
										'''simple docstring'''

										return self.scheduler.state_dict()

						def      A__    (		self :Union[str, Any]     ,			__snake_case :Dict			):
										'''simple docstring'''

										self.scheduler.load_state_dict(__snake_case			)

						def      A__    (		self :Tuple			):
										'''simple docstring'''

										return self.scheduler.get_lr()


						def      A__    (		self :Any     ,			*__snake_case :int     ,			**__snake_case :Union[str, Any]			):
										'''simple docstring'''

										return self.scheduler.print_lr(*__snake_case     ,			**__snake_case			)






"""simple docstring"""

from ...configuration_utils import PretrainedConfig
from ...utils import logging


SCREAMING_SNAKE_CASE_            =    logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_            =    {
    '''tiiuae/falcon-40b''': '''https://huggingface.co/tiiuae/falcon-40b/resolve/main/config.json''',
    '''tiiuae/falcon-7b''': '''https://huggingface.co/tiiuae/falcon-7b/resolve/main/config.json''',
}
class   _UpperCAmelCase				(		SCREAMING_SNAKE_CASE_				):
      __SCREAMING_SNAKE_CASE :				Optional[Any]				   =						"falcon"
      __SCREAMING_SNAKE_CASE :				List[Any]				   =						["past_key_values"]




      def __init__(				self ,							lowercase_=6_5_0_2_4 ,							lowercase_=4_5_4_4 ,							lowercase_=3_2 ,							lowercase_=7_1 ,							lowercase_=1E-5 ,							lowercase_=0.0_2 ,							lowercase_=True ,							lowercase_=0.0 ,							lowercase_=0.0 ,							lowercase_=None ,							lowercase_=False ,							lowercase_=False ,							lowercase_=True ,							lowercase_=True ,							lowercase_=False ,							lowercase_=1_1 ,							lowercase_=1_1 ,							**lowercase_ ,							)    ->       Any:
          UpperCAmelCase          =       vocab_size
          # Backward compatibility with n_embed kwarg
          UpperCAmelCase          =       kwargs.pop('n_embed' ,							lowercase_				)
          UpperCAmelCase          =       hidden_size if n_embed is None else n_embed
          UpperCAmelCase          =       num_hidden_layers
          UpperCAmelCase          =       num_attention_heads
          UpperCAmelCase          =       layer_norm_epsilon
          UpperCAmelCase          =       initializer_range
          UpperCAmelCase          =       use_cache
          UpperCAmelCase          =       hidden_dropout
          UpperCAmelCase          =       attention_dropout

          UpperCAmelCase          =       bos_token_id
          UpperCAmelCase          =       eos_token_id
          UpperCAmelCase          =       num_attention_heads if num_kv_heads is None else num_kv_heads
          UpperCAmelCase          =       alibi
          UpperCAmelCase          =       new_decoder_architecture
          UpperCAmelCase          =       multi_query  # Ignored when new_decoder_architecture is True
          UpperCAmelCase          =       parallel_attn
          UpperCAmelCase          =       bias

          super().__init__(bos_token_id=lowercase_ ,							eos_token_id=lowercase_ ,							**lowercase_				)




      @property
      def      a_   (				self				)    ->       List[Any]:
          return self.hidden_size // self.num_attention_heads




      @property
      def      a_   (				self				)    ->       List[str]:
          return not self.alibi



