





"""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       	=     'bert-base-cased'
A       	=     'fp16'
A       	=     'bf16'
A       	=     [FPaa, BFaa]




@require_fsdp
@require_cuda
class  __lowercase	(	a__							):







					'''simple docstring'''




					def        _lowerCamelCase						(  self   ):
									super().setUp()

									__a       :  Optional[int]					=     dict(
									    ACCELERATE_USE_FSDP='''true'''		,	MASTER_ADDR='''localhost'''		,	MASTER_PORT='''10999'''		,	RANK='''0'''		,	LOCAL_RANK='''0'''		,	WORLD_SIZE='''1'''		,	)







					def        _lowerCamelCase						(  self   ):
									from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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







					def        _lowerCamelCase						(  self   ):
									from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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

									for i, state_dict_type in enumerate(SCREAMING_SNAKE_CASE__   ):
													__a       :  List[str]					=     self.dist_env.copy()
													__a       :  List[str]					=     state_dict_type
													with mockenv_context(**SCREAMING_SNAKE_CASE__   ):
																	__a       :  Union[str, 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   ):
									__a       :  Any					=     AutoModel.from_pretrained(SCREAMING_SNAKE_CASE__   )
									for policy in FSDP_AUTO_WRAP_POLICY:
													__a       :  int					=     self.dist_env.copy()
													__a       :  Optional[int]					=     policy
													if policy == "TRANSFORMER_BASED_WRAP":
																	__a       :  Union[str, Any]					=     """BertLayer"""
													elif policy == "SIZE_BASED_WRAP":
																	__a       :  Union[str, Any]					=     """2000"""
													with mockenv_context(**SCREAMING_SNAKE_CASE__   ):
																	__a       :  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   )

									__a       :  List[str]					=     self.dist_env.copy()
									__a       :  str					=     """TRANSFORMER_BASED_WRAP"""
									__a       :  List[Any]					=     """T5Layer"""
									with mockenv_context(**SCREAMING_SNAKE_CASE__   ):
													__a       :  Optional[int]					=     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   )   )

									__a       :  Dict					=     self.dist_env.copy()
									__a       :  List[Any]					=     """SIZE_BASED_WRAP"""
									__a       :  int					=     """0"""
									with mockenv_context(**SCREAMING_SNAKE_CASE__   ):
													__a       :  str					=     FullyShardedDataParallelPlugin()
													fsdp_plugin.set_auto_wrap_policy(SCREAMING_SNAKE_CASE__   )
													self.assertIsNone(fsdp_plugin.auto_wrap_policy   )







					def        _lowerCamelCase						(  self   ):
									from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
									from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

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

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






@require_fsdp
@require_multi_gpu
@slow
class  __lowercase	(	a__							):







					'''simple docstring'''




					def        _lowerCamelCase						(  self   ):
									super().setUp()
									__a       :  Optional[int]					=     0.8_2
									__a       :  Optional[int]					=     [
									    """fsdp_shard_grad_op_transformer_based_wrap""",
									    """fsdp_full_shard_transformer_based_wrap""",
									]
									__a       :  Dict					=     {
									    """multi_gpu_fp16""": 3200,
									    """fsdp_shard_grad_op_transformer_based_wrap_fp16""": 2000,
									    """fsdp_full_shard_transformer_based_wrap_fp16""": 1900,
									    # 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
									}
									__a       :  List[str]					=     160
									__a       :  Any					=     160

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







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

																	__a       :  Dict					=     cmd_config[:-1]
																	__a       :  Any					=     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   ):
									__a       :  str					=     os.path.join(self.test_scripts_folder		,	'''test_peak_memory_usage.py'''   )
									__a       :  int					=     [
									    """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():
													__a       :  List[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"""




from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase__      :    Optional[Any]    							= logging.get_logger(__name__)

UpperCAmelCase__      :    Dict    							= {
    'MIT/ast-finetuned-audioset-10-10-0.4593': (
        'https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593/resolve/main/config.json'
    ),
}



class 	lowerCAmelCase_						(a__  ):


				"""simple docstring"""



				__UpperCamelCase      :						Optional[int]					       =   '''audio-spectrogram-transformer'''



				def __init__(self   ,       SCREAMING_SNAKE_CASE__=7_68   ,       SCREAMING_SNAKE_CASE__=12   ,       SCREAMING_SNAKE_CASE__=12   ,       SCREAMING_SNAKE_CASE__=30_72   ,       SCREAMING_SNAKE_CASE__="gelu"   ,       SCREAMING_SNAKE_CASE__=0.0   ,       SCREAMING_SNAKE_CASE__=0.0   ,       SCREAMING_SNAKE_CASE__=0.02   ,       SCREAMING_SNAKE_CASE__=1E-12   ,       SCREAMING_SNAKE_CASE__=16   ,       SCREAMING_SNAKE_CASE__=True   ,       SCREAMING_SNAKE_CASE__=10   ,       SCREAMING_SNAKE_CASE__=10   ,       SCREAMING_SNAKE_CASE__=10_24   ,       SCREAMING_SNAKE_CASE__=1_28   ,       **SCREAMING_SNAKE_CASE__   ,       )					->		Tuple:



											"""simple docstring"""




											super().__init__(**SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   hidden_size
											SCREAMING_SNAKE_CASE__						:       str           =   num_hidden_layers
											SCREAMING_SNAKE_CASE__						:       int           =   num_attention_heads
											SCREAMING_SNAKE_CASE__						:       Tuple           =   intermediate_size
											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   hidden_act
											SCREAMING_SNAKE_CASE__						:       Any           =   hidden_dropout_prob
											SCREAMING_SNAKE_CASE__						:       List[Any]           =   attention_probs_dropout_prob
											SCREAMING_SNAKE_CASE__						:       int           =   initializer_range
											SCREAMING_SNAKE_CASE__						:       int           =   layer_norm_eps
											SCREAMING_SNAKE_CASE__						:       Dict           =   patch_size
											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   qkv_bias
											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   frequency_stride
											SCREAMING_SNAKE_CASE__						:       Any           =   time_stride
											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   max_length
											SCREAMING_SNAKE_CASE__						:       Any           =   num_mel_bins




