




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

_A													=      """bert-base-cased"""
_A													=      """fp16"""
_A													=      """bf16"""
_A													=      [FPaa, BFaa]


@require_fsdp
@require_cuda
class _lowerCamelCase    (					a__						):






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

         """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 	_lowerCAmelCase   (		self       :		Optional[Any]   )			->						List[str]:

         """simple docstring"""
         from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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






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

         """simple docstring"""
         from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

         for i, prefetch_policy in enumerate(SCREAMING_SNAKE_CASE_   ):
               lowerCAmelCase__      :			Dict   =	self.dist_env.copy()
               lowerCAmelCase__      :			List[str]   =	prefetch_policy
               with mockenv_context(**SCREAMING_SNAKE_CASE_   ):
                     lowerCAmelCase__      :			Tuple   =	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[int]   )			->						Tuple:

         """simple docstring"""
         from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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

         """simple docstring"""
         lowerCAmelCase__      :			int   =	AutoModel.from_pretrained(SCREAMING_SNAKE_CASE_   )
         for policy in FSDP_AUTO_WRAP_POLICY:
               lowerCAmelCase__      :			Optional[Any]   =	self.dist_env.copy()
               lowerCAmelCase__      :			List[Any]   =	policy
               if policy == "TRANSFORMER_BASED_WRAP":
                     lowerCAmelCase__      :			Optional[int]   =	'BertLayer'
               elif policy == "SIZE_BASED_WRAP":
                     lowerCAmelCase__      :			List[str]   =	'2000'
               with mockenv_context(**SCREAMING_SNAKE_CASE_   ):
                     lowerCAmelCase__      :			Union[str, Any]   =	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__      :			Optional[Any]   =	self.dist_env.copy()
         lowerCAmelCase__      :			Dict   =	'TRANSFORMER_BASED_WRAP'
         lowerCAmelCase__      :			Optional[int]   =	'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__      :			Tuple   =	'SIZE_BASED_WRAP'
         lowerCAmelCase__      :			Optional[Any]   =	'0'
         with mockenv_context(**SCREAMING_SNAKE_CASE_   ):
               lowerCAmelCase__      :			str   =	FullyShardedDataParallelPlugin()
               fsdp_plugin.set_auto_wrap_policy(SCREAMING_SNAKE_CASE_   )
               self.assertIsNone(fsdp_plugin.auto_wrap_policy   )






   def 	_lowerCAmelCase   (		self       :		List[str]   )			->						Tuple:

         """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__      :			Dict   =	self.dist_env.copy()
               lowerCAmelCase__      :			str   =	mp_dtype
               with mockenv_context(**SCREAMING_SNAKE_CASE_   ):
                     lowerCAmelCase__      :			List[Any]   =	Accelerator()
                     if mp_dtype == "fp16":
                           lowerCAmelCase__      :			Dict   =	torch.floataa
                     elif mp_dtype == "bf16":
                           lowerCAmelCase__      :			str   =	torch.bfloataa
                     lowerCAmelCase__      :			List[Any]   =	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 	_lowerCAmelCase   (		self       :		Tuple   )			->						int:

         """simple docstring"""
         from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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






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






   def 	_lowerCAmelCase   (		self       :		Optional[int]   )			->						Dict:

         """simple docstring"""
         super().setUp()
         lowerCAmelCase__      :			List[str]   =	0.82
         lowerCAmelCase__      :			List[Any]   =	[
             'fsdp_shard_grad_op_transformer_based_wrap',
             'fsdp_full_shard_transformer_based_wrap',
         ]
         lowerCAmelCase__      :			List[Any]   =	{
             '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
         }
         lowerCAmelCase__      :			Dict   =	1_60
         lowerCAmelCase__      :			Optional[int]   =	1_60

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






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

         """simple docstring"""
         lowerCAmelCase__      :			Optional[int]   =	os.path.join(self.test_scripts_folder   ,     """test_performance.py"""   )
         lowerCAmelCase__      :			str   =	['accelerate', 'launch', '--num_processes=2', '--num_machines=1', '--machine_rank=0', '--use_fsdp']
         for config in self.performance_configs:
               lowerCAmelCase__      :			List[Any]   =	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 	_lowerCAmelCase   (		self       :		Union[str, Any]   )			->						Optional[int]:

         """simple docstring"""
         lowerCAmelCase__      :			str   =	os.path.join(self.test_scripts_folder   ,     """test_checkpointing.py"""   )
         lowerCAmelCase__      :			List[Any]   =	[
             '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__      :			int   =	cmd.copy()
               cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""   )
               if strategy != "FULL_SHARD":
                     continue
               lowerCAmelCase__      :			List[str]   =	len(SCREAMING_SNAKE_CASE_   )
               for state_dict_type in FSDP_STATE_DICT_TYPE:
                     lowerCAmelCase__      :			Tuple   =	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__      :			Optional[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(SCREAMING_SNAKE_CASE_   ,     env=os.environ.copy()   )






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

         """simple docstring"""
         lowerCAmelCase__      :			Optional[int]   =	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(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"""


import json
import os
from typing import Optional, Tuple

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


_A													=      logging.get_logger(__name__)

_A													=      {"""vocab_file""": """vocab.json"""}

_A													=      {
    """vocab_file""": {
        """mgp-str""": """https://huggingface.co/alibaba-damo/mgp-str-base/blob/main/vocab.json""",
    }
}

