






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)

_UpperCAmelCase :			Optional[Any]              =	"bert-base-cased"
_UpperCAmelCase :			Dict              =	"fp16"
_UpperCAmelCase :			int              =	"bf16"
_UpperCAmelCase :			Optional[Any]              =	[FPaa, BFaa]


@require_fsdp
@require_cuda
class     lowercase (		_SCREAMING_SNAKE_CASE       ):



		def __UpperCamelCase			(			self	)  -> str:
				"""simple docstring"""
				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	)  -> List[str]:
				"""simple docstring"""
				from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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



		def __UpperCamelCase			(			self	)  -> str:
				"""simple docstring"""
				from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

				for i, prefetch_policy in enumerate(A_	):
						UpperCamelCase          =					self.dist_env.copy()
						UpperCamelCase          =					prefetch_policy
						with mockenv_context(**A_	):
								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	)  -> Optional[Any]:
				"""simple docstring"""
				from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

				for i, state_dict_type in enumerate(A_	):
						UpperCamelCase          =					self.dist_env.copy()
						UpperCamelCase          =					state_dict_type
						with mockenv_context(**A_	):
								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	)  -> Any:
				"""simple docstring"""
				UpperCamelCase          =					AutoModel.from_pretrained(A_	)
				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(**A_	):
								UpperCamelCase          =					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	)

				UpperCamelCase          =					self.dist_env.copy()
				UpperCamelCase          =					'TRANSFORMER_BASED_WRAP'
				UpperCamelCase          =					'T5Layer'
				with mockenv_context(**A_	):
						UpperCamelCase          =					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	)	)

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



		def __UpperCamelCase			(			self	)  -> List[Any]:
				"""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:
						UpperCamelCase          =					self.dist_env.copy()
						UpperCamelCase          =					mp_dtype
						with mockenv_context(**A_	):
								UpperCamelCase          =					Accelerator()
								if mp_dtype == "fp16":
										UpperCamelCase          =					torch.floataa
								elif mp_dtype == "bf16":
										UpperCamelCase          =					torch.bfloataa
								UpperCamelCase          =					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 __UpperCamelCase			(			self	)  -> Union[str, Any]:
				"""simple docstring"""
				from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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






@require_fsdp
@require_multi_gpu
@slow
class     lowercase (		_SCREAMING_SNAKE_CASE       ):



		def __UpperCamelCase			(			self	)  -> Union[str, Any]:
				"""simple docstring"""
				super().setUp()
				UpperCamelCase          =					0.82
				UpperCamelCase          =					[
				    'fsdp_shard_grad_op_transformer_based_wrap',
				    'fsdp_full_shard_transformer_based_wrap',
				]
				UpperCamelCase          =					{
				    '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          =					160
				UpperCamelCase          =					160

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



		def __UpperCamelCase			(			self	)  -> Any:
				"""simple docstring"""
				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(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 __UpperCamelCase			(			self	)  -> List[str]:
				"""simple docstring"""
				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(A_	):
						UpperCamelCase          =					cmd.copy()
						cmd_config.append(F'''--fsdp_sharding_strategy={i+1}'''	)
						if strategy != "FULL_SHARD":
								continue
						UpperCamelCase          =					len(A_	)
						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(A_    ,				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(A_    ,				env=os.environ.copy()	)



		def __UpperCamelCase			(			self	)  -> Union[str, Any]:
				"""simple docstring"""
				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(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()	)


_UpperCAmelCase :			str              =	[
    999,
    800,
    799,
    600,
    599,
    500,
    400,
    399,
    377,
    355,
    333,
    311,
    288,
    266,
    244,
    222,
    200,
    199,
    177,
    155,
    133,
    111,
    88,
    66,
    44,
    22,
    0,
]

_UpperCAmelCase :			Any              =	[
    999,
    976,
    952,
    928,
    905,
    882,
    858,
    857,
    810,
    762,
    715,
    714,
    572,
    429,
    428,
    286,
    285,
    238,
    190,
    143,
    142,
    118,
    95,
    71,
    47,
    24,
    0,
]

_UpperCAmelCase :			List[Any]              =	[
    999,
    988,
    977,
    966,
    955,
    944,
    933,
    922,
    911,
    900,
    899,
    879,
    859,
    840,
    820,
    800,
    799,
    766,
    733,
    700,
    699,
    650,
    600,
    599,
    500,
    499,
    400,
    399,
    350,
    300,
    299,
    266,
    233,
    200,
    199,
    179,
    159,
    140,
    120,
    100,
    99,
    88,
    77,
    66,
    55,
    44,
    33,
    22,
    11,
    0,
]

_UpperCAmelCase :			Tuple              =	[
    999,
    995,
    992,
    989,
    985,
    981,
    978,
    975,
    971,
    967,
    964,
    961,
    957,
    956,
    951,
    947,
    942,
    937,
    933,
    928,
    923,
    919,
    914,
    913,
    908,
    903,
    897,
    892,
    887,
    881,
    876,
    871,
    870,
    864,
    858,
    852,
    846,
    840,
    834,
    828,
    827,
    820,
    813,
    806,
    799,
    792,
    785,
    784,
    777,
    770,
    763,
    756,
    749,
    742,
    741,
    733,
    724,
    716,
    707,
    699,
    698,
    688,
    677,
    666,
    656,
    655,
    645,
    634,
    623,
    613,
    612,
    598,
    584,
    570,
    569,
    555,
    541,
    527,
    526,
    505,
    484,
    483,
    462,
    440,
    439,
    396,
    395,
    352,
    351,
    308,
    307,
    264,
    263,
    220,
    219,
    176,
    132,
    88,
    44,
    0,
]

_UpperCAmelCase :			Union[str, Any]              =	[
    999,
    997,
    995,
    992,
    990,
    988,
    986,
    984,
    981,
    979,
    977,
    975,
    972,
    970,
    968,
    966,
    964,
    961,
    959,
    957,
    956,
    954,
    951,
    949,
    946,
    944,
    941,
    939,
    936,
    934,
    931,
    929,
    926,
    924,
    921,
    919,
    916,
    914,
    913,
    910,
    907,
    905,
    902,
    899,
    896,
    893,
    891,
    888,
    885,
    882,
    879,
    877,
    874,
    871,
    870,
    867,
    864,
    861,
    858,
    855,
    852,
    849,
    846,
    843,
    840,
    837,
    834,
    831,
    828,
    827,
    824,
    821,
    817,
    814,
    811,
    808,
    804,
    801,
    798,
    795,
    791,
    788,
    785,
    784,
    780,
    777,
    774,
    770,
    766,
    763,
    760,
    756,
    752,
    749,
    746,
    742,
    741,
    737,
    733,
    730,
    726,
    722,
    718,
    714,
    710,
    707,
    703,
    699,
    698,
    694,
    690,
    685,
    681,
    677,
    673,
    669,
    664,
    660,
    656,
    655,
    650,
    646,
    641,
    636,
    632,
    627,
    622,
    618,
    613,
    612,
    607,
    602,
    596,
    591,
    586,
    580,
    575,
    570,
    569,
    563,
    557,
    551,
    545,
    539,
    533,
    527,
    526,
    519,
    512,
    505,
    498,
    491,
    484,
    483,
    474,
    466,
    457,
    449,
    440,
    439,
    428,
    418,
    407,
    396,
    395,
    381,
    366,
    352,
    351,
    330,
    308,
    307,
    286,
    264,
    263,
    242,
    220,
    219,
    176,
    175,
    132,
    131,
    88,
    44,
    0,
]

_UpperCAmelCase :			List[str]              =	[
    999,
    991,
    982,
    974,
    966,
    958,
    950,
    941,
    933,
    925,
    916,
    908,
    900,
    899,
    874,
    850,
    825,
    800,
    799,
    700,
    600,
    500,
    400,
    300,
    200,
    100,
    0,
]

_UpperCAmelCase :			Tuple              =	[
    999,
    992,
    985,
    978,
    971,
    964,
    957,
    949,
    942,
    935,
    928,
    921,
    914,
    907,
    900,
    899,
    879,
    859,
    840,
    820,
    800,
    799,
    766,
    733,
    700,
    699,
    650,
    600,
    599,
    500,
    499,
    400,
    399,
    300,
    299,
    200,
    199,
    100,
    99,
    0,
]

_UpperCAmelCase :			Dict              =	[
    999,
    996,
    992,
    989,
    985,
    982,
    979,
    975,
    972,
    968,
    965,
    961,
    958,
    955,
    951,
    948,
    944,
    941,
    938,
    934,
    931,
    927,
    924,
    920,
    917,
    914,
    910,
    907,
    903,
    900,
    899,
    891,
    884,
    876,
    869,
    861,
    853,
    846,
    838,
    830,
    823,
    815,
    808,
    800,
    799,
    788,
    777,
    766,
    755,
    744,
    733,
    722,
    711,
    700,
    699,
    688,
    677,
    666,
    655,
    644,
    633,
    622,
    611,
    600,
    599,
    585,
    571,
    557,
    542,
    528,
    514,
    500,
    499,
    485,
    471,
    457,
    442,
    428,
    414,
    400,
    399,
    379,
    359,
    340,
    320,
    300,
    299,
    279,
    259,
    240,
    220,
    200,
    199,
    166,
    133,
    100,
    99,
    66,
    33,
    0,
]
