



"""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)

A_						:    List[str]				=      'bert-base-cased'
A_						:    Optional[int]				=      'fp16'
A_						:    str				=      'bf16'
A_						:    List[str]				=      [FPaa, BFaa]
@require_fsdp
@require_cuda
class     lowerCAmelCase__						(	_lowerCamelCase ):


	'''simple docstring'''







	def 			_lowerCAmelCase     (      self				:       Optional[int]    )						->	Optional[int]:





					"""simple docstring"""


					super().setUp()

					SCREAMING_SNAKE_CASE						:							List[Any]							=     dict(
					    ACCELERATE_USE_FSDP='true'     ,       MASTER_ADDR='localhost'     ,       MASTER_PORT='10999'     ,       RANK='0'     ,       LOCAL_RANK='0'     ,       WORLD_SIZE='1'     ,       )





	def 			_lowerCAmelCase     (      self				:       Tuple    )						->	List[Any]:





					"""simple docstring"""


					from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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





	def 			_lowerCAmelCase     (      self				:       List[Any]    )						->	Dict:





					"""simple docstring"""


					from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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





	def 			_lowerCAmelCase     (      self				:       Union[str, Any]    )						->	str:





					"""simple docstring"""


					from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

					for i, state_dict_type in enumerate(_SCREAMING_SNAKE_CASE    ):
									SCREAMING_SNAKE_CASE						:							Any							=     self.dist_env.copy()
									SCREAMING_SNAKE_CASE						:							Optional[Any]							=     state_dict_type
									with mockenv_context(**_SCREAMING_SNAKE_CASE    ):
													SCREAMING_SNAKE_CASE						:							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 			_lowerCAmelCase     (      self				:       Any    )						->	Dict:





					"""simple docstring"""


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

					SCREAMING_SNAKE_CASE						:							int							=     self.dist_env.copy()
					SCREAMING_SNAKE_CASE						:							Any							=     'TRANSFORMER_BASED_WRAP'
					SCREAMING_SNAKE_CASE						:							str							=     'T5Layer'
					with mockenv_context(**_SCREAMING_SNAKE_CASE    ):
									SCREAMING_SNAKE_CASE						:							List[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    )    )

					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     self.dist_env.copy()
					SCREAMING_SNAKE_CASE						:							Optional[int]							=     'SIZE_BASED_WRAP'
					SCREAMING_SNAKE_CASE						:							List[Any]							=     '0'
					with mockenv_context(**_SCREAMING_SNAKE_CASE    ):
									SCREAMING_SNAKE_CASE						:							List[Any]							=     FullyShardedDataParallelPlugin()
									fsdp_plugin.set_auto_wrap_policy(_SCREAMING_SNAKE_CASE    )
									self.assertIsNone(fsdp_plugin.auto_wrap_policy    )





	def 			_lowerCAmelCase     (      self				:       int    )						->	str:





					"""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:
									SCREAMING_SNAKE_CASE						:							int							=     self.dist_env.copy()
									SCREAMING_SNAKE_CASE						:							List[Any]							=     mp_dtype
									with mockenv_context(**_SCREAMING_SNAKE_CASE    ):
													SCREAMING_SNAKE_CASE						:							Any							=     Accelerator()
													if mp_dtype == "fp16":
																	SCREAMING_SNAKE_CASE						:							int							=     torch.floataa
													elif mp_dtype == "bf16":
																	SCREAMING_SNAKE_CASE						:							int							=     torch.bfloataa
													SCREAMING_SNAKE_CASE						:							Tuple							=     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 			_lowerCAmelCase     (      self				:       Tuple    )						->	str:





					"""simple docstring"""


					from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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




@require_fsdp
@require_multi_gpu
@slow
class     lowerCAmelCase__						(	_lowerCamelCase ):


	'''simple docstring'''







	def 			_lowerCAmelCase     (      self				:       List[str]    )						->	Dict:





					"""simple docstring"""


					super().setUp()
					SCREAMING_SNAKE_CASE						:							Dict							=     0.8_2
					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     [
					    'fsdp_shard_grad_op_transformer_based_wrap',
					    'fsdp_full_shard_transformer_based_wrap',
					]
					SCREAMING_SNAKE_CASE						:							Tuple							=     {
					    '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
					}
					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     160
					SCREAMING_SNAKE_CASE						:							Tuple							=     160

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





	def 			_lowerCAmelCase     (      self				:       Optional[int]    )						->	Dict:





					"""simple docstring"""


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





					"""simple docstring"""


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

													SCREAMING_SNAKE_CASE						:							Tuple							=     cmd_config[:-1]
													SCREAMING_SNAKE_CASE						:							str							=     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 			_lowerCAmelCase     (      self				:       str    )						->	str:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Tuple							=     os.path.join(self.test_scripts_folder     ,       'test_peak_memory_usage.py'    )
					SCREAMING_SNAKE_CASE						:							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():
									SCREAMING_SNAKE_CASE						:							Union[str, 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"""

def  __snake_case							(	__A      :						list							,		__A      :						int							,		__A      :						int = 0							,		__A      :						int = 0     )     ->       int:






				'''simple docstring'''
				SCREAMING_SNAKE_CASE						:							Any							=     right or len(__A     ) - 1
				if left > right:
								return -1
				elif list_data[left] == key:
								return left
				elif list_data[right] == key:
								return right
				else:
								return search(__A							,		__A							,		left + 1							,		right - 1     )


if __name__ == "__main__":
							import doctest

							doctest.testmod()





