





import inspect
import os

import torch
from transformers import AutoModel
from transformers.testing_utils import mockenv_context
from transformers.trainer_utils import set_seed

import accelerate
from accelerate.accelerator import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils.testing import (
    AccelerateTestCase,
    TempDirTestCase,
    execute_subprocess_async,
    require_cuda,
    require_fsdp,
    require_multi_gpu,
    slow,
)
from accelerate.utils.constants import (
    FSDP_AUTO_WRAP_POLICY,
    FSDP_BACKWARD_PREFETCH,
    FSDP_SHARDING_STRATEGY,
    FSDP_STATE_DICT_TYPE,
)
from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin
from accelerate.utils.other import patch_environment


set_seed(42)

a__      :     Dict         =		'bert-base-cased'
a__      :     List[str]         =		'fp16'
a__      :     Optional[Any]         =		'bf16'
a__      :     Dict         =		[FPaa, BFaa]





@require_fsdp
@require_cuda
class UpperCAmelCase_				( _SCREAMING_SNAKE_CASE ):


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       super().setUp()

       A_          =       dict(
           ACCELERATE_USE_FSDP='''true''' ,MASTER_ADDR='''localhost''' ,MASTER_PORT='''10999''' ,RANK='''0''' ,LOCAL_RANK='''0''' ,WORLD_SIZE='''1''' ,)


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

       for i, strategy in enumerate(__snake_case ):
            A_          =       self.dist_env.copy()
            A_          =       f'{i + 1}'
            A_          =       strategy
            with mockenv_context(**__snake_case ):
                 A_          =       FullyShardedDataParallelPlugin()
                 self.assertEqual(fsdp_plugin.sharding_strategy ,ShardingStrategy(i + 1 ) )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

       for i, prefetch_policy in enumerate(__snake_case ):
            A_          =       self.dist_env.copy()
            A_          =       prefetch_policy
            with mockenv_context(**__snake_case ):
                 A_          =       FullyShardedDataParallelPlugin()
                 if prefetch_policy == "NO_PREFETCH":
                      self.assertIsNone(fsdp_plugin.backward_prefetch )
                 else:
                      self.assertEqual(fsdp_plugin.backward_prefetch ,BackwardPrefetch(i + 1 ) )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

       for i, state_dict_type in enumerate(__snake_case ):
            A_          =       self.dist_env.copy()
            A_          =       state_dict_type
            with mockenv_context(**__snake_case ):
                 A_          =       FullyShardedDataParallelPlugin()
                 self.assertEqual(fsdp_plugin.state_dict_type ,StateDictType(i + 1 ) )
                 if state_dict_type == "FULL_STATE_DICT":
                      self.assertTrue(fsdp_plugin.state_dict_config.offload_to_cpu )
                      self.assertTrue(fsdp_plugin.state_dict_config.ranka_only )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       AutoModel.from_pretrained(__snake_case )
       for policy in FSDP_AUTO_WRAP_POLICY:
            A_          =       self.dist_env.copy()
            A_          =       policy
            if policy == "TRANSFORMER_BASED_WRAP":
                 A_          =       '''BertLayer'''
            elif policy == "SIZE_BASED_WRAP":
                 A_          =       '''2000'''
            with mockenv_context(**__snake_case ):
                 A_          =       FullyShardedDataParallelPlugin()
                 fsdp_plugin.set_auto_wrap_policy(__snake_case )
                 if policy == "NO_WRAP":
                      self.assertIsNone(fsdp_plugin.auto_wrap_policy )
                 else:
                      self.assertIsNotNone(fsdp_plugin.auto_wrap_policy )

       A_          =       self.dist_env.copy()
       A_          =       '''TRANSFORMER_BASED_WRAP'''
       A_          =       '''T5Layer'''
       with mockenv_context(**__snake_case ):
            A_          =       FullyShardedDataParallelPlugin()
            with self.assertRaises(__snake_case ) as cm:
                 fsdp_plugin.set_auto_wrap_policy(__snake_case )
            self.assertTrue('''Could not find the transformer layer class to wrap in the model.''' in str(cm.exception ) )

       A_          =       self.dist_env.copy()
       A_          =       '''SIZE_BASED_WRAP'''
       A_          =       '''0'''
       with mockenv_context(**__snake_case ):
            A_          =       FullyShardedDataParallelPlugin()
            fsdp_plugin.set_auto_wrap_policy(__snake_case )
            self.assertIsNone(fsdp_plugin.auto_wrap_policy )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
       from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

       for mp_dtype in dtypes:
            A_          =       self.dist_env.copy()
            A_          =       mp_dtype
            with mockenv_context(**__snake_case ):
                 A_          =       Accelerator()
                 if mp_dtype == "fp16":
                      A_          =       torch.floataa
                 elif mp_dtype == "bf16":
                      A_          =       torch.bfloataa
                 A_          =       MixedPrecision(param_dtype=__snake_case ,reduce_dtype=__snake_case ,buffer_dtype=__snake_case )
                 self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy ,__snake_case )
                 if mp_dtype == FPaa:
                      self.assertTrue(isinstance(accelerator.scaler ,__snake_case ) )
                 elif mp_dtype == BFaa:
                      self.assertIsNone(accelerator.scaler )
                 AcceleratorState._reset_state(__snake_case )




  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

       for flag in [True, False]:
            A_          =       self.dist_env.copy()
            A_          =       str(__snake_case ).lower()
            with mockenv_context(**__snake_case ):
                 A_          =       FullyShardedDataParallelPlugin()
                 self.assertEqual(fsdp_plugin.cpu_offload ,CPUOffload(offload_params=__snake_case ) )





