

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 							UpperCAmelCase_						(     __lowercase  ):
							"""simple docstring"""







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

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


							def 							__lowercase							(     self			)       ->       str:
												from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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


							def 							__lowercase							(     self			)       ->       str:
												from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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


							def 							__lowercase							(     self			)       ->       List[Any]:
												from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

												for i, state_dict_type in enumerate(_a			):
																	_a      :       List[Any]   =						self.dist_env.copy()
																	_a      :       List[Any]   =						state_dict_type
																	with mockenv_context(**_a			):
																						_a      :       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 							__lowercase							(     self			)       ->       List[Any]:
												_a      :       List[Any]   =						AutoModel.from_pretrained(_a			)
												for policy in FSDP_AUTO_WRAP_POLICY:
																	_a      :       Any   =						self.dist_env.copy()
																	_a      :       str   =						policy
																	if policy == "TRANSFORMER_BASED_WRAP":
																						_a      :       Optional[Any]   =						'''BertLayer'''
																	elif policy == "SIZE_BASED_WRAP":
																						_a      :       Optional[int]   =						'''2000'''
																	with mockenv_context(**_a			):
																						_a      :       List[Any]   =						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			)

												_a      :       Optional[int]   =						self.dist_env.copy()
												_a      :       Union[str, Any]   =						'''TRANSFORMER_BASED_WRAP'''
												_a      :       Optional[Any]   =						'''T5Layer'''
												with mockenv_context(**_a			):
																	_a      :       Tuple   =						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			)			)

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


							def 							__lowercase							(     self			)       ->       Optional[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:
																	_a      :       List[str]   =						self.dist_env.copy()
																	_a      :       Optional[Any]   =						mp_dtype
																	with mockenv_context(**_a			):
																						_a      :       Tuple   =						Accelerator()
																						if mp_dtype == "fp16":
																											_a      :       Union[str, Any]   =						torch.floataa
																						elif mp_dtype == "bf16":
																											_a      :       str   =						torch.bfloataa
																						_a      :       Tuple   =						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 							__lowercase							(     self			)       ->       Optional[Any]:
												from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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



@require_fsdp
@require_multi_gpu
@slow
class 							UpperCAmelCase_						(     __lowercase  ):
							"""simple docstring"""







							def 							__lowercase							(     self			)       ->       Union[str, Any]:
												super().setUp()
												_a      :       List[Any]   =						0.82
												_a      :       str   =						[
												    '''fsdp_shard_grad_op_transformer_based_wrap''',
												    '''fsdp_full_shard_transformer_based_wrap''',
												]
												_a      :       str   =						{
												    '''multi_gpu_fp16''': 3_2_0_0,
												    '''fsdp_shard_grad_op_transformer_based_wrap_fp16''': 2_0_0_0,
												    '''fsdp_full_shard_transformer_based_wrap_fp16''': 1_9_0_0,
												    # 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
												}
												_a      :       Any   =						1_6_0
												_a      :       str   =						1_6_0

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


							def 							__lowercase							(     self			)       ->       List[Any]:
												_a      :       str   =						os.path.join(self.test_scripts_folder			,				'''test_performance.py'''			)
												_a      :       Union[str, Any]   =						['''accelerate''', '''launch''', '''--num_processes=2''', '''--num_machines=1''', '''--machine_rank=0''', '''--use_fsdp''']
												for config in self.performance_configs:
																	_a      :       Optional[int]   =						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 							__lowercase							(     self			)       ->       str:
												_a      :       List[Any]   =						os.path.join(self.test_scripts_folder			,				'''test_checkpointing.py'''			)
												_a      :       Dict   =						[
												    '''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			):
																	_a      :       int   =						cmd.copy()
																	cmd_config.append(F"""--fsdp_sharding_strategy={i+1}"""			)
																	if strategy != "FULL_SHARD":
																						continue
																	_a      :       int   =						len(_a			)
																	for state_dict_type in FSDP_STATE_DICT_TYPE:
																						_a      :       Dict   =						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()			)

																						_a      :       str   =						cmd_config[:-1]
																						_a      :       Union[str, 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 							__lowercase							(     self			)       ->       int:
												_a      :       Optional[int]   =						os.path.join(self.test_scripts_folder			,				'''test_peak_memory_usage.py'''			)
												_a      :       Tuple   =						[
												    '''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():
																	_a      :       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()			)






import os
import unittest

from transformers import LxmertTokenizer, LxmertTokenizerFast
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin

@require_tokenizers
class 							UpperCAmelCase_						(     __lowercase							,			unittest.TestCase  ):
							"""simple docstring"""







							UpperCAmelCase__			:		Any             =       LxmertTokenizer
							UpperCAmelCase__			:		Optional[Any]             =       LxmertTokenizerFast
							UpperCAmelCase__			:		Any             =       True
							UpperCAmelCase__			:		Dict             =       True


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

												_a      :       int   =						[
												    '''[UNK]''',
												    '''[CLS]''',
												    '''[SEP]''',
												    '''want''',
												    '''##want''',
												    '''##ed''',
												    '''wa''',
												    '''un''',
												    '''runn''',
												    '''##ing''',
												    ''',''',
												    '''low''',
												    '''lowest''',
												]
												_a      :       Any   =						os.path.join(self.tmpdirname			,				VOCAB_FILES_NAMES['''vocab_file''']			)
												with open(self.vocab_file			,				'''w'''			,				encoding='''utf-8'''			) as vocab_writer:
																	vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]			)			)


							def 							__lowercase							(     self			,				_a			)       ->       List[str]:
												_a      :       Tuple   =						'''UNwant\u00E9d,running'''
												_a      :       str   =						'''unwanted, running'''
												return input_text, output_text


							def 							__lowercase							(     self			)       ->       List[Any]:
												_a      :       str   =						self.tokenizer_class(self.vocab_file			)

												_a      :       str   =						tokenizer.tokenize('''UNwant\u00E9d,running'''			)
												self.assertListEqual(_a			,				['''un''', '''##want''', '''##ed''', ''',''', '''runn''', '''##ing''']			)
												self.assertListEqual(tokenizer.convert_tokens_to_ids(_a			)			,				[7, 4, 5, 1_0, 8, 9]			)




							def 							__lowercase							(     self			)       ->       List[Any]:
												if not self.test_rust_tokenizer:
																	return

												_a      :       Optional[Any]   =						self.get_tokenizer()
												_a      :       str   =						self.get_rust_tokenizer()

												_a      :       Optional[Any]   =						'''I was born in 92000, and this is falsé.'''

												_a      :       Optional[Any]   =						tokenizer.tokenize(_a			)
												_a      :       List[Any]   =						rust_tokenizer.tokenize(_a			)
												self.assertListEqual(_a			,				_a			)

												_a      :       List[Any]   =						tokenizer.encode(_a			,				add_special_tokens=_a			)
												_a      :       Any   =						rust_tokenizer.encode(_a			,				add_special_tokens=_a			)
												self.assertListEqual(_a			,				_a			)

												_a      :       Dict   =						self.get_rust_tokenizer()
												_a      :       Optional[int]   =						tokenizer.encode(_a			)
												_a      :       Dict   =						rust_tokenizer.encode(_a			)
												self.assertListEqual(_a			,				_a			)



