



'''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		(     A		):





				'''simple docstring'''





				def 				__lowerCamelCase     (							self  :  Any):







											'''simple docstring'''

											super().setUp()

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

				def 				__lowerCamelCase     (							self  :  str):







											'''simple docstring'''

											from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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

				def 				__lowerCamelCase     (							self  :  List[str]):







											'''simple docstring'''

											from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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

				def 				__lowerCamelCase     (							self  :  Optional[Any]):







											'''simple docstring'''

											from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

											for i, state_dict_type in enumerate(_UpperCAmelCase):
																		__lowercase       	=self.dist_env.copy()
																		__lowercase       	=state_dict_type
																		with mockenv_context(**_UpperCAmelCase):
																									__lowercase       	=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 				__lowerCamelCase     (							self  :  List[Any]):







											'''simple docstring'''

											__lowercase       	=AutoModel.from_pretrained(_UpperCAmelCase)
											for policy in FSDP_AUTO_WRAP_POLICY:
																		__lowercase       	=self.dist_env.copy()
																		__lowercase       	=policy
																		if policy == "TRANSFORMER_BASED_WRAP":
																									__lowercase       	='BertLayer'
																		elif policy == "SIZE_BASED_WRAP":
																									__lowercase       	='2000'
																		with mockenv_context(**_UpperCAmelCase):
																									__lowercase       	=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)

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

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

				def 				__lowerCamelCase     (							self  :  List[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:
																		__lowercase       	=self.dist_env.copy()
																		__lowercase       	=mp_dtype
																		with mockenv_context(**_UpperCAmelCase):
																									__lowercase       	=Accelerator()
																									if mp_dtype == "fp16":
																																__lowercase       	=torch.floataa
																									elif mp_dtype == "bf16":
																																__lowercase       	=torch.bfloataa
																									__lowercase       	=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 				__lowerCamelCase     (							self  :  Optional[Any]):







											'''simple docstring'''

											from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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





@require_fsdp
@require_multi_gpu
@slow
class _UpperCamelCase		(     A		):





				'''simple docstring'''





				def 				__lowerCamelCase     (							self  :  List[Any]):







											'''simple docstring'''

											super().setUp()
											__lowercase       	=0.82
											__lowercase       	=[
											    'fsdp_shard_grad_op_transformer_based_wrap',
											    'fsdp_full_shard_transformer_based_wrap',
											]
											__lowercase       	={
											    '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
											}
											__lowercase       	=1_6_0
											__lowercase       	=1_6_0

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

				def 				__lowerCamelCase     (							self  :  Tuple):







											'''simple docstring'''

											__lowercase       	=os.path.join(self.test_scripts_folder  ,       'test_performance.py')
											__lowercase       	=['accelerate', 'launch', '--num_processes=2', '--num_machines=1', '--machine_rank=0', '--use_fsdp']
											for config in self.performance_configs:
																		__lowercase       	=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 				__lowerCamelCase     (							self  :  Any):







											'''simple docstring'''

											__lowercase       	=os.path.join(self.test_scripts_folder  ,       'test_checkpointing.py')
											__lowercase       	=[
											    '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):
																		__lowercase       	=cmd.copy()
																		cmd_config.append(f"""--fsdp_sharding_strategy={i+1}""")
																		if strategy != "FULL_SHARD":
																									continue
																		__lowercase       	=len(_UpperCAmelCase)
																		for state_dict_type in FSDP_STATE_DICT_TYPE:
																									__lowercase       	=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())

																									__lowercase       	=cmd_config[:-1]
																									__lowercase       	=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 				__lowerCamelCase     (							self  :  Dict):







											'''simple docstring'''

											__lowercase       	=os.path.join(self.test_scripts_folder  ,       'test_peak_memory_usage.py')
											__lowercase       	=[
											    '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():
																		__lowercase       	=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 json
import os
import unittest

from transformers.models.blenderbot_small.tokenization_blenderbot_small import (
    VOCAB_FILES_NAMES,
    BlenderbotSmallTokenizer,
)

from ...test_tokenization_common import TokenizerTesterMixin





class _UpperCamelCase		(     A      ,       unittest.TestCase		):





				'''simple docstring'''





				lowerCAmelCase__					    =    BlenderbotSmallTokenizer
				lowerCAmelCase__					    =    False

				def 				__lowerCamelCase     (							self  :  str):







											'''simple docstring'''

											super().setUp()

											__lowercase       	=['__start__', 'adapt', 'act', 'ap@@', 'te', '__end__', '__unk__']
											__lowercase       	=dict(zip(_lowerCAmelCase  ,       range(len(_lowerCAmelCase))))

											__lowercase       	=['#version: 0.2', 'a p', 't e</w>', 'ap t</w>', 'a d', 'ad apt</w>', 'a c', 'ac t</w>', '']
											__lowercase       	={'unk_token': '__unk__', 'bos_token': '__start__', 'eos_token': '__end__'}

											__lowercase       	=os.path.join(self.tmpdirname  ,       VOCAB_FILES_NAMES['vocab_file'])
											__lowercase       	=os.path.join(self.tmpdirname  ,       VOCAB_FILES_NAMES['merges_file'])
											with open(self.vocab_file  ,       'w'  ,       encoding='utf-8') as fp:
																		fp.write(json.dumps(_lowerCAmelCase) + '\n')
											with open(self.merges_file  ,       'w'  ,       encoding='utf-8') as fp:
																		fp.write('\n'.join(_lowerCAmelCase))

				def 				__lowerCamelCase     (							self  :  int  ,       **_lowerCAmelCase  :  Tuple):







											'''simple docstring'''

											kwargs.update(self.special_tokens_map)
											return BlenderbotSmallTokenizer.from_pretrained(self.tmpdirname  ,       **_lowerCAmelCase)

				def 				__lowerCamelCase     (							self  :  Tuple  ,       _lowerCAmelCase  :  Any):







											'''simple docstring'''

											__lowercase       	='adapt act apte'
											__lowercase       	='adapt act apte'
											return input_text, output_text

				def 				__lowerCamelCase     (							self  :  Union[str, Any]):







											'''simple docstring'''

											__lowercase       	=BlenderbotSmallTokenizer(self.vocab_file  ,       self.merges_file  ,       **self.special_tokens_map)
											__lowercase       	='adapt act apte'
											__lowercase       	=['adapt', 'act', 'ap@@', 'te']
											__lowercase       	=tokenizer.tokenize(_lowerCAmelCase)
											self.assertListEqual(_lowerCAmelCase  ,       _lowerCAmelCase)

											__lowercase       	=[tokenizer.bos_token] + tokens + [tokenizer.eos_token]

											__lowercase       	=[0, 1, 2, 3, 4, 5]
											self.assertListEqual(tokenizer.convert_tokens_to_ids(_lowerCAmelCase)  ,       _lowerCAmelCase)

				def 				__lowerCamelCase     (							self  :  Union[str, Any]):







											'''simple docstring'''

											__lowercase       	=BlenderbotSmallTokenizer.from_pretrained('facebook/blenderbot-90M')
											assert tok('sam').input_ids == [1_3_8_4]
											__lowercase       	='I am a small frog.'
											__lowercase       	=tok([src_text]  ,       padding=_lowerCAmelCase  ,       truncation=_lowerCAmelCase)['input_ids']
											__lowercase       	=tok.batch_decode(_lowerCAmelCase  ,       skip_special_tokens=_lowerCAmelCase  ,       clean_up_tokenization_spaces=_lowerCAmelCase)[0]
											assert src_text != decoded  # I wish it did!
											assert decoded == "i am a small frog ."

				def 				__lowerCamelCase     (							self  :  str):







											'''simple docstring'''

											__lowercase       	=BlenderbotSmallTokenizer.from_pretrained('facebook/blenderbot-90M')
											__lowercase       	='I am a small frog .'
											__lowercase       	='.'
											__lowercase       	=tok(_lowerCAmelCase)['input_ids']
											__lowercase       	=tok(_lowerCAmelCase)['input_ids']

											assert encoded[-1] == encoded_dot[0]



