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

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






@require_fsdp
@require_cuda
class 						A   (   UpperCAmelCase   ):





							def       snake_case__						(      self				:							List[Any]	)							->	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       snake_case__						(      self				:							str	)							->	int:
								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       snake_case__						(      self				:							List[Any]	)							->	Optional[int]:
								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       snake_case__						(      self				:							Optional[Any]	)							->	List[Any]:
								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       snake_case__						(      self				:							int	)							->	Optional[Any]:
								__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       snake_case__						(      self				:							List[str]	)							->	Dict:
								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       snake_case__						(      self				:							List[str]	)							->	List[Any]:
								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 						A   (   UpperCAmelCase   ):





							def       snake_case__						(      self				:							Dict	)							->	Optional[Any]:
								super().setUp()
								__UpperCAmelCase							=			0.8_2
								__UpperCAmelCase							=			[
								    '''fsdp_shard_grad_op_transformer_based_wrap''',
								    '''fsdp_full_shard_transformer_based_wrap''',
								]
								__UpperCAmelCase							=			{
								    '''multi_gpu_fp16''': 3_2_0_0,
								    '''fsdp_shard_grad_op_transformer_based_wrap_fp16''': 2_0_0_0,
								    '''fsdp_full_shard_transformer_based_wrap_fp16''': 1_9_0_0,
								    # 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_6_0
								__UpperCAmelCase							=			1_6_0

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





							def       snake_case__						(      self				:							Tuple	)							->	List[str]:
								__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       snake_case__						(      self				:							str	)							->	int:
								__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       snake_case__						(      self				:							int	)							->	List[Any]:
								__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()	)

'''simple docstring'''


import unittest

import torch
from torch import nn

from accelerate.test_utils import require_cuda
from accelerate.utils.memory import find_executable_batch_size, release_memory





def       lowerCAmelCase					(		):


	"""simple docstring"""
	raise RuntimeError('''CUDA out of memory.'''       )






class 						A   (   nn.Module   ):





							def __init__(      self				:							Optional[Any]	)							->	int:
								super().__init__()
								__UpperCAmelCase							=			nn.Linear(3     ,		4	)
								__UpperCAmelCase							=			nn.BatchNormad(4	)
								__UpperCAmelCase							=			nn.Linear(4     ,		5	)






							def       snake_case__						(      self				:							List[str]     ,		__a				:							Optional[int]	)							->	Optional[int]:
								return self.lineara(self.batchnorm(self.lineara(__a	)	)	)






class 						A   (   unittest.TestCase   ):





							def       snake_case__						(      self				:							Optional[int]	)							->	Any:
								__UpperCAmelCase							=			[]

								@find_executable_batch_size(starting_batch_size=1_2_8	)
								def mock_training_loop_function(__a				:							Union[str, Any]	):
									nonlocal batch_sizes
									batch_sizes.append(__a	)
									if batch_size != 8:
										raise_fake_out_of_memory()

								mock_training_loop_function()
								self.assertListEqual(__a     ,		[1_2_8, 6_4, 3_2, 1_6, 8]	)





							def       snake_case__						(      self				:							str	)							->	int:
								__UpperCAmelCase							=			[]

								@find_executable_batch_size(starting_batch_size=1_2_8	)
								def mock_training_loop_function(__a				:							str     ,		__a				:							Optional[int]	):
									nonlocal batch_sizes
									batch_sizes.append(__a	)
									if batch_size != 8:
										raise_fake_out_of_memory()
									return batch_size, arga

								__UpperCAmelCase				,			__UpperCAmelCase							=			mock_training_loop_function('''hello'''	)
								self.assertListEqual(__a     ,		[1_2_8, 6_4, 3_2, 1_6, 8]	)
								self.assertListEqual([bs, arga]     ,		[8, '''hello''']	)





							def       snake_case__						(      self				:							Any	)							->	int:
								@find_executable_batch_size(starting_batch_size=0	)
								def mock_training_loop_function(__a				:							Optional[int]	):
									pass

								with self.assertRaises(__a	) as cm:
									mock_training_loop_function()
									self.assertIn('''No executable batch size found, reached zero.'''     ,		cm.exception.args[0]	)





							def       snake_case__						(      self				:							Any	)							->	List[Any]:
								@find_executable_batch_size(starting_batch_size=1_6	)
								def mock_training_loop_function(__a				:							Dict	):
									if batch_size > 0:
										raise_fake_out_of_memory()
									pass

								with self.assertRaises(__a	) as cm:
									mock_training_loop_function()
									self.assertIn('''No executable batch size found, reached zero.'''     ,		cm.exception.args[0]	)





							def       snake_case__						(      self				:							List[Any]	)							->	List[str]:
								@find_executable_batch_size(starting_batch_size=1_2_8	)
								def mock_training_loop_function(__a				:							str     ,		__a				:							Union[str, Any]     ,		__a				:							int	):
									if batch_size != 8:
										raise raise_fake_out_of_memory()

								with self.assertRaises(__a	) as cm:
									mock_training_loop_function(1_2_8     ,		'''hello'''     ,		'''world'''	)
									self.assertIn('''Batch size was passed into `f`'''     ,		cm.exception.args[0]	)
									self.assertIn('''`f(arg1=\'hello\', arg2=\'world\')'''     ,		cm.exception.args[0]	)





							def       snake_case__						(      self				:							Tuple	)							->	Optional[Any]:
								@find_executable_batch_size(starting_batch_size=1_6	)
								def mock_training_loop_function(__a				:							Tuple	):
									raise ValueError('''Oops, we had an error!'''	)

								with self.assertRaises(__a	) as cm:
									mock_training_loop_function()
									self.assertIn('''Oops, we had an error!'''     ,		cm.exception.args[0]	)






							@require_cuda
							def       snake_case__						(      self				:							Any	)							->	List[Any]:
								__UpperCAmelCase							=			torch.cuda.memory_allocated()
								__UpperCAmelCase							=			ModelForTest()
								model.cuda()
								self.assertGreater(torch.cuda.memory_allocated()     ,		__a	)
								__UpperCAmelCase							=			release_memory(__a	)
								self.assertEqual(torch.cuda.memory_allocated()     ,		__a	)
