



"""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)

__snake_case							:   List[Any]      =    'bert-base-cased'
__snake_case							:   Tuple      =    'fp16'
__snake_case							:   str      =    'bf16'
__snake_case							:   Optional[int]      =    [FPaa, BFaa]



@require_fsdp
@require_cuda
class 			A__   (  __SCREAMING_SNAKE_CASE							):




	'''simple docstring'''





	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str])				->       Union[str, Any]:

				"""simple docstring"""

				super().setUp()

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

	def 							_SCREAMING_SNAKE_CASE   (     self:		int)				->       Optional[Any]:

				"""simple docstring"""

				from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

				for i, strategy in enumerate(_SCREAMING_SNAKE_CASE):
							__lowerCAmelCase      :  Tuple												=       self.dist_env.copy()
							__lowerCAmelCase      :  Dict												=       F"""{i + 1}"""
							__lowerCAmelCase      :  Union[str, Any]												=       strategy
							with mockenv_context(**_SCREAMING_SNAKE_CASE):
										__lowerCAmelCase      :  str												=       FullyShardedDataParallelPlugin()
										self.assertEqual(fsdp_plugin.sharding_strategy       ,			ShardingStrategy(i + 1))

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str])				->       Union[str, Any]:

				"""simple docstring"""

				from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

				for i, prefetch_policy in enumerate(_SCREAMING_SNAKE_CASE):
							__lowerCAmelCase      :  int												=       self.dist_env.copy()
							__lowerCAmelCase      :  List[str]												=       prefetch_policy
							with mockenv_context(**_SCREAMING_SNAKE_CASE):
										__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 							_SCREAMING_SNAKE_CASE   (     self:		int)				->       Union[str, Any]:

				"""simple docstring"""

				from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

				for i, state_dict_type in enumerate(_SCREAMING_SNAKE_CASE):
							__lowerCAmelCase      :  Union[str, Any]												=       self.dist_env.copy()
							__lowerCAmelCase      :  Optional[Any]												=       state_dict_type
							with mockenv_context(**_SCREAMING_SNAKE_CASE):
										__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 							_SCREAMING_SNAKE_CASE   (     self:		Dict)				->       Dict:

				"""simple docstring"""

				__lowerCAmelCase      :  List[str]												=       AutoModel.from_pretrained(_SCREAMING_SNAKE_CASE)
				for policy in FSDP_AUTO_WRAP_POLICY:
							__lowerCAmelCase      :  List[str]												=       self.dist_env.copy()
							__lowerCAmelCase      :  List[str]												=       policy
							if policy == "TRANSFORMER_BASED_WRAP":
										__lowerCAmelCase      :  List[str]												=       "BertLayer"
							elif policy == "SIZE_BASED_WRAP":
										__lowerCAmelCase      :  Any												=       "2000"
							with mockenv_context(**_SCREAMING_SNAKE_CASE):
										__lowerCAmelCase      :  Dict												=       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)

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

				__lowerCAmelCase      :  Any												=       self.dist_env.copy()
				__lowerCAmelCase      :  Any												=       "SIZE_BASED_WRAP"
				__lowerCAmelCase      :  str												=       "0"
				with mockenv_context(**_SCREAMING_SNAKE_CASE):
							__lowerCAmelCase      :  Tuple												=       FullyShardedDataParallelPlugin()
							fsdp_plugin.set_auto_wrap_policy(_SCREAMING_SNAKE_CASE)
							self.assertIsNone(fsdp_plugin.auto_wrap_policy)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[Any])				->       Union[str, Any]:

				"""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      :  Union[str, Any]												=       self.dist_env.copy()
							__lowerCAmelCase      :  Tuple												=       mp_dtype
							with mockenv_context(**_SCREAMING_SNAKE_CASE):
										__lowerCAmelCase      :  Any												=       Accelerator()
										if mp_dtype == "fp16":
													__lowerCAmelCase      :  Optional[Any]												=       torch.floataa
										elif mp_dtype == "bf16":
													__lowerCAmelCase      :  List[Any]												=       torch.bfloataa
										__lowerCAmelCase      :  int												=       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 							_SCREAMING_SNAKE_CASE   (     self:		Optional[Any])				->       int:

				"""simple docstring"""

				from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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






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




	'''simple docstring'''





	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str])				->       List[str]:

				"""simple docstring"""

				super().setUp()
				__lowerCAmelCase      :  Optional[Any]												=       0.82
				__lowerCAmelCase      :  Dict												=       [
				    "fsdp_shard_grad_op_transformer_based_wrap",
				    "fsdp_full_shard_transformer_based_wrap",
				]
				__lowerCAmelCase      :  List[Any]												=       {
				    "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
				}
				__lowerCAmelCase      :  Tuple												=       160
				__lowerCAmelCase      :  Any												=       160

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

	def 							_SCREAMING_SNAKE_CASE   (     self:		Any)				->       List[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  Union[str, Any]												=       os.path.join(self.test_scripts_folder       ,			"test_performance.py")
				__lowerCAmelCase      :  Tuple												=       ["accelerate", "launch", "--num_processes=2", "--num_machines=1", "--machine_rank=0", "--use_fsdp"]
				for config in self.performance_configs:
							__lowerCAmelCase      :  Tuple												=       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 							_SCREAMING_SNAKE_CASE   (     self:		List[Any])				->       Dict:

				"""simple docstring"""

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

										__lowerCAmelCase      :  Tuple												=       cmd_config[:-1]
										__lowerCAmelCase      :  List[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(_SCREAMING_SNAKE_CASE       ,			env=os.environ.copy())







	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any])				->       Any:

				"""simple docstring"""

				__lowerCAmelCase      :  Union[str, Any]												=       os.path.join(self.test_scripts_folder       ,			"test_peak_memory_usage.py")
				__lowerCAmelCase      :  str												=       [
				    "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      :  Dict												=       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"""






#
# This a `torch.distributed` diagnostics script that checks that all GPUs in the cluster (one or
# many nodes) can talk to each other via nccl and allocate gpu memory.
#
# To run first adjust the number of processes and nodes:
#
# python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# You may need to add --master_addr $MASTER_ADDR --master_port $MASTER_PORT if using a custom addr:port
#
# You can also use the rdzv API: --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT --rdzv_backend c10d
#
# use torch.distributed.launch instead of torch.distributed.run for torch < 1.9
#
# If you get a hanging in `barrier` calls you have some network issues, you may try to debug this with:
#
# NCCL_DEBUG=INFO python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# which should tell you what's going on behind the scenes.
#
#
# This script can be run via `srun` in the SLURM environment as well. Here is a SLURM script that
# runs on 2 nodes of 4 gpus per node:
#
# #SBATCH --job-name=test-nodes        # name
# #SBATCH --nodes=2                    # nodes
# #SBATCH --ntasks-per-node=1          # crucial - only 1 task per dist per node!
# #SBATCH --cpus-per-task=10           # number of cores per tasks
# #SBATCH --gres=gpu:4                 # number of gpus
# #SBATCH --time 0:05:00               # maximum execution time (HH:MM:SS)
# #SBATCH --output=%x-%j.out           # output file name
#
# GPUS_PER_NODE=4
# MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
# MASTER_PORT=6000
#
# srun --jobid $SLURM_JOBID bash -c 'python -m torch.distributed.run \
# --nproc_per_node $GPUS_PER_NODE --nnodes $SLURM_NNODES --node_rank $SLURM_PROCID \
# --master_addr $MASTER_ADDR --master_port $MASTER_PORT \
# torch-distributed-gpu-test.py'
#

import fcntl
import os
import socket

import torch
import torch.distributed as dist


def 			_lowercase				(  *__snake_case	)						->   Optional[Any]:
			with open(__snake_case					,"r"	) as fh:
						fcntl.flock(__snake_case					,fcntl.LOCK_EX	)
						try:
									print(*__snake_case	)
						finally:
									fcntl.flock(__snake_case					,fcntl.LOCK_UN	)


__snake_case							:   List[Any]      =    int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
__snake_case							:   List[str]      =    torch.device('cuda', local_rank)
__snake_case							:   Optional[Any]      =    socket.gethostname()

__snake_case							:   str      =    F"""[{hostname}-{local_rank}]"""

try:
	# test distributed
	dist.init_process_group('nccl')
	dist.all_reduce(torch.ones(1).to(device), op=dist.ReduceOp.SUM)
	dist.barrier()

	# test cuda is available and can allocate memory
	torch.cuda.is_available()
	torch.ones(1).cuda(local_rank)

	# global rank
	__snake_case							:   Tuple      =    dist.get_rank()
	__snake_case							:   Any      =    dist.get_world_size()

	printflock(F"""{gpu} is OK (global rank: {rank}/{world_size})""")

	dist.barrier()
	if rank == 0:
		printflock(F"""pt={torch.__version__}, cuda={torch.version.cuda}, nccl={torch.cuda.nccl.version()}""")

except Exception:
	printflock(F"""{gpu} is broken""")
	raise