
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 	_lowerCAmelCase					(	__a       ):




				def      __a						(					self     )				->      List[str]:
							super().setUp()

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




				def      __a						(					self     )				->      str:
							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      __a						(					self     )				->      Any:
							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      __a						(					self     )				->      List[Any]:
							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      __a						(					self     )				->      str:
							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      __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_						=	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      __a						(					self     )				->      List[Any]:
							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 	_lowerCAmelCase					(	__a       ):




				def      __a						(					self     )				->      Optional[int]:
							super().setUp()
							lowerCAmelCase_						=	0.82
							lowerCAmelCase_						=	[
							    "fsdp_shard_grad_op_transformer_based_wrap",
							    "fsdp_full_shard_transformer_based_wrap",
							]
							lowerCAmelCase_						=	{
							    "multi_gpu_fp16": 3_200,
							    "fsdp_shard_grad_op_transformer_based_wrap_fp16": 2_000,
							    "fsdp_full_shard_transformer_based_wrap_fp16": 1_900,
							    # 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_						=	160
							lowerCAmelCase_						=	160

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




				def      __a						(					self     )				->      Tuple:
							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      __a						(					self     )				->      int:
							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      __a						(					self     )				->      Optional[int]:
							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 ...configuration_utils import PretrainedConfig
from ...utils import logging


_A   =     logging.get_logger(__name__)

_A   =     {
    # See all MEGATRON_BERT models at https://huggingface.co/models?filter=bert
}





class 	_lowerCAmelCase					(	__a       ):
				_lowercase					='''megatron-bert'''




				def __init__(					self				,				_UpperCamelCase=29_056				,				_UpperCamelCase=1_024				,				_UpperCamelCase=24				,				_UpperCamelCase=16				,				_UpperCamelCase=4_096				,				_UpperCamelCase="gelu"				,				_UpperCamelCase=0.1				,				_UpperCamelCase=0.1				,				_UpperCamelCase=512				,				_UpperCamelCase=2				,				_UpperCamelCase=0.02				,				_UpperCamelCase=1e-1_2				,				_UpperCamelCase=0				,				_UpperCamelCase="absolute"				,				_UpperCamelCase=True				,				**_UpperCamelCase				,				)				->      int:
							super().__init__(pad_token_id=_UpperCamelCase				,				**_UpperCamelCase     )

							lowerCAmelCase_						=	vocab_size
							lowerCAmelCase_						=	hidden_size
							lowerCAmelCase_						=	num_hidden_layers
							lowerCAmelCase_						=	num_attention_heads
							lowerCAmelCase_						=	hidden_act
							lowerCAmelCase_						=	intermediate_size
							lowerCAmelCase_						=	hidden_dropout_prob
							lowerCAmelCase_						=	attention_probs_dropout_prob
							lowerCAmelCase_						=	max_position_embeddings
							lowerCAmelCase_						=	type_vocab_size
							lowerCAmelCase_						=	initializer_range
							lowerCAmelCase_						=	layer_norm_eps
							lowerCAmelCase_						=	position_embedding_type
							lowerCAmelCase_						=	use_cache