_A													=      {"""mgp-str""": 2_7}


class _lowerCamelCase    (					a_						):
   _lowerCamelCase    :Dict     				=     VOCAB_FILES_NAMES
   _lowerCamelCase    :Union[str, Any]     				=     PRETRAINED_VOCAB_FILES_MAP
   _lowerCamelCase    :Optional[int]     				=     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES






   def __init__(		self       :		Optional[int]   ,     UpperCamelCase       :		int   ,     UpperCamelCase       :		Union[str, Any]="[GO]"   ,     UpperCamelCase       :		Any="[GO]"   ,     UpperCamelCase       :		Tuple="[s]"   ,     UpperCamelCase       :		List[Any]="[GO]"   ,     **UpperCamelCase       :		Dict   )			->						Union[str, Any]:

         """simple docstring"""
         super().__init__(
             unk_token=UpperCamelCase   ,     bos_token=UpperCamelCase   ,     eos_token=UpperCamelCase   ,     pad_token=UpperCamelCase   ,     **UpperCamelCase   ,     )

         with open(UpperCamelCase   ,     encoding="""utf-8"""   ) as vocab_handle:
               lowerCAmelCase__      :			Optional[Any]   =	json.load(UpperCamelCase   )
         lowerCAmelCase__      :			Union[str, Any]   =	{v: k for k, v in self.vocab.items()}






   @property
   def 	_lowerCAmelCase   (		self       :		int   )			->						Dict:

         """simple docstring"""
         return len(self.vocab   )






   def 	_lowerCAmelCase   (		self       :		Dict   )			->						Any:

         """simple docstring"""
         return dict(self.vocab   ,     **self.added_tokens_encoder   )






   def 	_lowerCAmelCase   (		self       :		List[Any]   ,     UpperCamelCase       :		Tuple   )			->						Optional[Any]:

         """simple docstring"""
         lowerCAmelCase__      :			Optional[Any]   =	[]
         for s in text:
               char_tokens.extend(UpperCamelCase   )
         return char_tokens






   def 	_lowerCAmelCase   (		self       :		List[Any]   ,     UpperCamelCase       :		str   )			->						Tuple:

         """simple docstring"""
         return self.vocab.get(UpperCamelCase   ,     self.vocab.get(self.unk_token   )   )






   def 	_lowerCAmelCase   (		self       :		List[str]   ,     UpperCamelCase       :		int   )			->						Any:

         """simple docstring"""
         return self.decoder.get(UpperCamelCase   )






   def 	_lowerCAmelCase   (		self       :		Dict   ,     UpperCamelCase       :		str   ,     UpperCamelCase       :		Optional[str] = None   )			->						Tuple[str]:

         """simple docstring"""
         if not os.path.isdir(UpperCamelCase   ):
               logger.error("""Vocabulary path ({}) should be a directory""".format(UpperCamelCase   )   )
               return
         lowerCAmelCase__      :			Union[str, Any]   =	os.path.join(
             UpperCamelCase   ,     (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]   )

         with open(UpperCamelCase   ,     """w"""   ,     encoding="""utf-8"""   ) as f:
               f.write(json.dumps(self.vocab   ,     indent=2   ,     sort_keys=UpperCamelCase   ,     ensure_ascii=UpperCamelCase   ) + """\n"""   )

         return (vocab_file,)





