


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(4_2)

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


@require_fsdp
@require_cuda
class     __UpperCAmelCase							(       __A							):


     """simple docstring"""


     def 	snake_case_  (	self    ):
        super().setUp()

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


     def 	snake_case_  (	self    ):
        from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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


     def 	snake_case_  (	self    ):
        from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

        for i, prefetch_policy in enumerate(__A    ):
           __a					=  self.dist_env.copy()
           __a					=  prefetch_policy
           with mockenv_context(**__A    ):
              __a					=  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    ):
        from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

        for i, state_dict_type in enumerate(__A    ):
           __a					=  self.dist_env.copy()
           __a					=  state_dict_type
           with mockenv_context(**__A    ):
              __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 	snake_case_  (	self    ):
        __a					=  AutoModel.from_pretrained(__A    )
        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(**__A    ):
              __a					=  FullyShardedDataParallelPlugin()
              fsdp_plugin.set_auto_wrap_policy(__A    )
              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(**__A    ):
           __a					=  FullyShardedDataParallelPlugin()
           with self.assertRaises(__A    ) as cm:
              fsdp_plugin.set_auto_wrap_policy(__A    )
           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(**__A    ):
           __a					=  FullyShardedDataParallelPlugin()
           fsdp_plugin.set_auto_wrap_policy(__A    )
           self.assertIsNone(fsdp_plugin.auto_wrap_policy    )


     def 	snake_case_  (	self    ):
        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(**__A    ):
              __a					=  Accelerator()
              if mp_dtype == "fp16":
                 __a					=  torch.floataa
              elif mp_dtype == "bf16":
                 __a					=  torch.bfloataa
              __a					=  MixedPrecision(param_dtype=__A  ,       reduce_dtype=__A  ,       buffer_dtype=__A    )
              self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy  ,       __A    )
              if mp_dtype == FPaa:
                 self.assertTrue(isinstance(accelerator.scaler  ,       __A    )    )
              elif mp_dtype == BFaa:
                 self.assertIsNone(accelerator.scaler    )
              AcceleratorState._reset_state(__A    )


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

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







@require_fsdp
@require_multi_gpu
@slow
class     __UpperCAmelCase							(       __A							):


     """simple docstring"""


     def 	snake_case_  (	self    ):
        super().setUp()
        __a					=  0.82
        __a					=  [
            """fsdp_shard_grad_op_transformer_based_wrap""",
            """fsdp_full_shard_transformer_based_wrap""",
        ]
        __a					=  {
            """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
        }
        __a					=  160
        __a					=  160

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


     def 	snake_case_  (	self    ):
        __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(__A    ):
              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(__A  ,       env=os.environ.copy()    )


     def 	snake_case_  (	self    ):
        __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(__A    ):
           __a					=  cmd.copy()
           cmd_config.append(f'''--fsdp_sharding_strategy={i+1}'''    )
           if strategy != "FULL_SHARD":
              continue
           __a					=  len(__A    )
           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(__A  ,       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(__A  ,       env=os.environ.copy()    )


     def 	snake_case_  (	self    ):
        __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(__A    ):
                 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(__A  ,       env=os.environ.copy()    )







class     __UpperCAmelCase							:


     """simple docstring"""


     def __init__(	self  ,       __A    ):
        __a					=  set_counts
        __a					=  max(__A    )
        __a					=  len(__A    )
        __a					=  [1] * num_sets
        __a					=  list(range(__A    )    )


     def 	snake_case_  (	self  ,       __A  ,       __A    ):
        __a					=  self.get_parent(__A    )
        __a					=  self.get_parent(__A    )

        if src_parent == dst_parent:
           return False

        if self.ranks[dst_parent] >= self.ranks[src_parent]:
           self.set_counts[dst_parent] += self.set_counts[src_parent]
           __a					=  0
           __a					=  dst_parent
           if self.ranks[dst_parent] == self.ranks[src_parent]:
              self.ranks[dst_parent] += 1
           __a					=  self.set_counts[dst_parent]
        else:
           self.set_counts[src_parent] += self.set_counts[dst_parent]
           __a					=  0
           __a					=  src_parent
           __a					=  self.set_counts[src_parent]

        __a					=  max(self.max_set  ,       __A    )
        return True


     def 	snake_case_  (	self  ,       __A    ):
        if self.parents[disj_set] == disj_set:
           return disj_set
        __a					=  self.get_parent(self.parents[disj_set]    )
        return self.parents[disj_set]



