





'''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						: List[Any]		=   'bert-base-cased'
__lowerCamelCase						: str		=   'fp16'
__lowerCamelCase						: Optional[int]		=   'bf16'
__lowerCamelCase						: List[Any]		=   [FPaa, BFaa]







@require_fsdp
@require_cuda
class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







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

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


  def   UpperCamelCase__      (				self    :					int  )     ->						List[Any]:
    from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Any:
    from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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


  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Tuple:
    from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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

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

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


  def   UpperCamelCase__      (				self    :					Any  )     ->						List[Any]:
    from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
    from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

    for mp_dtype in dtypes:
      _UpperCamelCase       			=self.dist_env.copy()
      _UpperCamelCase       			=mp_dtype
      with mockenv_context(**UpperCamelCase__  ):
        _UpperCamelCase       			=Accelerator()
        if mp_dtype == "fp16":
          _UpperCamelCase       			=torch.floataa
        elif mp_dtype == "bf16":
          _UpperCamelCase       			=torch.bfloataa
        _UpperCamelCase       			=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   UpperCamelCase__      (				self    :					Dict  )     ->						Tuple:
    from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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







@require_fsdp
@require_multi_gpu
@slow
class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						List[str]:
    super().setUp()
    _UpperCamelCase       			=0.82
    _UpperCamelCase       			=[
        '''fsdp_shard_grad_op_transformer_based_wrap''',
        '''fsdp_full_shard_transformer_based_wrap''',
    ]
    _UpperCamelCase       			={
        '''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
    }
    _UpperCamelCase       			=160
    _UpperCamelCase       			=160

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


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

        _UpperCamelCase       			=cmd_config[:-1]
        _UpperCamelCase       			=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   UpperCamelCase__      (				self    :					Dict  )     ->						int:
    _UpperCamelCase       			=os.path.join(self.test_scripts_folder							,	'''test_peak_memory_usage.py'''  )
    _UpperCamelCase       			=[
        '''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():
      _UpperCamelCase       			=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 argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
__lowerCamelCase						: List[Any]		=   data_utils.TransfoXLTokenizer
__lowerCamelCase						: str		=   data_utils.TransfoXLCorpus
__lowerCamelCase						: Optional[int]		=   data_utils
__lowerCamelCase						: int		=   data_utils


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  if transfo_xl_dataset_file:
    # Convert a pre-processed corpus (see original TensorFlow repo)
    with open(__SCREAMING_SNAKE_CASE  ,      '''rb'''   ) as fp:
      _UpperCamelCase       			=pickle.load(__SCREAMING_SNAKE_CASE  ,      encoding='''latin1'''   )
    # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
    _UpperCamelCase       			=pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''pretrained_vocab_file''']
    print(f'''Save vocabulary to {pytorch_vocab_dump_path}'''   )
    _UpperCamelCase       			=corpus.vocab.__dict__
    torch.save(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )

    _UpperCamelCase       			=corpus.__dict__
    corpus_dict_no_vocab.pop('''vocab'''  ,      __SCREAMING_SNAKE_CASE   )
    _UpperCamelCase       			=pytorch_dump_folder_path + '''/''' + CORPUS_NAME
    print(f'''Save dataset to {pytorch_dataset_dump_path}'''   )
    torch.save(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )

  if tf_checkpoint_path:
    # Convert a pre-trained TensorFlow model
    _UpperCamelCase       			=os.path.abspath(__SCREAMING_SNAKE_CASE   )
    _UpperCamelCase       			=os.path.abspath(__SCREAMING_SNAKE_CASE   )

    print(f'''Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.'''   )
    # Initialise PyTorch model
    if transfo_xl_config_file == "":
      _UpperCamelCase       			=TransfoXLConfig()
    else:
      _UpperCamelCase       			=TransfoXLConfig.from_json_file(__SCREAMING_SNAKE_CASE   )
    print(f'''Building PyTorch model from configuration: {config}'''   )
    _UpperCamelCase       			=TransfoXLLMHeadModel(__SCREAMING_SNAKE_CASE   )

    _UpperCamelCase       			=load_tf_weights_in_transfo_xl(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )
    # Save pytorch-model
    _UpperCamelCase       			=os.path.join(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )
    _UpperCamelCase       			=os.path.join(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )
    print(f'''Save PyTorch model to {os.path.abspath(__SCREAMING_SNAKE_CASE   )}'''   )
    torch.save(model.state_dict()  ,      __SCREAMING_SNAKE_CASE   )
    print(f'''Save configuration file to {os.path.abspath(__SCREAMING_SNAKE_CASE   )}'''   )
    with open(__SCREAMING_SNAKE_CASE  ,      '''w'''  ,      encoding='''utf-8'''   ) as f:
      f.write(config.to_json_string()   )


if __name__ == "__main__":
      __lowerCamelCase						: Any		=   argparse.ArgumentParser()
      parser.add_argument(
          '--pytorch_dump_folder_path',
          default=None,
          type=str,
          required=True,
          help='Path to the folder to store the PyTorch model or dataset/vocab.',
      )
      parser.add_argument(
          '--tf_checkpoint_path',
          default='',
          type=str,
          help='An optional path to a TensorFlow checkpoint path to be converted.',
      )
      parser.add_argument(
          '--transfo_xl_config_file',
          default='',
          type=str,
          help=(
              'An optional config json file corresponding to the pre-trained BERT model. \n'
              'This specifies the model architecture.'
          ),
      )
      parser.add_argument(
          '--transfo_xl_dataset_file',
          default='',
          type=str,
          help='An optional dataset file to be converted in a vocabulary.',
      )
      __lowerCamelCase						: Optional[int]		=   parser.parse_args()
      convert_transfo_xl_checkpoint_to_pytorch(
          args.tf_checkpoint_path,
          args.transfo_xl_config_file,
          args.pytorch_dump_folder_path,
          args.transfo_xl_dataset_file,
      )






