
'''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(42)

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




@require_fsdp
@require_cuda
class 	A_			(    _A							):






     '''simple docstring'''



     def        UpperCAmelCase_	(		self			:					Optional[int]	)						->     Tuple:
      super().setUp()

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


     def        UpperCAmelCase_	(		self			:					str	)						->     Tuple:
      from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

      for i, strategy in enumerate(__SCREAMING_SNAKE_CASE	):
       UpperCAmelCase	:     Tuple						 =			self.dist_env.copy()
       UpperCAmelCase	:     List[str]						 =			f"""{i + 1}"""
       UpperCAmelCase	:     List[str]						 =			strategy
       with mockenv_context(**__SCREAMING_SNAKE_CASE	):
        UpperCAmelCase	:     str						 =			FullyShardedDataParallelPlugin()
        self.assertEqual(fsdp_plugin.sharding_strategy					,      ShardingStrategy(i + 1	)	)


     def        UpperCAmelCase_	(		self			:					Any	)						->     Union[str, Any]:
      from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

      for i, prefetch_policy in enumerate(__SCREAMING_SNAKE_CASE	):
       UpperCAmelCase	:     List[str]						 =			self.dist_env.copy()
       UpperCAmelCase	:     Dict						 =			prefetch_policy
       with mockenv_context(**__SCREAMING_SNAKE_CASE	):
        UpperCAmelCase	:     str						 =			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			:					Union[str, Any]	)						->     Optional[int]:
      from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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


     def        UpperCAmelCase_	(		self			:					List[str]	)						->     Optional[Any]:
      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	:     List[str]						 =			self.dist_env.copy()
       UpperCAmelCase	:     Union[str, Any]						 =			mp_dtype
       with mockenv_context(**__SCREAMING_SNAKE_CASE	):
        UpperCAmelCase	:     Tuple						 =			Accelerator()
        if mp_dtype == "fp16":
         UpperCAmelCase	:     Any						 =			torch.floataa
        elif mp_dtype == "bf16":
         UpperCAmelCase	:     str						 =			torch.bfloataa
        UpperCAmelCase	:     Optional[int]						 =			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        UpperCAmelCase_	(		self			:					Tuple	)						->     Optional[int]:
      from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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




@require_fsdp
@require_multi_gpu
@slow
class 	A_			(    _A							):






     '''simple docstring'''



     def        UpperCAmelCase_	(		self			:					Tuple	)						->     int:
      super().setUp()
      UpperCAmelCase	:     Tuple						 =			0.82
      UpperCAmelCase	:     List[Any]						 =			[
          'fsdp_shard_grad_op_transformer_based_wrap',
          'fsdp_full_shard_transformer_based_wrap',
      ]
      UpperCAmelCase	:     List[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
      }
      UpperCAmelCase	:     Union[str, Any]						 =			160
      UpperCAmelCase	:     Tuple						 =			160

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


     def        UpperCAmelCase_	(		self			:					str	)						->     Tuple:
      UpperCAmelCase	:     Optional[Any]						 =			os.path.join(self.test_scripts_folder					,      'test_performance.py'	)
      UpperCAmelCase	:     Optional[Any]						 =			['accelerate', 'launch', '--num_processes=2', '--num_machines=1', '--machine_rank=0', '--use_fsdp']
      for config in self.performance_configs:
       UpperCAmelCase	:     Dict						 =			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        UpperCAmelCase_	(		self			:					Tuple	)						->     Optional[int]:
      UpperCAmelCase	:     Tuple						 =			os.path.join(self.test_scripts_folder					,      'test_checkpointing.py'	)
      UpperCAmelCase	:     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(__SCREAMING_SNAKE_CASE	):
       UpperCAmelCase	:     str						 =			cmd.copy()
       cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""	)
       if strategy != "FULL_SHARD":
        continue
       UpperCAmelCase	:     Union[str, Any]						 =			len(__SCREAMING_SNAKE_CASE	)
       for state_dict_type in FSDP_STATE_DICT_TYPE:
        UpperCAmelCase	:     Dict						 =			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[int]						 =			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        UpperCAmelCase_	(		self			:					Union[str, Any]	)						->     Optional[Any]:
      UpperCAmelCase	:     int						 =			os.path.join(self.test_scripts_folder					,      'test_peak_memory_usage.py'	)
      UpperCAmelCase	:     str						 =			[
          '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	:     Optional[Any]						 =			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()	)







'''simple docstring'''

import argparse
import json

from tqdm import tqdm




def 				UpperCamelCase(					):
 UpperCAmelCase	:     List[Any]						 =			argparse.ArgumentParser()

 # Required parameters
 parser.add_argument(
     '--src_path'  ,	type=UpperCAmelCase_  ,	default='biencoder-nq-dev.json'  ,	help='Path to raw DPR training data'  ,	)
 parser.add_argument(
     '--evaluation_set'  ,	type=UpperCAmelCase_  ,	help='where to store parsed evaluation_set file'  ,	)
 parser.add_argument(
     '--gold_data_path'  ,	type=UpperCAmelCase_  ,	help='where to store parsed gold_data_path file'  ,	)
 UpperCAmelCase	:     int						 =			parser.parse_args()

 with open(args.src_path  ,	'r'     ) as src_file, open(args.evaluation_set  ,	'w'     ) as eval_file, open(
     args.gold_data_path  ,	'w'     ) as gold_file:
  UpperCAmelCase	:     int						 =			json.load(UpperCAmelCase_     )
  for dpr_record in tqdm(UpperCAmelCase_     ):
   UpperCAmelCase	:     Any						 =			dpr_record['question']
   UpperCAmelCase	:     List[str]						 =			[context['title'] for context in dpr_record['positive_ctxs']]
   eval_file.write(question + '\n'     )
   gold_file.write('\t'.join(UpperCAmelCase_     ) + '\n'     )


if __name__ == "__main__":
 main()




