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

__lowerCamelCase											=  "bert-base-cased"
__lowerCamelCase											=  "fp16"
__lowerCamelCase											=  "bf16"
__lowerCamelCase											=  [FPaa, BFaa]
@require_fsdp
@require_cuda
class     UpperCamelCase__(			__lowerCamelCase ):





						def 				snake_case__	(       self     )					->				Dict:
													super().setUp()

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





						def 				snake_case__	(       self     )					->				Tuple:
													from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

													for i, strategy in enumerate(UpperCAmelCase_     ):
																				A__	     =							self.dist_env.copy()
																				A__	     =							f'''{i + 1}'''
																				A__	     =							strategy
																				with mockenv_context(**UpperCAmelCase_     ):
																											A__	     =							FullyShardedDataParallelPlugin()
																											self.assertEqual(fsdp_plugin.sharding_strategy		,ShardingStrategy(i + 1     )     )





						def 				snake_case__	(       self     )					->				Optional[int]:
													from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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





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

													for i, state_dict_type in enumerate(UpperCAmelCase_     ):
																				A__	     =							self.dist_env.copy()
																				A__	     =							state_dict_type
																				with mockenv_context(**UpperCAmelCase_     ):
																											A__	     =							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 				snake_case__	(       self     )					->				Dict:
													A__	     =							AutoModel.from_pretrained(UpperCAmelCase_     )
													for policy in FSDP_AUTO_WRAP_POLICY:
																				A__	     =							self.dist_env.copy()
																				A__	     =							policy
																				if policy == "TRANSFORMER_BASED_WRAP":
																											A__	     =							'BertLayer'
																				elif policy == "SIZE_BASED_WRAP":
																											A__	     =							'2000'
																				with mockenv_context(**UpperCAmelCase_     ):
																											A__	     =							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     )

													A__	     =							self.dist_env.copy()
													A__	     =							'TRANSFORMER_BASED_WRAP'
													A__	     =							'T5Layer'
													with mockenv_context(**UpperCAmelCase_     ):
																				A__	     =							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     )     )

													A__	     =							self.dist_env.copy()
													A__	     =							'SIZE_BASED_WRAP'
													A__	     =							'0'
													with mockenv_context(**UpperCAmelCase_     ):
																				A__	     =							FullyShardedDataParallelPlugin()
																				fsdp_plugin.set_auto_wrap_policy(UpperCAmelCase_     )
																				self.assertIsNone(fsdp_plugin.auto_wrap_policy     )





						def 				snake_case__	(       self     )					->				Dict:
													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__	     =							self.dist_env.copy()
																				A__	     =							mp_dtype
																				with mockenv_context(**UpperCAmelCase_     ):
																											A__	     =							Accelerator()
																											if mp_dtype == "fp16":
																																		A__	     =							torch.floataa
																											elif mp_dtype == "bf16":
																																		A__	     =							torch.bfloataa
																											A__	     =							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 				snake_case__	(       self     )					->				Optional[int]:
													from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

													for flag in [True, False]:
																				A__	     =							self.dist_env.copy()
																				A__	     =							str(UpperCAmelCase_     ).lower()
																				with mockenv_context(**UpperCAmelCase_     ):
																											A__	     =							FullyShardedDataParallelPlugin()
																											self.assertEqual(fsdp_plugin.cpu_offload		,CPUOffload(offload_params=UpperCAmelCase_     )     )






@require_fsdp
@require_multi_gpu
@slow
class     UpperCamelCase__(			__lowerCamelCase ):





						def 				snake_case__	(       self     )					->				Union[str, Any]:
													super().setUp()
													A__	     =							0.8_2
													A__	     =							[
													    'fsdp_shard_grad_op_transformer_based_wrap',
													    'fsdp_full_shard_transformer_based_wrap',
													]
													A__	     =							{
													    '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
													}
													A__	     =							1_60
													A__	     =							1_60

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





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

																											A__	     =							cmd_config[:-1]
																											A__	     =							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 				snake_case__	(       self     )					->				str:
													A__	     =							os.path.join(self.test_scripts_folder		,'test_peak_memory_usage.py'     )
													A__	     =							[
													    '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__	     =							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()     )






"""simple docstring"""




import os
import unittest

from transformers.models.cpmant.tokenization_cpmant import VOCAB_FILES_NAMES, CpmAntTokenizer
from transformers.testing_utils import require_jieba, tooslow

from ...test_tokenization_common import TokenizerTesterMixin
@require_jieba
class     UpperCamelCase__(			__A		,						unittest.TestCase ):
						lowerCAmelCase__  :						Optional[int]													=		CpmAntTokenizer
						lowerCAmelCase__  :						List[str]													=		False





						def 				snake_case__	(       self     )					->				Dict:
													super().setUp()

													A__	     =							[
													    '<d>',
													    '</d>',
													    '<s>',
													    '</s>',
													    '</_>',
													    '<unk>',
													    '<pad>',
													    '</n>',
													    '我',
													    '是',
													    'C',
													    'P',
													    'M',
													    'A',
													    'n',
													    't',
													]
													A__	     =							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]     )     )





						@tooslow
						def 				snake_case__	(       self     )					->				List[Any]:
													A__	     =							CpmAntTokenizer.from_pretrained('openbmb/cpm-ant-10b'     )
													A__	     =							'今天天气真好！'
													A__	     =							['今天', '天气', '真', '好', '！']
													A__	     =							tokenizer.tokenize(__UpperCAmelCase     )
													self.assertListEqual(__UpperCAmelCase		,__UpperCAmelCase     )
													A__	     =							'今天天气真好！'
													A__	     =							[tokenizer.bos_token] + tokens

													A__	     =							[6, 98_02, 1_49_62, 20_82, 8_31, 2_44]
													self.assertListEqual(tokenizer.convert_tokens_to_ids(__UpperCAmelCase     )		,__UpperCAmelCase     )

													A__	     =							tokenizer.decode(__UpperCAmelCase     )
													self.assertEqual(__UpperCAmelCase		,__UpperCAmelCase     )





