





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)

_lowercase		:  Any			       ="""bert-base-cased"""
_lowercase		:  List[Any]			       ="""fp16"""
_lowercase		:  List[Any]			       ="""bf16"""
_lowercase		:  str			       =[FPaa, BFaa]






@require_fsdp
@require_cuda
class 						UpperCamelCase_  (  snake_case__							):



					def   __a     (		self     :     int						):
							super().setUp()

							lowerCamelCase_	:     Union[str, Any] 						=      dict(
							    ACCELERATE_USE_FSDP='true'       ,      MASTER_ADDR='localhost'       ,      MASTER_PORT='10999'       ,      RANK='0'       ,      LOCAL_RANK='0'       ,      WORLD_SIZE='1'       ,      )



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

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



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

							for i, prefetch_policy in enumerate(lowerCamelCase						):
									lowerCamelCase_	:     Optional[Any] 						=      self.dist_env.copy()
									lowerCamelCase_	:     Optional[Any] 						=      prefetch_policy
									with mockenv_context(**lowerCamelCase						):
											lowerCamelCase_	:     Dict 						=      FullyShardedDataParallelPlugin()
											if prefetch_policy == "NO_PREFETCH":
													self.assertIsNone(fsdp_plugin.backward_prefetch						)
											else:
													self.assertEqual(fsdp_plugin.backward_prefetch       ,      BackwardPrefetch(i + 1						)						)



					def   __a     (		self     :     Any						):
							from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

							for i, state_dict_type in enumerate(lowerCamelCase						):
									lowerCamelCase_	:     Optional[Any] 						=      self.dist_env.copy()
									lowerCamelCase_	:     Any 						=      state_dict_type
									with mockenv_context(**lowerCamelCase						):
											lowerCamelCase_	:     int 						=      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   __a     (		self     :     str						):
							lowerCamelCase_	:     List[Any] 						=      AutoModel.from_pretrained(lowerCamelCase						)
							for policy in FSDP_AUTO_WRAP_POLICY:
									lowerCamelCase_	:     Union[str, Any] 						=      self.dist_env.copy()
									lowerCamelCase_	:     Union[str, Any] 						=      policy
									if policy == "TRANSFORMER_BASED_WRAP":
											lowerCamelCase_	:     Optional[Any] 						=      'BertLayer'
									elif policy == "SIZE_BASED_WRAP":
											lowerCamelCase_	:     Optional[int] 						=      '2000'
									with mockenv_context(**lowerCamelCase						):
											lowerCamelCase_	:     Dict 						=      FullyShardedDataParallelPlugin()
											fsdp_plugin.set_auto_wrap_policy(lowerCamelCase						)
											if policy == "NO_WRAP":
													self.assertIsNone(fsdp_plugin.auto_wrap_policy						)
											else:
													self.assertIsNotNone(fsdp_plugin.auto_wrap_policy						)

							lowerCamelCase_	:     List[Any] 						=      self.dist_env.copy()
							lowerCamelCase_	:     Union[str, Any] 						=      'TRANSFORMER_BASED_WRAP'
							lowerCamelCase_	:     Any 						=      'T5Layer'
							with mockenv_context(**lowerCamelCase						):
									lowerCamelCase_	:     List[Any] 						=      FullyShardedDataParallelPlugin()
									with self.assertRaises(lowerCamelCase						) as cm:
											fsdp_plugin.set_auto_wrap_policy(lowerCamelCase						)
									self.assertTrue('Could not find the transformer layer class to wrap in the model.' in str(cm.exception						)						)

							lowerCamelCase_	:     Dict 						=      self.dist_env.copy()
							lowerCamelCase_	:     int 						=      'SIZE_BASED_WRAP'
							lowerCamelCase_	:     str 						=      '0'
							with mockenv_context(**lowerCamelCase						):
									lowerCamelCase_	:     Optional[int] 						=      FullyShardedDataParallelPlugin()
									fsdp_plugin.set_auto_wrap_policy(lowerCamelCase						)
									self.assertIsNone(fsdp_plugin.auto_wrap_policy						)



					def   __a     (		self     :     Tuple						):
							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_	:     Optional[int] 						=      self.dist_env.copy()
									lowerCamelCase_	:     Tuple 						=      mp_dtype
									with mockenv_context(**lowerCamelCase						):
											lowerCamelCase_	:     Optional[Any] 						=      Accelerator()
											if mp_dtype == "fp16":
													lowerCamelCase_	:     List[Any] 						=      torch.floataa
											elif mp_dtype == "bf16":
													lowerCamelCase_	:     Any 						=      torch.bfloataa
											lowerCamelCase_	:     Dict 						=      MixedPrecision(param_dtype=lowerCamelCase       ,      reduce_dtype=lowerCamelCase       ,      buffer_dtype=lowerCamelCase						)
											self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy       ,      lowerCamelCase						)
											if mp_dtype == FPaa:
													self.assertTrue(isinstance(accelerator.scaler       ,      lowerCamelCase						)						)
											elif mp_dtype == BFaa:
													self.assertIsNone(accelerator.scaler						)
											AcceleratorState._reset_state(lowerCamelCase						)





					def   __a     (		self     :     Dict						):
							from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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






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



					def   __a     (		self     :     List[str]						):
							super().setUp()
							lowerCamelCase_	:     str 						=      0.82
							lowerCamelCase_	:     List[str] 						=      [
							    'fsdp_shard_grad_op_transformer_based_wrap',
							    'fsdp_full_shard_transformer_based_wrap',
							]
							lowerCamelCase_	:     Any 						=      {
							    '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_	:     Optional[int] 						=      1_60
							lowerCamelCase_	:     Any 						=      1_60

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



					def   __a     (		self     :     Any						):
							lowerCamelCase_	:     int 						=      os.path.join(self.test_scripts_folder       ,      'test_performance.py'						)
							lowerCamelCase_	:     Optional[Any] 						=      ['accelerate', 'launch', '--num_processes=2', '--num_machines=1', '--machine_rank=0', '--use_fsdp']
							for config in self.performance_configs:
									lowerCamelCase_	:     List[Any] 						=      cmd.copy()
									for i, strategy in enumerate(lowerCamelCase						):
											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(lowerCamelCase       ,      env=os.environ.copy()						)



					def   __a     (		self     :     Dict						):
							lowerCamelCase_	:     Tuple 						=      os.path.join(self.test_scripts_folder       ,      'test_checkpointing.py'						)
							lowerCamelCase_	:     List[Any] 						=      [
							    '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(lowerCamelCase						):
									lowerCamelCase_	:     Union[str, Any] 						=      cmd.copy()
									cmd_config.append(F"--fsdp_sharding_strategy={i+1}"						)
									if strategy != "FULL_SHARD":
											continue
									lowerCamelCase_	:     Tuple 						=      len(lowerCamelCase						)
									for state_dict_type in FSDP_STATE_DICT_TYPE:
											lowerCamelCase_	:     int 						=      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(lowerCamelCase       ,      env=os.environ.copy()						)

											lowerCamelCase_	:     Optional[Any] 						=      cmd_config[:-1]
											lowerCamelCase_	:     Dict 						=      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(lowerCamelCase       ,      env=os.environ.copy()						)





					def   __a     (		self     :     Optional[int]						):
							lowerCamelCase_	:     Optional[Any] 						=      os.path.join(self.test_scripts_folder       ,      'test_peak_memory_usage.py'						)
							lowerCamelCase_	:     Optional[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():
									lowerCamelCase_	:     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(lowerCamelCase						):
													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(lowerCamelCase       ,      env=os.environ.copy()						)







_lowercase		:  Optional[Any]			       ="""
# Transformers installation
! pip install transformers datasets
# To install from source instead of the last release, comment the command above and uncomment the following one.
# ! pip install git+https://github.com/huggingface/transformers.git
"""

_lowercase		:  Union[str, Any]			       =[{"""type""": """code""", """content""": INSTALL_CONTENT}]
_lowercase		:  List[Any]			       ={
    """{processor_class}""": """FakeProcessorClass""",
    """{model_class}""": """FakeModelClass""",
    """{object_class}""": """FakeObjectClass""",
}