@require_fsdp
@require_multi_gpu
@slow
class UpperCAmelCase_				( _SCREAMING_SNAKE_CASE ):


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       super().setUp()
       A_          =       0.82
       A_          =       [
           '''fsdp_shard_grad_op_transformer_based_wrap''',
           '''fsdp_full_shard_transformer_based_wrap''',
       ]
       A_          =       {
           '''multi_gpu_fp16''': 3_2_0_0,
           '''fsdp_shard_grad_op_transformer_based_wrap_fp16''': 2_0_0_0,
           '''fsdp_full_shard_transformer_based_wrap_fp16''': 1_9_0_0,
           # Disabling below test as it overwhelms the RAM memory usage
           # on CI self-hosted runner leading to tests getting killed.
           # "fsdp_full_shard_cpu_offload_transformer_based_wrap_fp32": 1500,  # fp16 was leading to indefinite hang
       }
       A_          =       1_6_0
       A_          =       1_6_0

       A_          =       inspect.getfile(accelerate.test_utils )
       A_          =       os.path.sep.join(mod_file.split(os.path.sep )[:-1] + ['''scripts''', '''external_deps'''] )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       os.path.join(self.test_scripts_folder ,'''test_performance.py''' )
       A_          =       ['''accelerate''', '''launch''', '''--num_processes=2''', '''--num_machines=1''', '''--machine_rank=0''', '''--use_fsdp''']
       for config in self.performance_configs:
            A_          =       cmd.copy()
            for i, strategy in enumerate(__snake_case ):
                 if strategy.lower() in config:
                      cmd_config.append(f'--fsdp_sharding_strategy={i+1}' )
                      break

            if "fp32" in config:
                 cmd_config.append('''--mixed_precision=no''' )
            else:
                 cmd_config.append('''--mixed_precision=fp16''' )

            if "cpu_offload" in config:
                 cmd_config.append('''--fsdp_offload_params=True''' )

            for policy in FSDP_AUTO_WRAP_POLICY:
                 if policy.lower() in config:
                      cmd_config.append(f'--fsdp_auto_wrap_policy={policy}' )
                      break

            if policy == "TRANSFORMER_BASED_WRAP":
                 cmd_config.append('''--fsdp_transformer_layer_cls_to_wrap=BertLayer''' )
            elif policy == "SIZE_BASED_WRAP":
                 cmd_config.append('''--fsdp_min_num_params=2000''' )

            cmd_config.extend(
                [
                    self.test_file_path,
                    f'--output_dir={self.tmpdir}',
                    f'--performance_lower_bound={self.performance_lower_bound}',
                ] )
            with patch_environment(omp_num_threads=1 ):
                 execute_subprocess_async(__snake_case ,env=os.environ.copy() )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       os.path.join(self.test_scripts_folder ,'''test_checkpointing.py''' )
       A_          =       [
           '''accelerate''',
           '''launch''',
           '''--num_processes=2''',
           '''--num_machines=1''',
           '''--machine_rank=0''',
           '''--use_fsdp''',
           '''--mixed_precision=fp16''',
           '''--fsdp_transformer_layer_cls_to_wrap=BertLayer''',
       ]

       for i, strategy in enumerate(__snake_case ):
            A_          =       cmd.copy()
            cmd_config.append(f'--fsdp_sharding_strategy={i+1}' )
            if strategy != "FULL_SHARD":
                 continue
            A_          =       len(__snake_case )
            for state_dict_type in FSDP_STATE_DICT_TYPE:
                 A_          =       cmd_config[:state_dict_config_index]
                 cmd_config.append(f'--fsdp_state_dict_type={state_dict_type}' )
                 cmd_config.extend(
                     [
                         self.test_file_path,
                         f'--output_dir={self.tmpdir}',
                         '''--partial_train_epoch=1''',
                     ] )
                 with patch_environment(omp_num_threads=1 ):
                      execute_subprocess_async(__snake_case ,env=os.environ.copy() )

                 A_          =       cmd_config[:-1]
                 A_          =       os.path.join(self.tmpdir ,'''epoch_0''' )
                 cmd_config.extend(
                     [
                         f'--resume_from_checkpoint={resume_from_checkpoint}',
                     ] )
                 with patch_environment(omp_num_threads=1 ):
                      execute_subprocess_async(__snake_case ,env=os.environ.copy() )




  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       os.path.join(self.test_scripts_folder ,'''test_peak_memory_usage.py''' )
       A_          =       [
           '''accelerate''',
           '''launch''',
           '''--num_processes=2''',
           '''--num_machines=1''',
           '''--machine_rank=0''',
       ]
       for spec, peak_mem_upper_bound in self.peak_memory_usage_upper_bound.items():
            A_          =       cmd.copy()
            if "fp16" in spec:
                 cmd_config.extend(['''--mixed_precision=fp16'''] )
            else:
                 cmd_config.extend(['''--mixed_precision=no'''] )

            if "multi_gpu" in spec:
                 continue
            else:
                 cmd_config.extend(['''--use_fsdp'''] )
                 for i, strategy in enumerate(__snake_case ):
                      if strategy.lower() in spec:
                           cmd_config.append(f'--fsdp_sharding_strategy={i+1}' )
                           break

                 if "cpu_offload" in spec:
                      cmd_config.append('''--fsdp_offload_params=True''' )

                 for policy in FSDP_AUTO_WRAP_POLICY:
                      if policy.lower() in spec:
                           cmd_config.append(f'--fsdp_auto_wrap_policy={policy}' )
                           break

                 if policy == "TRANSFORMER_BASED_WRAP":
                      cmd_config.append('''--fsdp_transformer_layer_cls_to_wrap=BertLayer''' )
                 elif policy == "SIZE_BASED_WRAP":
                      cmd_config.append('''--fsdp_min_num_params=2000''' )

            cmd_config.extend(
                [
                    self.test_file_path,
                    f'--output_dir={self.tmpdir}',
                    f'--peak_memory_upper_bound={peak_mem_upper_bound}',
                    f'--n_train={self.n_train}',
                    f'--n_val={self.n_val}',
                ] )
            with patch_environment(omp_num_threads=1 ):
                 execute_subprocess_async(__snake_case ,env=os.environ.copy() )










