



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)

__UpperCAmelCase	       =	'''bert-base-cased'''
__UpperCAmelCase	       =	'''fp16'''
__UpperCAmelCase	       =	'''bf16'''
__UpperCAmelCase	       =	[FPaa, BFaa]



@require_fsdp
@require_cuda
class     lowerCAmelCase_		(  a__							):
       def  snake_case_	(      self	)					->		Optional[int]:
          super().setUp()

          UpperCamelCase   :  Optional[Any]						     =     dict(
              ACCELERATE_USE_FSDP='true',      MASTER_ADDR='localhost',      MASTER_PORT='10999',      RANK='0',      LOCAL_RANK='0',      WORLD_SIZE='1',      )
       def  snake_case_	(      self	)					->		List[str]:
          from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

          for i, strategy in enumerate(SCREAMING_SNAKE_CASE_	):
             UpperCamelCase   :  Union[str, Any]						     =     self.dist_env.copy()
             UpperCamelCase   :  Optional[Any]						     =     F"""{i + 1}"""
             UpperCamelCase   :  Dict						     =     strategy
             with mockenv_context(**SCREAMING_SNAKE_CASE_	):
                UpperCamelCase   :  Optional[Any]						     =     FullyShardedDataParallelPlugin()
                self.assertEqual(fsdp_plugin.sharding_strategy,      ShardingStrategy(i + 1	)	)
       def  snake_case_	(      self	)					->		Any:
          from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

          for i, prefetch_policy in enumerate(SCREAMING_SNAKE_CASE_	):
             UpperCamelCase   :  Dict						     =     self.dist_env.copy()
             UpperCamelCase   :  List[str]						     =     prefetch_policy
             with mockenv_context(**SCREAMING_SNAKE_CASE_	):
                UpperCamelCase   :  Tuple						     =     FullyShardedDataParallelPlugin()
                if prefetch_policy == "NO_PREFETCH":
                   self.assertIsNone(fsdp_plugin.backward_prefetch	)
                else:
                   self.assertEqual(fsdp_plugin.backward_prefetch,      BackwardPrefetch(i + 1	)	)
       def  snake_case_	(      self	)					->		Tuple:
          from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

          for i, state_dict_type in enumerate(SCREAMING_SNAKE_CASE_	):
             UpperCamelCase   :  str						     =     self.dist_env.copy()
             UpperCamelCase   :  Union[str, Any]						     =     state_dict_type
             with mockenv_context(**SCREAMING_SNAKE_CASE_	):
                UpperCamelCase   :  Any						     =     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  snake_case_	(      self	)					->		Any:
          UpperCamelCase   :  int						     =     AutoModel.from_pretrained(SCREAMING_SNAKE_CASE_	)
          for policy in FSDP_AUTO_WRAP_POLICY:
             UpperCamelCase   :  Optional[Any]						     =     self.dist_env.copy()
             UpperCamelCase   :  List[Any]						     =     policy
             if policy == "TRANSFORMER_BASED_WRAP":
                UpperCamelCase   :  Optional[int]						     =     'BertLayer'
             elif policy == "SIZE_BASED_WRAP":
                UpperCamelCase   :  List[str]						     =     '2000'
             with mockenv_context(**SCREAMING_SNAKE_CASE_	):
                UpperCamelCase   :  Union[str, Any]						     =     FullyShardedDataParallelPlugin()
                fsdp_plugin.set_auto_wrap_policy(SCREAMING_SNAKE_CASE_	)
                if policy == "NO_WRAP":
                   self.assertIsNone(fsdp_plugin.auto_wrap_policy	)
                else:
                   self.assertIsNotNone(fsdp_plugin.auto_wrap_policy	)

          UpperCamelCase   :  Optional[Any]						     =     self.dist_env.copy()
          UpperCamelCase   :  Dict						     =     'TRANSFORMER_BASED_WRAP'
          UpperCamelCase   :  Optional[int]						     =     'T5Layer'
          with mockenv_context(**SCREAMING_SNAKE_CASE_	):
             UpperCamelCase   :  Optional[Any]						     =     FullyShardedDataParallelPlugin()
             with self.assertRaises(SCREAMING_SNAKE_CASE_	) as cm:
                fsdp_plugin.set_auto_wrap_policy(SCREAMING_SNAKE_CASE_	)
             self.assertTrue('Could not find the transformer layer class to wrap in the model.' in str(cm.exception	)	)

          UpperCamelCase   :  Any						     =     self.dist_env.copy()
          UpperCamelCase   :  Tuple						     =     'SIZE_BASED_WRAP'
          UpperCamelCase   :  Optional[Any]						     =     '0'
          with mockenv_context(**SCREAMING_SNAKE_CASE_	):
             UpperCamelCase   :  str						     =     FullyShardedDataParallelPlugin()
             fsdp_plugin.set_auto_wrap_policy(SCREAMING_SNAKE_CASE_	)
             self.assertIsNone(fsdp_plugin.auto_wrap_policy	)
       def  snake_case_	(      self	)					->		Tuple:
          from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
          from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

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






       def  snake_case_	(      self	)					->		int:
          from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

          for flag in [True, False]:
             UpperCamelCase   :  Any						     =     self.dist_env.copy()
             UpperCamelCase   :  Union[str, Any]						     =     str(SCREAMING_SNAKE_CASE_	).lower()
             with mockenv_context(**SCREAMING_SNAKE_CASE_	):
                UpperCamelCase   :  List[str]						     =     FullyShardedDataParallelPlugin()
                self.assertEqual(fsdp_plugin.cpu_offload,      CPUOffload(offload_params=SCREAMING_SNAKE_CASE_	)	)






