






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


lowerCamelCase					    =	logging.get_logger(__name__)

lowerCamelCase					    =	{
    '''facebook/opt-125m''': '''https://huggingface.co/facebook/opt-125m/blob/main/config.json''',
    '''facebook/opt-350m''': '''https://huggingface.co/facebook/opt-350m/blob/main/config.json''',
    '''facebook/opt-1.3b''': '''https://huggingface.co/facebook/opt-1.3b/blob/main/config.json''',
    '''facebook/opt-2.7b''': '''https://huggingface.co/facebook/opt-2.7b/blob/main/config.json''',
    '''facebook/opt-6.7b''': '''https://huggingface.co/facebook/opt-6.7b/blob/main/config.json''',
    '''facebook/opt-13b''': '''https://huggingface.co/facebook/opt-13b/blob/main/config.json''',
}

class        snake_case_				(							lowercase_   ):







   """simple docstring"""



   __lowerCAmelCase  =    """opt"""
   __lowerCAmelCase  =    ["""past_key_values"""]

   def __init__(     self:						int ,							snake_case__:						Dict=50_272 ,							snake_case__:						Optional[Any]=768 ,							snake_case__:						Any=12 ,							snake_case__:						int=3_072 ,							snake_case__:						str=2_048 ,							snake_case__:						int=True ,							snake_case__:						Any=False ,							snake_case__:						List[Any]=None ,							snake_case__:						List[str]=0.1 ,							snake_case__:						Any=0.0 ,							snake_case__:						str=12 ,							snake_case__:						Optional[Any]="relu" ,							snake_case__:						Dict=0.0 ,							snake_case__:						int=0.02 ,							snake_case__:						List[Any]=True ,							snake_case__:						Optional[int]=1 ,							snake_case__:						List[str]=2 ,							snake_case__:						Optional[Any]=2 ,							snake_case__:						str=True ,							snake_case__:						Tuple=True ,							**snake_case__:						Optional[Any] ,							)	->		Any:


       """simple docstring"""



       super().__init__(
           pad_token_id=snake_case__ ,							bos_token_id=snake_case__ ,							eos_token_id=snake_case__ ,							**snake_case__ ,							)
       _lowercase     			=						vocab_size
       _lowercase     			=						max_position_embeddings
       _lowercase     			=						num_attention_heads
       _lowercase     			=						word_embed_proj_dim if word_embed_proj_dim is not None else hidden_size
       _lowercase     			=						ffn_dim
       _lowercase     			=						hidden_size
       _lowercase     			=						num_hidden_layers
       _lowercase     			=						dropout
       _lowercase     			=						attention_dropout
       _lowercase     			=						activation_function
       _lowercase     			=						init_std
       _lowercase     			=						layerdrop
       _lowercase     			=						use_cache
       _lowercase     			=						do_layer_norm_before
       # We keep these variables at `True` for backward compatibility.
       _lowercase     			=						enable_bias
       _lowercase     			=						layer_norm_elementwise_affine

       # Note that the only purpose of `_remove_final_layer_norm` is to keep backward compatibility
       # with checkpoints that have been fine-tuned before transformers v4.20.1
       # see https://github.com/facebookresearch/metaseq/pull/164
       _lowercase     			=						_remove_final_layer_norm








import torch

from accelerate import Accelerator, DistributedType

class        snake_case_				:







   """simple docstring"""



   def __enter__(     self:						int						)	->		Tuple:


       """simple docstring"""



       if self.enabled:
           _lowercase     			=						self.model.no_sync()
           self.model_sync_obj.__enter__()

       return self

   def __exit__(     self:						str ,							snake_case__:						Tuple ,							snake_case__:						int ,							snake_case__:						Tuple						)	->		Any:


       """simple docstring"""



       if self.enabled:
           # Average all models on exit
           self._sync_and_avg_model_params()
           self.model_sync_obj.__exit__(snake_case__ ,							snake_case__ ,							snake_case__						)

   def __init__(     self:						int ,							snake_case__:						Accelerator ,							snake_case__:						torch.nn.Module ,							snake_case__:						int ,							snake_case__:						bool = True						)	->		Optional[Any]:


       """simple docstring"""



       if accelerator.distributed_type not in [
           DistributedType.NO,
           DistributedType.MULTI_CPU,
           DistributedType.MULTI_GPU,
       ]:
           raise NotImplementedError('LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'						)
       _lowercase     			=						enabled and accelerator.distributed_type != DistributedType.NO
       _lowercase     			=						0
       if self.enabled:
           _lowercase     			=						accelerator
           _lowercase     			=						model
           _lowercase     			=						local_sgd_steps

   def       SCREAMING_SNAKE_CASE__					(     self:						str						)	->		List[Any]:


       """simple docstring"""



       self.num_steps += 1
       if not self.enabled:
           return

       if self.num_steps % self.local_sgd_steps == 0:
           self._sync_and_avg_model_params()

   def       SCREAMING_SNAKE_CASE__					(     self:						Optional[int]						)	->		Dict:


       """simple docstring"""



       self.accelerator.wait_for_everyone()
       with self.accelerator.autocast():
           for param in self.model.parameters():
               _lowercase     			=						self.accelerator.reduce(param.data ,							reduction='mean'						)