from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


a__      :     str         =		{
    'configuration_roformer': ['ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP', 'RoFormerConfig', 'RoFormerOnnxConfig'],
    'tokenization_roformer': ['RoFormerTokenizer'],
}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      a__      :     List[Any]         =		['RoFormerTokenizerFast']

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      a__      :     List[str]         =		[
          'ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
          'RoFormerForCausalLM',
          'RoFormerForMaskedLM',
          'RoFormerForMultipleChoice',
          'RoFormerForQuestionAnswering',
          'RoFormerForSequenceClassification',
          'RoFormerForTokenClassification',
          'RoFormerLayer',
          'RoFormerModel',
          'RoFormerPreTrainedModel',
          'load_tf_weights_in_roformer',
      ]


try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      a__      :     List[Any]         =		[
          'TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
          'TFRoFormerForCausalLM',
          'TFRoFormerForMaskedLM',
          'TFRoFormerForMultipleChoice',
          'TFRoFormerForQuestionAnswering',
          'TFRoFormerForSequenceClassification',
          'TFRoFormerForTokenClassification',
          'TFRoFormerLayer',
          'TFRoFormerModel',
          'TFRoFormerPreTrainedModel',
      ]


try:
      if not is_flax_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      a__      :     Dict         =		[
          'FLAX_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
          'FlaxRoFormerForMaskedLM',
          'FlaxRoFormerForMultipleChoice',
          'FlaxRoFormerForQuestionAnswering',
          'FlaxRoFormerForSequenceClassification',
          'FlaxRoFormerForTokenClassification',
          'FlaxRoFormerModel',
          'FlaxRoFormerPreTrainedModel',
      ]


if TYPE_CHECKING:
      from .configuration_roformer import ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, RoFormerConfig, RoFormerOnnxConfig
      from .tokenization_roformer import RoFormerTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_roformer_fast import RoFormerTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_roformer import (
                ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                RoFormerForCausalLM,
                RoFormerForMaskedLM,
                RoFormerForMultipleChoice,
                RoFormerForQuestionAnswering,
                RoFormerForSequenceClassification,
                RoFormerForTokenClassification,
                RoFormerLayer,
                RoFormerModel,
                RoFormerPreTrainedModel,
                load_tf_weights_in_roformer,
            )

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_roformer import (
                TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFRoFormerForCausalLM,
                TFRoFormerForMaskedLM,
                TFRoFormerForMultipleChoice,
                TFRoFormerForQuestionAnswering,
                TFRoFormerForSequenceClassification,
                TFRoFormerForTokenClassification,
                TFRoFormerLayer,
                TFRoFormerModel,
                TFRoFormerPreTrainedModel,
            )

      try:
            if not is_flax_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_flax_roformer import (
                FLAX_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                FlaxRoFormerForMaskedLM,
                FlaxRoFormerForMultipleChoice,
                FlaxRoFormerForQuestionAnswering,
                FlaxRoFormerForSequenceClassification,
                FlaxRoFormerForTokenClassification,
                FlaxRoFormerModel,
                FlaxRoFormerPreTrainedModel,
            )


else:
      import sys

      a__      :     Union[str, Any]         =		_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



