




"""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(4_2)

UpperCamelCase__						=			'''bert-base-cased'''
UpperCamelCase__						=			'''fp16'''
UpperCamelCase__						=			'''bf16'''
UpperCamelCase__						=			[FPaa, BFaa]




@require_fsdp
@require_cuda
class   a__			(       UpperCamelCase_  ):





						def       __UpperCamelCase     (							self      :	Optional[Any])						->       List[Any]:




											"""simple docstring"""







											super().setUp()

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





						def       __UpperCamelCase     (							self      :	List[str])						->       Any:




											"""simple docstring"""







											from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

											for i, strategy in enumerate(a__):
																_lowerCAmelCase:List[str]									=  self.dist_env.copy()
																_lowerCAmelCase:int									=  F'{i + 1}'
																_lowerCAmelCase:List[str]									=  strategy
																with mockenv_context(**a__):
																					_lowerCAmelCase:str									=  FullyShardedDataParallelPlugin()
																					self.assertEqual(fsdp_plugin.sharding_strategy     ,ShardingStrategy(i + 1))





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




											"""simple docstring"""







											from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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




											"""simple docstring"""







											from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

											for i, state_dict_type in enumerate(a__):
																_lowerCAmelCase:Union[str, Any]									=  self.dist_env.copy()
																_lowerCAmelCase:str									=  state_dict_type
																with mockenv_context(**a__):
																					_lowerCAmelCase:Optional[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       __UpperCamelCase     (							self      :	List[Any])						->       List[Any]:




											"""simple docstring"""







											_lowerCAmelCase:int									=  AutoModel.from_pretrained(a__)
											for policy in FSDP_AUTO_WRAP_POLICY:
																_lowerCAmelCase:Optional[int]									=  self.dist_env.copy()
																_lowerCAmelCase:str									=  policy
																if policy == "TRANSFORMER_BASED_WRAP":
																					_lowerCAmelCase:Any									=  '''BertLayer'''
																elif policy == "SIZE_BASED_WRAP":
																					_lowerCAmelCase:Any									=  '''2000'''
																with mockenv_context(**a__):
																					_lowerCAmelCase:List[str]									=  FullyShardedDataParallelPlugin()
																					fsdp_plugin.set_auto_wrap_policy(a__)
																					if policy == "NO_WRAP":
																										self.assertIsNone(fsdp_plugin.auto_wrap_policy)
																					else:
																										self.assertIsNotNone(fsdp_plugin.auto_wrap_policy)

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

											_lowerCAmelCase:List[str]									=  self.dist_env.copy()
											_lowerCAmelCase:Optional[int]									=  '''SIZE_BASED_WRAP'''
											_lowerCAmelCase:Union[str, Any]									=  '''0'''
											with mockenv_context(**a__):
																_lowerCAmelCase:Union[str, Any]									=  FullyShardedDataParallelPlugin()
																fsdp_plugin.set_auto_wrap_policy(a__)
																self.assertIsNone(fsdp_plugin.auto_wrap_policy)





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







						def       __UpperCamelCase     (							self      :	str)						->       List[Any]:




											"""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[str]									=  str(a__).lower()
																with mockenv_context(**a__):
																					_lowerCAmelCase:int									=  FullyShardedDataParallelPlugin()
																					self.assertEqual(fsdp_plugin.cpu_offload     ,CPUOffload(offload_params=a__))




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





						def       __UpperCamelCase     (							self      :	int)						->       Tuple:




											"""simple docstring"""







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

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





						def       __UpperCamelCase     (							self      :	Tuple)						->       Any:




											"""simple docstring"""







											_lowerCAmelCase:Dict									=  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:str									=  cmd.copy()
																for i, strategy in enumerate(a__):
																					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(a__     ,env=os.environ.copy())





						def       __UpperCamelCase     (							self      :	Optional[int])						->       List[str]:




											"""simple docstring"""







											_lowerCAmelCase:Any									=  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(a__):
																_lowerCAmelCase:int									=  cmd.copy()
																cmd_config.append(F'--fsdp_sharding_strategy={i+1}')
																if strategy != "FULL_SHARD":
																					continue
																_lowerCAmelCase:Optional[Any]									=  len(a__)
																for state_dict_type in FSDP_STATE_DICT_TYPE:
																					_lowerCAmelCase:Optional[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(a__     ,env=os.environ.copy())

																					_lowerCAmelCase:Tuple									=  cmd_config[:-1]
																					_lowerCAmelCase:Any									=  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(a__     ,env=os.environ.copy())







						def       __UpperCamelCase     (							self      :	Dict)						->       Any:




											"""simple docstring"""







											_lowerCAmelCase:List[str]									=  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:int									=  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(a__):
																										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(a__     ,env=os.environ.copy())












"""simple docstring"""






import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp





def 							UpperCAmelCase		(		):
					_lowerCAmelCase:Optional[int]									=  ArgumentParser(
					    description=(
					        '''PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes'''
					    )    )

					# Optional arguments for the launch helper
					parser.add_argument('''--num_cores'''      ,       type=snake_case      ,       default=1      ,       help='''Number of TPU cores to use (1 or 8).'''    )

					# positional
					parser.add_argument(
					    '''training_script'''      ,       type=snake_case      ,       help=(
					        '''The full path to the single TPU training '''
					        '''program/script to be launched in parallel, '''
					        '''followed by all the arguments for the '''
					        '''training script'''
					    )      ,       )

					# rest from the training program
					parser.add_argument('''training_script_args'''      ,       nargs=snake_case    )

					return parser.parse_args()





def 							UpperCAmelCase		(		):
					_lowerCAmelCase:int									=  parse_args()

					# Import training_script as a module.
					_lowerCAmelCase:Optional[int]									=  Path(args.training_script    )
					sys.path.append(str(script_fpath.parent.resolve()    )    )
					_lowerCAmelCase:Optional[int]									=  script_fpath.stem
					_lowerCAmelCase:Any									=  importlib.import_module(snake_case    )

					# Patch sys.argv
					_lowerCAmelCase:Dict									=  [args.training_script] + args.training_script_args + ['''--tpu_num_cores''', str(args.num_cores    )]

					xmp.spawn(mod._mp_fn      ,       args=()      ,       nprocs=args.num_cores    )


if __name__ == "__main__":
				main()






