
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)

lowercase  :					Dict						    =							"""bert-base-cased"""
lowercase  :					int						    =							"""fp16"""
lowercase  :					Optional[int]						    =							"""bf16"""
lowercase  :					str						    =							[FPaa, BFaa]

@require_fsdp
@require_cuda
class __snake_case		(      lowerCAmelCase     ):
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           super().setUp()

           lowercase							:     str       	=      dict(
               ACCELERATE_USE_FSDP="""true"""		,MASTER_ADDR="""localhost"""		,MASTER_PORT="""10999"""		,RANK="""0"""		,LOCAL_RANK="""0"""		,WORLD_SIZE="""1"""		,)
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

           for i, strategy in enumerate(snake_case      ):
                  lowercase							:     Tuple       	=      self.dist_env.copy()
                  lowercase							:     Dict       	=      f"{i + 1}"
                  lowercase							:     Optional[int]       	=      strategy
                  with mockenv_context(**snake_case      ):
                         lowercase							:     Any       	=      FullyShardedDataParallelPlugin()
                         self.assertEqual(fsdp_plugin.sharding_strategy		,ShardingStrategy(i + 1      )      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

           for i, prefetch_policy in enumerate(snake_case      ):
                  lowercase							:     Dict       	=      self.dist_env.copy()
                  lowercase							:     List[str]       	=      prefetch_policy
                  with mockenv_context(**snake_case      ):
                         lowercase							:     Tuple       	=      FullyShardedDataParallelPlugin()
                         if prefetch_policy == "NO_PREFETCH":
                                self.assertIsNone(fsdp_plugin.backward_prefetch      )
                         else:
                                self.assertEqual(fsdp_plugin.backward_prefetch		,BackwardPrefetch(i + 1      )      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

           for i, state_dict_type in enumerate(snake_case      ):
                  lowercase							:     int       	=      self.dist_env.copy()
                  lowercase							:     List[str]       	=      state_dict_type
                  with mockenv_context(**snake_case      ):
                         lowercase							:     Optional[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       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Optional[Any]       	=      AutoModel.from_pretrained(snake_case      )
           for policy in FSDP_AUTO_WRAP_POLICY:
                  lowercase							:     str       	=      self.dist_env.copy()
                  lowercase							:     Optional[Any]       	=      policy
                  if policy == "TRANSFORMER_BASED_WRAP":
                         lowercase							:     List[Any]       	=      """BertLayer"""
                  elif policy == "SIZE_BASED_WRAP":
                         lowercase							:     List[str]       	=      """2000"""
                  with mockenv_context(**snake_case      ):
                         lowercase							:     Optional[Any]       	=      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      )

           lowercase							:     Union[str, Any]       	=      self.dist_env.copy()
           lowercase							:     Optional[int]       	=      """TRANSFORMER_BASED_WRAP"""
           lowercase							:     Optional[int]       	=      """T5Layer"""
           with mockenv_context(**snake_case      ):
                  lowercase							:     str       	=      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      )      )

           lowercase							:     List[str]       	=      self.dist_env.copy()
           lowercase							:     Tuple       	=      """SIZE_BASED_WRAP"""
           lowercase							:     List[Any]       	=      """0"""
           with mockenv_context(**snake_case      ):
                  lowercase							:     Any       	=      FullyShardedDataParallelPlugin()
                  fsdp_plugin.set_auto_wrap_policy(snake_case      )
                  self.assertIsNone(fsdp_plugin.auto_wrap_policy      )
    def       _SCREAMING_SNAKE_CASE       (  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:
                  lowercase							:     str       	=      self.dist_env.copy()
                  lowercase							:     List[Any]       	=      mp_dtype
                  with mockenv_context(**snake_case      ):
                         lowercase							:     Dict       	=      Accelerator()
                         if mp_dtype == "fp16":
                                lowercase							:     str       	=      torch.floataa
                         elif mp_dtype == "bf16":
                                lowercase							:     Optional[int]       	=      torch.bfloataa
                         lowercase							:     Optional[Any]       	=      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       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

           for flag in [True, False]:
                  lowercase							:     Dict       	=      self.dist_env.copy()
                  lowercase							:     Optional[int]       	=      str(snake_case      ).lower()
                  with mockenv_context(**snake_case      ):
                         lowercase							:     Optional[Any]       	=      FullyShardedDataParallelPlugin()
                         self.assertEqual(fsdp_plugin.cpu_offload		,CPUOffload(offload_params=snake_case      )      )






@require_fsdp
@require_multi_gpu
@slow
class __snake_case		(      lowerCAmelCase     ):
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           super().setUp()
           lowercase							:     Optional[int]       	=      0.82
           lowercase							:     List[str]       	=      [
               """fsdp_shard_grad_op_transformer_based_wrap""",
               """fsdp_full_shard_transformer_based_wrap""",
           ]
           lowercase							:     int       	=      {
               """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
           }
           lowercase							:     Dict       	=      160
           lowercase							:     Optional[Any]       	=      160

           lowercase							:     int       	=      inspect.getfile(accelerate.test_utils      )
           lowercase							:     Any       	=      os.path.sep.join(mod_file.split(os.path.sep      )[:-1] + ["""scripts""", """external_deps"""]      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Dict       	=      os.path.join(self.test_scripts_folder		,"""test_performance.py"""      )
           lowercase							:     int       	=      ["""accelerate""", """launch""", """--num_processes=2""", """--num_machines=1""", """--machine_rank=0""", """--use_fsdp"""]
           for config in self.performance_configs:
                  lowercase							:     Optional[int]       	=      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       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Union[str, Any]       	=      os.path.join(self.test_scripts_folder		,"""test_checkpointing.py"""      )
           lowercase							:     int       	=      [
               """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      ):
                  lowercase							:     Any       	=      cmd.copy()
                  cmd_config.append(f"--fsdp_sharding_strategy={i+1}"      )
                  if strategy != "FULL_SHARD":
                         continue
                  lowercase							:     Union[str, Any]       	=      len(snake_case      )
                  for state_dict_type in FSDP_STATE_DICT_TYPE:
                         lowercase							:     Optional[int]       	=      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()      )

                         lowercase							:     Tuple       	=      cmd_config[:-1]
                         lowercase							:     Tuple       	=      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       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Union[str, Any]       	=      os.path.join(self.test_scripts_folder		,"""test_peak_memory_usage.py"""      )
           lowercase							:     List[Any]       	=      [
               """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():
                  lowercase							:     Dict       	=      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 scipy.stats import pearsonr, spearmanr
from sklearn.metrics import fa_score, matthews_corrcoef

import datasets


lowercase  :					str						    =							"""\
@inproceedings{wang2019glue,
  title={{GLUE}: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding},
  author={Wang, Alex and Singh, Amanpreet and Michael, Julian and Hill, Felix and Levy, Omer and Bowman, Samuel R.},
  note={In the Proceedings of ICLR.},
  year={2019}
}
"""

lowercase  :					Dict						    =							"""\
GLUE, the General Language Understanding Evaluation benchmark
(https://gluebenchmark.com/) is a collection of resources for training,
evaluating, and analyzing natural language understanding systems.
"""

lowercase  :					int						    =							"""
Compute GLUE evaluation metric associated to each GLUE dataset.
Args:
    predictions: list of predictions to score.
        Each translation should be tokenized into a list of tokens.
    references: list of lists of references for each translation.
        Each reference should be tokenized into a list of tokens.
Returns: depending on the GLUE subset, one or several of:
    \"accuracy\": Accuracy
    \"f1\": F1 score
    \"pearson\": Pearson Correlation
    \"spearmanr\": Spearman Correlation
    \"matthews_correlation\": Matthew Correlation
Examples:

    >>> glue_metric = datasets.load_metric('glue', 'sst2')  # 'sst2' or any of [\"mnli\", \"mnli_mismatched\", \"mnli_matched\", \"qnli\", \"rte\", \"wnli\", \"hans\"]
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'accuracy': 1.0}

    >>> glue_metric = datasets.load_metric('glue', 'mrpc')  # 'mrpc' or 'qqp'
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'accuracy': 1.0, 'f1': 1.0}

    >>> glue_metric = datasets.load_metric('glue', 'stsb')
    >>> references = [0., 1., 2., 3., 4., 5.]
    >>> predictions = [0., 1., 2., 3., 4., 5.]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print({\"pearson\": round(results[\"pearson\"], 2), \"spearmanr\": round(results[\"spearmanr\"], 2)})
    {'pearson': 1.0, 'spearmanr': 1.0}

    >>> glue_metric = datasets.load_metric('glue', 'cola')
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'matthews_correlation': 1.0}
"""







def 						_snake_case(  SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				) ->    Optional[int]:
       return float((preds == labels).mean()				)







def 						_snake_case(  SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				) ->    Optional[Any]:
       lowercase							:     Any       	=      simple_accuracy(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)
       lowercase							:     Union[str, Any]       	=      float(fa_score(y_true=SCREAMING_SNAKE_CASE__					,	y_pred=SCREAMING_SNAKE_CASE__				)				)
       return {
           "accuracy": acc,
           "f1": fa,
       }







def 						_snake_case(  SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				) ->    List[Any]:
       lowercase							:     Union[str, Any]       	=      float(pearsonr(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)[0]				)
       lowercase							:     Dict       	=      float(spearmanr(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)[0]				)
       return {
           "pearson": pearson_corr,
           "spearmanr": spearman_corr,
       }




@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION					,	_KWARGS_DESCRIPTION     )
class __snake_case		(      datasets.Metric     ):
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           if self.config_name not in [
               "sst2",
               "mnli",
               "mnli_mismatched",
               "mnli_matched",
               "cola",
               "stsb",
               "mrpc",
               "qqp",
               "qnli",
               "rte",
               "wnli",
               "hans",
           ]:
                  raise KeyError(
                      """You should supply a configuration name selected in """
                      """[\"sst2\", \"mnli\", \"mnli_mismatched\", \"mnli_matched\", """
                      """\"cola\", \"stsb\", \"mrpc\", \"qqp\", \"qnli\", \"rte\", \"wnli\", \"hans\"]"""      )
           return datasets.MetricInfo(
               description=_DESCRIPTION		,citation=_CITATION		,inputs_description=_KWARGS_DESCRIPTION		,features=datasets.Features(
                   {
                       """predictions""": datasets.Value("""int64""" if self.config_name != """stsb""" else """float32"""      ),
                       """references""": datasets.Value("""int64""" if self.config_name != """stsb""" else """float32"""      ),
                   }      )		,codebase_urls=[]		,reference_urls=[]		,format="""numpy"""		,)
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case      ):



           '''simple docstring'''







           if self.config_name == "cola":
                  return {"matthews_correlation": matthews_corrcoef(snake_case		,snake_case      )}
           elif self.config_name == "stsb":
                  return pearson_and_spearman(snake_case		,snake_case      )
           elif self.config_name in ["mrpc", "qqp"]:
                  return acc_and_fa(snake_case		,snake_case      )
           elif self.config_name in ["sst2", "mnli", "mnli_mismatched", "mnli_matched", "qnli", "rte", "wnli", "hans"]:
                  return {"accuracy": simple_accuracy(snake_case		,snake_case      )}
           else:
                  raise KeyError(
                      """You should supply a configuration name selected in """
                      """[\"sst2\", \"mnli\", \"mnli_mismatched\", \"mnli_matched\", """
                      """\"cola\", \"stsb\", \"mrpc\", \"qqp\", \"qnli\", \"rte\", \"wnli\", \"hans\"]"""      )
