



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

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







@require_fsdp
@require_cuda
class 				A_      (					SCREAMING_SNAKE_CASE_    ):



						"""simple docstring"""






						def      UpperCAmelCase__					( self			:Tuple )    ->      List[Any]:
												super().setUp()

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






						def      UpperCAmelCase__					( self			:Optional[int] )    ->      List[Any]:
												from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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






						def      UpperCAmelCase__					( self			:Optional[Any] )    ->      str:
												from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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

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

												UpperCAmelCase					=			self.dist_env.copy()
												UpperCAmelCase					=			'TRANSFORMER_BASED_WRAP'
												UpperCAmelCase					=			'T5Layer'
												with mockenv_context(**lowercase_ ):
																		UpperCAmelCase					=			FullyShardedDataParallelPlugin()
																		with self.assertRaises(lowercase_ ) as cm:
																								fsdp_plugin.set_auto_wrap_policy(lowercase_ )
																		self.assertTrue('Could not find the transformer layer class to wrap in the model.' in str(cm.exception ) )

												UpperCAmelCase					=			self.dist_env.copy()
												UpperCAmelCase					=			'SIZE_BASED_WRAP'
												UpperCAmelCase					=			'0'
												with mockenv_context(**lowercase_ ):
																		UpperCAmelCase					=			FullyShardedDataParallelPlugin()
																		fsdp_plugin.set_auto_wrap_policy(lowercase_ )
																		self.assertIsNone(fsdp_plugin.auto_wrap_policy )






						def      UpperCAmelCase__					( self			:Dict )    ->      List[str]:
												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					=			self.dist_env.copy()
																		UpperCAmelCase					=			mp_dtype
																		with mockenv_context(**lowercase_ ):
																								UpperCAmelCase					=			Accelerator()
																								if mp_dtype == "fp16":
																														UpperCAmelCase					=			torch.floataa
																								elif mp_dtype == "bf16":
																														UpperCAmelCase					=			torch.bfloataa
																								UpperCAmelCase					=			MixedPrecision(param_dtype=lowercase_       ,  reduce_dtype=lowercase_       ,  buffer_dtype=lowercase_ )
																								self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy       ,  lowercase_ )
																								if mp_dtype == FPaa:
																														self.assertTrue(isinstance(accelerator.scaler       ,  lowercase_ ) )
																								elif mp_dtype == BFaa:
																														self.assertIsNone(accelerator.scaler )
																								AcceleratorState._reset_state(lowercase_ )






						def      UpperCAmelCase__					( self			:Any )    ->      Tuple:
												from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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







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



						"""simple docstring"""






						def      UpperCAmelCase__					( self			:Dict )    ->      Union[str, Any]:
												super().setUp()
												UpperCAmelCase					=			0.82
												UpperCAmelCase					=			[
												    'fsdp_shard_grad_op_transformer_based_wrap',
												    'fsdp_full_shard_transformer_based_wrap',
												]
												UpperCAmelCase					=			{
												    'multi_gpu_fp16': 32_00,
												    'fsdp_shard_grad_op_transformer_based_wrap_fp16': 20_00,
												    'fsdp_full_shard_transformer_based_wrap_fp16': 19_00,
												    # 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					=			1_60
												UpperCAmelCase					=			1_60

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






						def      UpperCAmelCase__					( self			:Tuple )    ->      Optional[int]:
												UpperCAmelCase					=			os.path.join(self.test_scripts_folder       ,  'test_performance.py' )
												UpperCAmelCase					=			['accelerate', 'launch', '--num_processes=2', '--num_machines=1', '--machine_rank=0', '--use_fsdp']
												for config in self.performance_configs:
																		UpperCAmelCase					=			cmd.copy()
																		for i, strategy in enumerate(lowercase_ ):
																								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(lowercase_       ,  env=os.environ.copy() )






						def      UpperCAmelCase__					( self			:str )    ->      Any:
												UpperCAmelCase					=			os.path.join(self.test_scripts_folder       ,  'test_checkpointing.py' )
												UpperCAmelCase					=			[
												    '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(lowercase_ ):
																		UpperCAmelCase					=			cmd.copy()
																		cmd_config.append(f"""--fsdp_sharding_strategy={i+1}""" )
																		if strategy != "FULL_SHARD":
																								continue
																		UpperCAmelCase					=			len(lowercase_ )
																		for state_dict_type in FSDP_STATE_DICT_TYPE:
																								UpperCAmelCase					=			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(lowercase_       ,  env=os.environ.copy() )

																								UpperCAmelCase					=			cmd_config[:-1]
																								UpperCAmelCase					=			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(lowercase_       ,  env=os.environ.copy() )






						def      UpperCAmelCase__					( self			:List[str] )    ->      List[str]:
												UpperCAmelCase					=			os.path.join(self.test_scripts_folder       ,  'test_peak_memory_usage.py' )
												UpperCAmelCase					=			[
												    '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					=			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(lowercase_ ):
																														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(lowercase_       ,  env=os.environ.copy() )










"""simple docstring"""


def 			_lowerCAmelCase				(       lowercase_						):

						if not isinstance(lowercase_		,     lowercase_						):
												UpperCAmelCase					=			F"""Input value of [number={number}] must be an integer"""
												raise TypeError(lowercase_						)

						if number < 1:
												UpperCAmelCase					=			F"""Input value of [number={number}] must be > 0"""
												raise ValueError(lowercase_						)

						UpperCAmelCase					=			1

						for i in range(1		,     lowercase_						):
												current_number *= 4 * i - 2
												current_number //= i + 1

						return current_number


if __name__ == "__main__":
				import doctest

				doctest.testmod()





