
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__					:	Dict							       =					'bert-base-cased'
UpperCAmelCase__					:	Optional[int]							       =					'fp16'
UpperCAmelCase__					:	Union[str, Any]							       =					'bf16'
UpperCAmelCase__					:	List[Any]							       =					[FPaa, BFaa]



@require_fsdp
@require_cuda
class       UpperCamelCase_      (    snake_case_ ):




							'''simple docstring'''

							def 							lowerCAmelCase__			(					self)			->						Union[str, Any]:
										super().setUp()

										UpperCamelCase__     :		Tuple	 =					dict(
										    ACCELERATE_USE_FSDP='true'  ,    MASTER_ADDR='localhost'  ,    MASTER_PORT='10999'  ,    RANK='0'  ,    LOCAL_RANK='0'  ,    WORLD_SIZE='1'  ,    )
							def 							lowerCAmelCase__			(					self)			->						str:
										from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

										for i, strategy in enumerate(__lowercase):
													UpperCamelCase__     :		Optional[Any]	 =					self.dist_env.copy()
													UpperCamelCase__     :		Any	 =					F"""{i + 1}"""
													UpperCamelCase__     :		Union[str, Any]	 =					strategy
													with mockenv_context(**__lowercase):
																UpperCamelCase__     :		Any	 =					FullyShardedDataParallelPlugin()
																self.assertEqual(fsdp_plugin.sharding_strategy  ,    ShardingStrategy(i + 1))
							def 							lowerCAmelCase__			(					self)			->						List[str]:
										from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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

										for i, state_dict_type in enumerate(__lowercase):
													UpperCamelCase__     :		List[Any]	 =					self.dist_env.copy()
													UpperCamelCase__     :		Tuple	 =					state_dict_type
													with mockenv_context(**__lowercase):
																UpperCamelCase__     :		Optional[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)			->						List[Any]:
										UpperCamelCase__     :		str	 =					AutoModel.from_pretrained(__lowercase)
										for policy in FSDP_AUTO_WRAP_POLICY:
													UpperCamelCase__     :		List[Any]	 =					self.dist_env.copy()
													UpperCamelCase__     :		Tuple	 =					policy
													if policy == "TRANSFORMER_BASED_WRAP":
																UpperCamelCase__     :		Dict	 =					'BertLayer'
													elif policy == "SIZE_BASED_WRAP":
																UpperCamelCase__     :		Tuple	 =					'2000'
													with mockenv_context(**__lowercase):
																UpperCamelCase__     :		Optional[int]	 =					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__     :		Tuple	 =					self.dist_env.copy()
										UpperCamelCase__     :		Tuple	 =					'TRANSFORMER_BASED_WRAP'
										UpperCamelCase__     :		Optional[Any]	 =					'T5Layer'
										with mockenv_context(**__lowercase):
													UpperCamelCase__     :		int	 =					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__     :		Union[str, Any]	 =					self.dist_env.copy()
										UpperCamelCase__     :		str	 =					'SIZE_BASED_WRAP'
										UpperCamelCase__     :		int	 =					'0'
										with mockenv_context(**__lowercase):
													UpperCamelCase__     :		Any	 =					FullyShardedDataParallelPlugin()
													fsdp_plugin.set_auto_wrap_policy(__lowercase)
													self.assertIsNone(fsdp_plugin.auto_wrap_policy)
							def 							lowerCAmelCase__			(					self)			->						int:
										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__     :		Dict	 =					self.dist_env.copy()
													UpperCamelCase__     :		Optional[int]	 =					mp_dtype
													with mockenv_context(**__lowercase):
																UpperCamelCase__     :		Optional[Any]	 =					Accelerator()
																if mp_dtype == "fp16":
																			UpperCamelCase__     :		List[str]	 =					torch.floataa
																elif mp_dtype == "bf16":
																			UpperCamelCase__     :		List[Any]	 =					torch.bfloataa
																UpperCamelCase__     :		List[str]	 =					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 							lowerCAmelCase__			(					self)			->						List[Any]:
										from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

										for flag in [True, False]:
													UpperCamelCase__     :		int	 =					self.dist_env.copy()
													UpperCamelCase__     :		Dict	 =					str(__lowercase).lower()
													with mockenv_context(**__lowercase):
																UpperCamelCase__     :		Optional[Any]	 =					FullyShardedDataParallelPlugin()
																self.assertEqual(fsdp_plugin.cpu_offload  ,    CPUOffload(offload_params=__lowercase))



