




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

a  :				Optional[Any]									=  'bert-base-cased'
a  :				List[str]									=  'fp16'
a  :				Any									=  'bf16'
a  :				List[str]									=  [FPaa, BFaa]






@require_fsdp
@require_cuda
class 						a				(   __a     ):


				def   A_ (		self    :					Dict					):
									super().setUp()

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


				def   A_ (		self    :					List[Any]					):
									from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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


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

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


				def   A_ (		self    :					Any					):
									from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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

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

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


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

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







@require_fsdp
@require_multi_gpu
@slow
class 						a				(   __a     ):


				def   A_ (		self    :					Optional[Any]					):
									super().setUp()
									snake_case_									=				0.82
									snake_case_									=				[
									    '''fsdp_shard_grad_op_transformer_based_wrap''',
									    '''fsdp_full_shard_transformer_based_wrap''',
									]
									snake_case_									=				{
									    '''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
									}
									snake_case_									=				160
									snake_case_									=				160

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


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

																			snake_case_									=				cmd_config[:-1]
																			snake_case_									=				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   A_ (		self    :					List[str]					):
									snake_case_									=				os.path.join(self.test_scripts_folder       ,				'''test_peak_memory_usage.py'''					)
									snake_case_									=				[
									    '''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():
														snake_case_									=				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 os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

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


UpperCamelCase__    					=      logging.get_logger(__name__)

UpperCamelCase__    					=      '''▁'''

UpperCamelCase__    					=      {'''vocab_file''': '''sentencepiece.bpe.model'''}

UpperCamelCase__    					=      {
    '''vocab_file''': {
        '''facebook/xglm-564M''': '''https://huggingface.co/facebook/xglm-564M/resolve/main/sentencepiece.bpe.model''',
    }
}

UpperCamelCase__    					=      {
    '''facebook/xglm-564M''': 2_0_4_8,
}




class 			lowerCamelCase_   (	__a      ):
      lowerCAmelCase__		     =     VOCAB_FILES_NAMES
      lowerCAmelCase__		     =     PRETRAINED_VOCAB_FILES_MAP
      lowerCAmelCase__		     =     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      lowerCAmelCase__		     =     ['input_ids', 'attention_mask']






      def __init__(					self				:    Optional[int]						,						_A				:    Optional[Any]						,						_A				:    Optional[Any]="<s>"						,						_A				:    List[str]="</s>"						,						_A				:    Optional[Any]="</s>"						,						_A				:    List[str]="<s>"						,						_A				:    Tuple="<unk>"						,						_A				:    List[str]="<pad>"						,						_A				:    Optional[Dict[str, Any]] = None						,						**_A				:    Union[str, Any]						,						):


          '''simple docstring'''


          UpperCAmelCase__			:		str													=				{} if sp_model_kwargs is None else sp_model_kwargs

          # Compatibility with the original tokenizer
          UpperCAmelCase__			:		Optional[int]													=				7
          UpperCAmelCase__			:		Optional[int]													=				[f"""<madeupword{i}>""" for i in range(self.num_madeup_words )]

          UpperCAmelCase__			:		Tuple													=				kwargs.get('''additional_special_tokens'''						,						[] )
          kwargs["additional_special_tokens"] += [
              word for word in madeup_words if word not in kwargs["additional_special_tokens"]
          ]

          super().__init__(
              bos_token=_A						,						eos_token=_A						,						unk_token=_A						,						sep_token=_A						,						cls_token=_A						,						pad_token=_A						,						sp_model_kwargs=self.sp_model_kwargs						,						**_A						,						)

          UpperCAmelCase__			:		Any													=				spm.SentencePieceProcessor(**self.sp_model_kwargs )
          self.sp_model.Load(str(_A ) )
          UpperCAmelCase__			:		List[str]													=				vocab_file

          # Original fairseq vocab and spm vocab must be "aligned":
          # Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
          # -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
          # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
          # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

          # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
          UpperCAmelCase__			:		Any													=				1

          # Mimic fairseq token-to-id alignment for the first 4 token
          UpperCAmelCase__			:		Any													=				{'''<s>''': 0, '''<pad>''': 1, '''</s>''': 2, '''<unk>''': 3}

          UpperCAmelCase__			:		int													=				len(self.sp_model )
          UpperCAmelCase__			:		Optional[int]													=				{f"""<madeupword{i}>""": sp_size + i + self.fairseq_offset for i in range(self.num_madeup_words )}
          self.fairseq_tokens_to_ids.update(_A )

          UpperCAmelCase__			:		Tuple													=				{v: k for k, v in self.fairseq_tokens_to_ids.items()}






      def __getstate__(					self				:    Any ):


          '''simple docstring'''


          UpperCAmelCase__			:		Dict													=				self.__dict__.copy()
          UpperCAmelCase__			:		Optional[int]													=				None
          UpperCAmelCase__			:		Optional[Any]													=				self.sp_model.serialized_model_proto()
          return state






      def __setstate__(					self				:    Union[str, Any]						,						_A				:    int ):


          '''simple docstring'''


          UpperCAmelCase__			:		List[Any]													=				d

          # for backward compatibility
          if not hasattr(self						,						'''sp_model_kwargs''' ):
              UpperCAmelCase__			:		Dict													=				{}

          UpperCAmelCase__			:		str													=				spm.SentencePieceProcessor(**self.sp_model_kwargs )
          self.sp_model.LoadFromSerializedProto(self.sp_model_proto )






      def lowercase_							(					self				:    Any						,						_A				:    List[int]						,						_A				:    Optional[List[int]] = None ):


          '''simple docstring'''


          if token_ids_a is None:
              return [self.sep_token_id] + token_ids_a
          UpperCAmelCase__			:		Dict													=				[self.sep_token_id]
          return sep + token_ids_a + sep + sep + token_ids_a






      def lowercase_							(					self				:    str						,						_A				:    List[int]						,						_A				:    Optional[List[int]] = None						,						_A				:    bool = False ):


          '''simple docstring'''


          if already_has_special_tokens:
              return super().get_special_tokens_mask(
                  token_ids_a=_A						,						token_ids_a=_A						,						already_has_special_tokens=_A )

          if token_ids_a is None:
              return [1] + ([0] * len(_A ))
          return [1] + ([0] * len(_A )) + [1, 1] + ([0] * len(_A ))






      def lowercase_							(					self				:    Any						,						_A				:    List[int]						,						_A				:    Optional[List[int]] = None ):


          '''simple docstring'''


          UpperCAmelCase__			:		Optional[int]													=				[self.sep_token_id]

          if token_ids_a is None:
              return len(sep + token_ids_a ) * [0]
          return len(sep + token_ids_a + sep + sep + token_ids_a ) * [0]






      @property
      def lowercase_							(					self				:    Any ):


          '''simple docstring'''


          return len(self.sp_model ) + self.fairseq_offset + self.num_madeup_words






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


          '''simple docstring'''


          UpperCAmelCase__			:		str													=				{self.convert_ids_to_tokens(_A ): i for i in range(self.vocab_size )}
          vocab.update(self.added_tokens_encoder )
          return vocab






      def lowercase_							(					self				:    Optional[Any]						,						_A				:    str ):


          '''simple docstring'''


          return self.sp_model.encode(_A						,						out_type=_A )






      def lowercase_							(					self				:    List[str]						,						_A				:    List[Any] ):


          '''simple docstring'''


          if token in self.fairseq_tokens_to_ids:
              return self.fairseq_tokens_to_ids[token]
          UpperCAmelCase__			:		Union[str, Any]													=				self.sp_model.PieceToId(_A )

          # Need to return unknown token if the SP model returned 0
          return spm_id + self.fairseq_offset if spm_id else self.unk_token_id






      def lowercase_							(					self				:    List[Any]						,						_A				:    str ):


          '''simple docstring'''


          if index in self.fairseq_ids_to_tokens:
              return self.fairseq_ids_to_tokens[index]
          return self.sp_model.IdToPiece(index - self.fairseq_offset )






      def lowercase_							(					self				:    int						,						_A				:    Optional[int] ):


          '''simple docstring'''


          UpperCAmelCase__			:		List[str]													=				''''''.join(_A ).replace(_A						,						''' ''' ).strip()
          return out_string






      def lowercase_							(					self				:    Any						,						_A				:    str						,						_A				:    Optional[str] = None ):


          '''simple docstring'''


          if not os.path.isdir(_A ):
              logger.error(f"""Vocabulary path ({save_directory}) should be a directory""" )
              return
          UpperCAmelCase__			:		List[str]													=				os.path.join(
              _A						,						(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file'''] )

          if os.path.abspath(self.vocab_file ) != os.path.abspath(_A ) and os.path.isfile(self.vocab_file ):
              copyfile(self.vocab_file						,						_A )
          elif not os.path.isfile(self.vocab_file ):
              with open(_A						,						'''wb''' ) as fi:
                  UpperCAmelCase__			:		List[Any]													=				self.sp_model.serialized_model_proto()
                  fi.write(_A )

          return (out_vocab_file,)
