

"""simple docstring"""

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)

_UpperCamelCase		     =    """bert-base-cased"""
_UpperCamelCase		     =    """fp16"""
_UpperCamelCase		     =    """bf16"""
_UpperCamelCase		     =    [FPaa, BFaa]






@require_fsdp
@require_cuda
class 			__a					(						__magic_name__       ):





       """simple docstring"""


       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        super().setUp()

        lowerCAmelCase__   :				List[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   ):
         lowerCAmelCase__   :				Dict 			=				self.dist_env.copy()
         lowerCAmelCase__   :				Any 			=				F"""{i + 1}"""
         lowerCAmelCase__   :				Tuple 			=				strategy
         with mockenv_context(**snake_case   ):
          lowerCAmelCase__   :				List[str] 			=				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   ):
         lowerCAmelCase__   :				int 			=				self.dist_env.copy()
         lowerCAmelCase__   :				Dict 			=				prefetch_policy
         with mockenv_context(**snake_case   ):
          lowerCAmelCase__   :				int 			=				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   ):
         lowerCAmelCase__   :				List[Any] 			=				self.dist_env.copy()
         lowerCAmelCase__   :				Union[str, Any] 			=				state_dict_type
         with mockenv_context(**snake_case   ):
          lowerCAmelCase__   :				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"""

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

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

        lowerCAmelCase__   :				Dict 			=				self.dist_env.copy()
        lowerCAmelCase__   :				str 			=				"SIZE_BASED_WRAP"
        lowerCAmelCase__   :				Tuple 			=				"0"
        with mockenv_context(**snake_case   ):
         lowerCAmelCase__   :				List[str] 			=				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:
         lowerCAmelCase__   :				Tuple 			=				self.dist_env.copy()
         lowerCAmelCase__   :				List[str] 			=				mp_dtype
         with mockenv_context(**snake_case   ):
          lowerCAmelCase__   :				List[str] 			=				Accelerator()
          if mp_dtype == "fp16":
           lowerCAmelCase__   :				Tuple 			=				torch.floataa
          elif mp_dtype == "bf16":
           lowerCAmelCase__   :				List[Any] 			=				torch.bfloataa
          lowerCAmelCase__   :				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]:
         lowerCAmelCase__   :				Any 			=				self.dist_env.copy()
         lowerCAmelCase__   :				Union[str, Any] 			=				str(snake_case   ).lower()
         with mockenv_context(**snake_case   ):
          lowerCAmelCase__   :				str 			=				FullyShardedDataParallelPlugin()
          self.assertEqual(fsdp_plugin.cpu_offload			,					CPUOffload(offload_params=snake_case   )   )






@require_fsdp
@require_multi_gpu
@slow
class 			__a					(						__magic_name__       ):





       """simple docstring"""


       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        super().setUp()
        lowerCAmelCase__   :				List[str] 			=				0.82
        lowerCAmelCase__   :				Any 			=				[
            "fsdp_shard_grad_op_transformer_based_wrap",
            "fsdp_full_shard_transformer_based_wrap",
        ]
        lowerCAmelCase__   :				Any 			=				{
            "multi_gpu_fp16": 3_200,
            "fsdp_shard_grad_op_transformer_based_wrap_fp16": 2_000,
            "fsdp_full_shard_transformer_based_wrap_fp16": 1_900,
            # 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
        }
        lowerCAmelCase__   :				Union[str, Any] 			=				160
        lowerCAmelCase__   :				Optional[Any] 			=				160

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






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[Any] 			=				os.path.join(self.test_scripts_folder			,					"test_performance.py"   )
        lowerCAmelCase__   :				Union[str, Any] 			=				["accelerate", "launch", "--num_processes=2", "--num_machines=1", "--machine_rank=0", "--use_fsdp"]
        for config in self.performance_configs:
         lowerCAmelCase__   :				Tuple 			=				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"""

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

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






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[Any] 			=				os.path.join(self.test_scripts_folder			,					"test_peak_memory_usage.py"   )
        lowerCAmelCase__   :				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():
         lowerCAmelCase__   :				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(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()   )









"""simple docstring"""

def       SCREAMING_SNAKE_CASE							( lowercase__					,					lowercase__   )			->						list:
 lowerCAmelCase__   :				List[str] 			=				len(lowercase__   )
 lowerCAmelCase__   :				Dict 			=				[]
 for i in range(len(lowercase__   ) - pat_len + 1   ):
  lowerCAmelCase__   :				Union[str, Any] 			=				True
  for j in range(lowercase__   ):
   if s[i + j] != pattern[j]:
    lowerCAmelCase__   :				int 			=				False
    break
  if match_found:
   position.append(lowercase__   )
 return position


if __name__ == "__main__":
    assert naive_pattern_search("""ABCDEFG""", """DE""") == [3]
    print(naive_pattern_search("""ABAAABCDBBABCDDEBCABC""", """ABC"""))






