





'''simple docstring'''






#
# This a `torch.distributed` diagnostics script that checks that all GPUs in the cluster (one or
# many nodes) can talk to each other via nccl and allocate gpu memory.
#
# To run first adjust the number of processes and nodes:
#
# python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# You may need to add --master_addr $MASTER_ADDR --master_port $MASTER_PORT if using a custom addr:port
#
# You can also use the rdzv API: --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT --rdzv_backend c10d
#
# use torch.distributed.launch instead of torch.distributed.run for torch < 1.9
#
# If you get a hanging in `barrier` calls you have some network issues, you may try to debug this with:
#
# NCCL_DEBUG=INFO python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# which should tell you what's going on behind the scenes.
#
#
# This script can be run via `srun` in the SLURM environment as well. Here is a SLURM script that
# runs on 2 nodes of 4 gpus per node:
#
# #SBATCH --job-name=test-nodes        # name
# #SBATCH --nodes=2                    # nodes
# #SBATCH --ntasks-per-node=1          # crucial - only 1 task per dist per node!
# #SBATCH --cpus-per-task=10           # number of cores per tasks
# #SBATCH --gres=gpu:4                 # number of gpus
# #SBATCH --time 0:05:00               # maximum execution time (HH:MM:SS)
# #SBATCH --output=%x-%j.out           # output file name
#
# GPUS_PER_NODE=4
# MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
# MASTER_PORT=6000
#
# srun --jobid $SLURM_JOBID bash -c 'python -m torch.distributed.run \
# --nproc_per_node $GPUS_PER_NODE --nnodes $SLURM_NNODES --node_rank $SLURM_PROCID \
# --master_addr $MASTER_ADDR --master_port $MASTER_PORT \
# torch-distributed-gpu-test.py'
#

import fcntl
import os
import socket

import torch
import torch.distributed as dist


def 					_a   (*__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  with open(__UpperCamelCase  ,      '''r'''   ) as fh:
    fcntl.flock(__UpperCamelCase  ,      fcntl.LOCK_EX   )
    try:
      print(*__UpperCamelCase   )
    finally:
      fcntl.flock(__UpperCamelCase  ,      fcntl.LOCK_UN   )


__lowerCamelCase						: Dict		=   int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
__lowerCamelCase						: Any		=   torch.device('cuda', local_rank)
__lowerCamelCase						: int		=   socket.gethostname()

__lowerCamelCase						: List[str]		=   F"""[{hostname}-{local_rank}]"""

try:
      # test distributed
      dist.init_process_group('nccl')
      dist.all_reduce(torch.ones(1).to(device), op=dist.ReduceOp.SUM)
      dist.barrier()

      # test cuda is available and can allocate memory
      torch.cuda.is_available()
      torch.ones(1).cuda(local_rank)

      # global rank
      __lowerCamelCase						: Optional[Any]		=   dist.get_rank()
      __lowerCamelCase						: Union[str, Any]		=   dist.get_world_size()

      printflock(F"""{gpu} is OK (global rank: {rank}/{world_size})""")

      dist.barrier()
      if rank == 0:
            printflock(F"""pt={torch.__version__}, cuda={torch.version.cuda}, nccl={torch.cuda.nccl.version()}""")

except Exception:
      printflock(F"""{gpu} is broken""")
      raise








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






