





import torch

from accelerate import Accelerator, DistributedType



class 				__snake_case  :


    '''simple docstring'''







    def __enter__(					self    ):





        '''simple docstring'''

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

        return self



    def __exit__(					self      ,							A_      ,							A_      ,							A_    ):





        '''simple docstring'''

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



    def __init__(					self      ,							A_      ,							A_      ,							A_      ,							A_ = True    ):





        '''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)'''    )
        SCREAMING_SNAKE_CASE__       							= enabled and accelerator.distributed_type != DistributedType.NO
        SCREAMING_SNAKE_CASE__       							= 0
        if self.enabled:
            SCREAMING_SNAKE_CASE__       							= accelerator
            SCREAMING_SNAKE_CASE__       							= model
            SCREAMING_SNAKE_CASE__       							= local_sgd_steps



    def    lowercase_				(					self    ):





        '''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    lowercase_				(					self    ):





        '''simple docstring'''

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







import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
  import transformer_engine.pytorch as te


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_=True   ,       lowerCAmelCase_=True   ,       lowerCAmelCase_=True  )							->	Tuple:
    if not is_fpa_available():
        raise ImportError('''Using `convert_model` requires transformer_engine to be installed.'''  )
    for name, module in model.named_children():
        if isinstance(lowerCAmelCase_   ,       nn.Linear  ) and to_transformer_engine and _convert_linear:
            # Return early if the linear layer weights are not multiples of 16
            if any(p % 1_6 != 0 for p in module.weight.shape  ):
                return
            SCREAMING_SNAKE_CASE__       							= module.bias is not None
            SCREAMING_SNAKE_CASE__       							= te.Linear(
                module.in_features   ,       module.out_features   ,       bias=lowerCAmelCase_   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            if has_bias:
                SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        elif isinstance(lowerCAmelCase_   ,       nn.LayerNorm  ) and to_transformer_engine and _convert_ln:
            SCREAMING_SNAKE_CASE__       							= te.LayerNorm(module.normalized_shape[0]   ,       eps=module.eps   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        elif isinstance(lowerCAmelCase_   ,       te.Linear  ) and not to_transformer_engine and _convert_linear:
            SCREAMING_SNAKE_CASE__       							= module.bias is not None
            SCREAMING_SNAKE_CASE__       							= nn.Linear(
                module.in_features   ,       module.out_features   ,       bias=lowerCAmelCase_   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            if has_bias:
                SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        elif isinstance(lowerCAmelCase_   ,       te.LayerNorm  ) and not to_transformer_engine and _convert_ln:
            SCREAMING_SNAKE_CASE__       							= nn.LayerNorm(module.normalized_shape[0]   ,       eps=module.eps   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        else:
            convert_model(
                lowerCAmelCase_   ,       to_transformer_engine=lowerCAmelCase_   ,       _convert_linear=_convert_linear   ,       _convert_ln=_convert_ln   ,       )



def        __snake_case       (						lowerCAmelCase_  )							->	Union[str, Any]:
    if not is_fpa_available():
        raise ImportError('''Using `has_transformer_engine_layers` requires transformer_engine to be installed.'''  )
    for m in model.modules():
        if isinstance(lowerCAmelCase_   ,       (te.LayerNorm, te.Linear)  ):
            return True
    return False
