






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    :    Dict         =     """bert-base-cased"""
__snake_case    :    List[str]         =     """fp16"""
__snake_case    :    Dict         =     """bf16"""
__snake_case    :    List[Any]         =     [FPaa, BFaa]


@require_fsdp
@require_cuda
class     __SCREAMING_SNAKE_CASE					(   __lowercase):






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												super().setUp()

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






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

												for i, prefetch_policy in enumerate(_UpperCamelCase   ):
																		lowerCAmelCase__    						=						self.dist_env.copy()
																		lowerCAmelCase__    						=						prefetch_policy
																		with mockenv_context(**_UpperCamelCase   ):
																								lowerCAmelCase__    						=						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   ):






												"""simple docstring"""
												from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

												for i, state_dict_type in enumerate(_UpperCamelCase   ):
																		lowerCAmelCase__    						=						self.dist_env.copy()
																		lowerCAmelCase__    						=						state_dict_type
																		with mockenv_context(**_UpperCamelCase   ):
																								lowerCAmelCase__    						=						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   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						AutoModel.from_pretrained(_UpperCamelCase   )
												for policy in FSDP_AUTO_WRAP_POLICY:
																		lowerCAmelCase__    						=						self.dist_env.copy()
																		lowerCAmelCase__    						=						policy
																		if policy == "TRANSFORMER_BASED_WRAP":
																								lowerCAmelCase__    						=						'BertLayer'
																		elif policy == "SIZE_BASED_WRAP":
																								lowerCAmelCase__    						=						'2000'
																		with mockenv_context(**_UpperCamelCase   ):
																								lowerCAmelCase__    						=						FullyShardedDataParallelPlugin()
																								fsdp_plugin.set_auto_wrap_policy(_UpperCamelCase   )
																								if policy == "NO_WRAP":
																														self.assertIsNone(fsdp_plugin.auto_wrap_policy   )
																								else:
																														self.assertIsNotNone(fsdp_plugin.auto_wrap_policy   )

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

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






						def       UpperCamelCase__      (	self   ):






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







						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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




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






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												super().setUp()
												lowerCAmelCase__    						=						0.82
												lowerCAmelCase__    						=						[
												    'fsdp_shard_grad_op_transformer_based_wrap',
												    'fsdp_full_shard_transformer_based_wrap',
												]
												lowerCAmelCase__    						=						{
												    '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
												}
												lowerCAmelCase__    						=						1_60
												lowerCAmelCase__    						=						1_60

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






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						os.path.join(self.test_scripts_folder     ,						'test_performance.py'   )
												lowerCAmelCase__    						=						['accelerate', 'launch', '--num_processes=2', '--num_machines=1', '--machine_rank=0', '--use_fsdp']
												for config in self.performance_configs:
																		lowerCAmelCase__    						=						cmd.copy()
																		for i, strategy in enumerate(_UpperCamelCase   ):
																								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(_UpperCamelCase     ,						env=os.environ.copy()   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						os.path.join(self.test_scripts_folder     ,						'test_checkpointing.py'   )
												lowerCAmelCase__    						=						[
												    '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(_UpperCamelCase   ):
																		lowerCAmelCase__    						=						cmd.copy()
																		cmd_config.append(F"--fsdp_sharding_strategy={i+1}"   )
																		if strategy != "FULL_SHARD":
																								continue
																		lowerCAmelCase__    						=						len(_UpperCamelCase   )
																		for state_dict_type in FSDP_STATE_DICT_TYPE:
																								lowerCAmelCase__    						=						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(_UpperCamelCase     ,						env=os.environ.copy()   )

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







						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						os.path.join(self.test_scripts_folder     ,						'test_peak_memory_usage.py'   )
												lowerCAmelCase__    						=						[
												    '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():
																		lowerCAmelCase__    						=						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(_UpperCamelCase   ):
																														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(_UpperCamelCase     ,						env=os.environ.copy()   )













from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


__snake_case    :    Dict         =     {"""configuration_speech_encoder_decoder""": ["""SpeechEncoderDecoderConfig"""]}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__snake_case    :    Union[str, Any]         =     ["""SpeechEncoderDecoderModel"""]

try:
					if not is_flax_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__snake_case    :    str         =     ["""FlaxSpeechEncoderDecoderModel"""]

if TYPE_CHECKING:
					from .configuration_speech_encoder_decoder import SpeechEncoderDecoderConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_speech_encoder_decoder import SpeechEncoderDecoderModel

					try:
										if not is_flax_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_flax_speech_encoder_decoder import FlaxSpeechEncoderDecoderModel

else:
					import sys

					__snake_case    :    Tuple         =     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