@require_fsdp
@require_multi_gpu
@slow
class     lowerCAmelCase_		(  a__							):
       def  snake_case_	(      self	)					->		Dict:
          super().setUp()
          UpperCamelCase   :  List[str]						     =     0.82
          UpperCamelCase   :  List[Any]						     =     [
              'fsdp_shard_grad_op_transformer_based_wrap',
              'fsdp_full_shard_transformer_based_wrap',
          ]
          UpperCamelCase   :  List[Any]						     =     {
              'multi_gpu_fp16': 3200,
              'fsdp_shard_grad_op_transformer_based_wrap_fp16': 2000,
              'fsdp_full_shard_transformer_based_wrap_fp16': 1900,
              # 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
          }
          UpperCamelCase   :  Dict						     =     160
          UpperCamelCase   :  Optional[int]						     =     160

          UpperCamelCase   :  Optional[int]						     =     inspect.getfile(accelerate.test_utils	)
          UpperCamelCase   :  Optional[int]						     =     os.path.sep.join(mod_file.split(os.path.sep	)[:-1] + ['scripts', 'external_deps']	)
       def  snake_case_	(      self	)					->		List[Any]:
          UpperCamelCase   :  Optional[int]						     =     os.path.join(self.test_scripts_folder,      'test_performance.py'	)
          UpperCamelCase   :  str						     =     ['accelerate', 'launch', '--num_processes=2', '--num_machines=1', '--machine_rank=0', '--use_fsdp']
          for config in self.performance_configs:
             UpperCamelCase   :  List[Any]						     =     cmd.copy()
             for i, strategy in enumerate(SCREAMING_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(SCREAMING_SNAKE_CASE_,      env=os.environ.copy()	)
       def  snake_case_	(      self	)					->		Optional[int]:
          UpperCamelCase   :  str						     =     os.path.join(self.test_scripts_folder,      'test_checkpointing.py'	)
          UpperCamelCase   :  List[Any]						     =     [
              '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(SCREAMING_SNAKE_CASE_	):
             UpperCamelCase   :  int						     =     cmd.copy()
             cmd_config.append(F"""--fsdp_sharding_strategy={i+1}"""	)
             if strategy != "FULL_SHARD":
                continue
             UpperCamelCase   :  List[str]						     =     len(SCREAMING_SNAKE_CASE_	)
             for state_dict_type in FSDP_STATE_DICT_TYPE:
                UpperCamelCase   :  Tuple						     =     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(SCREAMING_SNAKE_CASE_,      env=os.environ.copy()	)

                UpperCamelCase   :  Tuple						     =     cmd_config[:-1]
                UpperCamelCase   :  Optional[Any]						     =     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(SCREAMING_SNAKE_CASE_,      env=os.environ.copy()	)






       def  snake_case_	(      self	)					->		List[Any]:
          UpperCamelCase   :  Optional[int]						     =     os.path.join(self.test_scripts_folder,      'test_peak_memory_usage.py'	)
          UpperCamelCase   :  Optional[int]						     =     [
              '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():
             UpperCamelCase   :  int						     =     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(SCREAMING_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(SCREAMING_SNAKE_CASE_,      env=os.environ.copy()	)








# tests directory-specific settings - this file is run automatically
# by pytest before any tests are run

import sys
import warnings
from os.path import abspath, dirname, join


# allow having multiple repository checkouts and not needing to remember to rerun
# 'pip install -e .[dev]' when switching between checkouts and running tests.
__UpperCAmelCase	       =	abspath(join(dirname(dirname(dirname(__file__))), '''src'''))
sys.path.insert(1, git_repo_path)


# silence FutureWarning warnings in tests since often we can't act on them until
# they become normal warnings - i.e. the tests still need to test the current functionality
warnings.simplefilter(action='''ignore''', category=FutureWarning)
def 		UpperCamelCase     (       snake_case__	:   Dict     )     ->					List[str]:
   from transformers.testing_utils import pytest_addoption_shared

   pytest_addoption_shared(snake_case__     )


def 		UpperCamelCase     (       snake_case__	:   List[Any]     )     ->					Dict:
   from transformers.testing_utils import pytest_terminal_summary_main

   UpperCamelCase   :  Optional[int]						     =     terminalreporter.config.getoption('--make-reports'     )
   if make_reports:
      pytest_terminal_summary_main(snake_case__	,					id=snake_case__     )