@require_fsdp
@require_multi_gpu
@slow
class       UpperCamelCase_      (    snake_case_ ):




							'''simple docstring'''

							def 							lowerCAmelCase__			(					self)			->						int:
										super().setUp()
										UpperCamelCase__     :		Dict	 =					0.82
										UpperCamelCase__     :		int	 =					[
										    'fsdp_shard_grad_op_transformer_based_wrap',
										    'fsdp_full_shard_transformer_based_wrap',
										]
										UpperCamelCase__     :		Tuple	 =					{
										    '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__     :		Union[str, Any]	 =					1_60
										UpperCamelCase__     :		List[Any]	 =					1_60

										UpperCamelCase__     :		Union[str, Any]	 =					inspect.getfile(accelerate.test_utils)
										UpperCamelCase__     :		Optional[Any]	 =					os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ['scripts', 'external_deps'])
							def 							lowerCAmelCase__			(					self)			->						int:
										UpperCamelCase__     :		Tuple	 =					os.path.join(self.test_scripts_folder  ,    'test_performance.py')
										UpperCamelCase__     :		Optional[Any]	 =					['accelerate', 'launch', '--num_processes=2', '--num_machines=1', '--machine_rank=0', '--use_fsdp']
										for config in self.performance_configs:
													UpperCamelCase__     :		List[str]	 =					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 							lowerCAmelCase__			(					self)			->						List[Any]:
										UpperCamelCase__     :		Tuple	 =					os.path.join(self.test_scripts_folder  ,    'test_checkpointing.py')
										UpperCamelCase__     :		int	 =					[
										    '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__     :		str	 =					cmd.copy()
													cmd_config.append(F"""--fsdp_sharding_strategy={i+1}""")
													if strategy != "FULL_SHARD":
																continue
													UpperCamelCase__     :		Union[str, Any]	 =					len(__lowercase)
													for state_dict_type in FSDP_STATE_DICT_TYPE:
																UpperCamelCase__     :		Tuple	 =					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__     :		Union[str, Any]	 =					cmd_config[:-1]
																UpperCamelCase__     :		str	 =					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 							lowerCAmelCase__			(					self)			->						Any:
										UpperCamelCase__     :		Optional[Any]	 =					os.path.join(self.test_scripts_folder  ,    'test_peak_memory_usage.py')
										UpperCamelCase__     :		Union[str, Any]	 =					[
										    '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__     :		str	 =					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'''
from typing import TYPE_CHECKING

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


UpperCAmelCase_			:  Optional[Any]										=   {
    'configuration_transfo_xl': ['TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TransfoXLConfig'],
    'tokenization_transfo_xl': ['TransfoXLCorpus', 'TransfoXLTokenizer'],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCAmelCase_			:  Any										=   [
     'TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST',
     'AdaptiveEmbedding',
     'TransfoXLForSequenceClassification',
     'TransfoXLLMHeadModel',
     'TransfoXLModel',
     'TransfoXLPreTrainedModel',
     'load_tf_weights_in_transfo_xl',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCAmelCase_			:  Optional[Any]										=   [
     'TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST',
     'TFAdaptiveEmbedding',
     'TFTransfoXLForSequenceClassification',
     'TFTransfoXLLMHeadModel',
     'TFTransfoXLMainLayer',
     'TFTransfoXLModel',
     'TFTransfoXLPreTrainedModel',
 ]


if TYPE_CHECKING:
 from .configuration_transfo_xl import TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP, TransfoXLConfig
 from .tokenization_transfo_xl import TransfoXLCorpus, TransfoXLTokenizer

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_transfo_xl import (
      TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
      AdaptiveEmbedding,
      TransfoXLForSequenceClassification,
      TransfoXLLMHeadModel,
      TransfoXLModel,
      TransfoXLPreTrainedModel,
      load_tf_weights_in_transfo_xl,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_transfo_xl import (
      TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFAdaptiveEmbedding,
      TFTransfoXLForSequenceClassification,
      TFTransfoXLLMHeadModel,
      TFTransfoXLMainLayer,
      TFTransfoXLModel,
      TFTransfoXLPreTrainedModel,
  )

else:
 import sys

 UpperCAmelCase_			:  Optional[int]										=   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






