


"""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    _snake_case  (  *_snake_case :	Dict	):
      with open(_snake_case       ,      '''r'''	) as fh:
            fcntl.flock(_snake_case       ,      fcntl.LOCK_EX	)
            try:
                  print(*_snake_case	)
            finally:
                  fcntl.flock(_snake_case       ,      fcntl.LOCK_UN	)


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

snake_case__						:   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
 snake_case__						:   List[Any]			 = dist.get_rank()
 snake_case__						:   List[str]			 = 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





import math
import os
import re
import sys
import unittest
from pathlib import Path
from typing import Tuple
from unittest.mock import patch

from parameterized import parameterized

from transformers.testing_utils import (
    CaptureStderr,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_torch_dist_unique_port,
    require_apex,
    require_bitsandbytes,
    require_fairscale,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    require_torch_non_multi_gpu,
    slow,
)
from transformers.trainer_callback import TrainerState
from transformers.trainer_utils import set_seed


lowercase_									=   os.path.abspath(os.path.dirname(__file__))
with ExtendSysPath(f"{bindir}/../../examples/pytorch/translation"):
  from run_translation import main  # noqa


set_seed(4_2)
lowercase_									=   'sshleifer/student_marian_en_ro_6_1'
lowercase_									=   'sshleifer/tiny-mbart'





@require_torch
class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE							):




  def   A__     (				self							,   lowerCAmelCase=False							,   lowerCAmelCase=None							,   lowerCAmelCase=True							,   lowerCAmelCase=True							,   lowerCAmelCase=True							,   lowerCAmelCase=True							,   )  -> Dict:
        '''simple docstring'''






        _lowercase	       =self.run_trainer(
            eval_steps=1							,   max_len=12							,   model_name=lowerCAmelCase							,   num_train_epochs=1							,   distributed=lowerCAmelCase							,   extra_args_str=lowerCAmelCase							,   predict_with_generate=lowerCAmelCase							,   do_train=lowerCAmelCase							,   do_eval=lowerCAmelCase							,   do_predict=lowerCAmelCase							,   )
        _lowercase	       =TrainerState.load_from_json(os.path.join(lowerCAmelCase							,   'trainer_state.json'					)					).log_history

        if not do_eval:
              return

        _lowercase	       =[log for log in logs if 'eval_loss' in log.keys()]

        _lowercase	       =eval_metrics[0]
        if predict_with_generate:
              assert "eval_bleu" in first_step_stats

              _lowercase	       =eval_metrics[-1]
              assert isinstance(last_step_stats['eval_bleu']							,   lowerCAmelCase					)
              assert not math.isnan(float(last_step_stats['eval_loss']					)					), "eval_loss must not be `nan`"




  @require_torch_non_multi_gpu
  def   A__     (				self					)  -> Optional[Any]:
        '''simple docstring'''






        self.run_seqaseq_quick()




  @require_torch_multi_gpu
  def   A__     (				self					)  -> Tuple:
        '''simple docstring'''






        self.run_seqaseq_quick(distributed=lowerCAmelCase					)




  @require_torch_multi_gpu
  def   A__     (				self					)  -> Optional[Any]:
        '''simple docstring'''






        self.run_seqaseq_quick(distributed=lowerCAmelCase					)




  @unittest.skip('Requires an update of the env running those tests'					)
  @require_torch_multi_gpu
  @require_fairscale
  def   A__     (				self					)  -> Optional[Any]:
        '''simple docstring'''






        self.run_seqaseq_quick(distributed=lowerCAmelCase							,   extra_args_str='--sharded_ddp simple'					)




  @unittest.skip('Requires an update of the env running those tests'					)
  @require_torch_multi_gpu
  @require_fairscale
  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        self.run_seqaseq_quick(distributed=lowerCAmelCase							,   extra_args_str='--sharded_ddp simple --fp16'					)




  @unittest.skip('Requires an update of the env running those tests'					)
  @require_torch_multi_gpu
  @require_fairscale
  def   A__     (				self					)  -> Dict:
        '''simple docstring'''






        self.run_seqaseq_quick(distributed=lowerCAmelCase							,   extra_args_str='--sharded_ddp zero_dp_2'							,   predict_with_generate=lowerCAmelCase					)




  @unittest.skip('Requires an update of the env running those tests'					)
  @require_torch_multi_gpu
  @require_fairscale
  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        self.run_seqaseq_quick(
            distributed=lowerCAmelCase							,   extra_args_str='--sharded_ddp zero_dp_2 --fp16'							,   predict_with_generate=lowerCAmelCase					)




  @require_apex
  @require_torch_gpu
  def   A__     (				self					)  -> List[Any]:
        '''simple docstring'''






        self.run_seqaseq_quick(distributed=lowerCAmelCase							,   extra_args_str='--fp16 --fp16_backend=apex'					)
        # test 2nd time - was getting eval_loss': nan'
        # to reproduce the problem set distributed=False
        self.run_seqaseq_quick(distributed=lowerCAmelCase							,   extra_args_str='--fp16 --fp16_backend=apex'					)




  @parameterized.expand(['base', 'low', 'high', 'mixed']					)
  @require_torch_multi_gpu
  def   A__     (				self							,   lowerCAmelCase					)  -> Any:
        '''simple docstring'''






        _lowercase	       ={
            # test with the default log_level - should be info and thus log info once
            'base': {'extra_args_str': '', 'n_matches': 1},
            # test with low log_level and log_level_replica - should be noisy on all processes
            # now the info string should appear twice on 2 processes
            'low': {'extra_args_str': '--log_level debug --log_level_replica debug', 'n_matches': 2},
            # test with high log_level and low log_level_replica
            # now the info string should appear once only on the replica
            'high': {'extra_args_str': '--log_level error --log_level_replica debug', 'n_matches': 1},
            # test with high log_level and log_level_replica - should be quiet on all processes
            'mixed': {'extra_args_str': '--log_level error --log_level_replica error', 'n_matches': 0},
        }

        _lowercase	       =experiments[experiment_id]
        _lowercase	       ={'distributed': True, 'predict_with_generate': False, 'do_eval': False, 'do_predict': False}
        _lowercase	       ='Running training'
        with CaptureStderr() as cl:
              self.run_seqaseq_quick(**lowerCAmelCase							,   extra_args_str=data['extra_args_str']					)
        _lowercase	       =len(re.findall(lowerCAmelCase							,   cl.err					)					)
        self.assertEqual(lowerCAmelCase							,   data['n_matches']					)




  @slow
  def   A__     (				self					)  -> Union[str, Any]:
        '''simple docstring'''






        _lowercase	       =self.run_trainer(
            eval_steps=2							,   max_len=128							,   model_name=lowerCAmelCase							,   learning_rate=3e-4							,   num_train_epochs=10							,   distributed=lowerCAmelCase							,   )

        # Check metrics
        _lowercase	       =TrainerState.load_from_json(os.path.join(lowerCAmelCase							,   'trainer_state.json'					)					).log_history
        _lowercase	       =[log for log in logs if 'eval_loss' in log.keys()]
        _lowercase	       =eval_metrics[0]
        _lowercase	       =eval_metrics[-1]

        assert first_step_stats["eval_loss"] > last_step_stats["eval_loss"], "model learned nothing"
        assert isinstance(last_step_stats['eval_bleu']							,   lowerCAmelCase					)

        # test if do_predict saves generations and metrics
        _lowercase	       =os.listdir(lowerCAmelCase					)
        _lowercase	       ={os.path.basename(lowerCAmelCase					) for p in contents}
        assert "generated_predictions.txt" in contents
        assert "predict_results.json" in contents




  @slow
  @require_bitsandbytes
  def   A__     (				self					)  -> List[str]:
        '''simple docstring'''






        from transformers.training_args import OptimizerNames

        def train_and_return_metrics(lowerCAmelCase					) -> Tuple[int, float]:
              _lowercase	       ='--skip_memory_metrics 0'

              _lowercase	       =self.run_trainer(
                  max_len=128							,   model_name=lowerCAmelCase							,   learning_rate=3e-4							,   num_train_epochs=1							,   optim=lowerCAmelCase							,   distributed=lowerCAmelCase							,   extra_args_str=lowerCAmelCase							,   do_eval=lowerCAmelCase							,   do_predict=lowerCAmelCase							,   n_gpus_to_use=1							,   )

              # Check metrics
              _lowercase	       =TrainerState.load_from_json(Path(lowerCAmelCase							,   'trainer_state.json'					)					).log_history
              _lowercase	       =int(logs[0]['train_mem_gpu_peaked_delta'] / 2**20					)
              _lowercase	       =int(logs[0]['train_mem_gpu_alloc_delta'] / 2**20					)

              _lowercase	       =logs[0]['train_loss']
              return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

        _lowercase       ,       _lowercase       ,       _lowercase	       =train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value					)
        _lowercase       ,       _lowercase       ,       _lowercase	       =train_and_return_metrics(OptimizerNames.ADAMW_BNB.value					)

        _lowercase	       =gpu_alloc_mem_orig - gpu_alloc_mem_bnb

        _lowercase	       =gpu_peak_mem_orig + gpu_alloc_mem_orig
        _lowercase	       =gpu_peak_mem_bnb + gpu_alloc_mem_bnb
        _lowercase	       =gpu_total_mem_orig - gpu_total_mem_bnb

        # sshleifer/student_marian_en_ro_6_1 has 54M parameter, 29M of which is `nn.Embedding` which
        # doesn't get quantized and remains in fp32. Therefore we only have 25M parameters quantized
        # in 2 bytes and the diff in optim memory usage is derived as so:
        #
        # - normal 25*8=~200MB (8 bytes per param)
        # - bnb    25*2= ~50MB (2 bytes per param)
        #
        # Thus we should expect ~150MB total memory saved.
        #
        # Peak memory should be the same - the total should be different by about that same margin
        #
        # After leaving a small margin to accommodate for differences between gpus let's check
        # that we have at least 120MB in savings
        _lowercase	       =120

        # uncomment the following if this test starts failing - requires py38 for a new print feature
        # gpu_peak_mem_diff = gpu_peak_mem_orig - gpu_peak_mem_bnb
        # print(f"{gpu_alloc_mem_orig=}MB {gpu_peak_mem_orig=}MB {gpu_alloc_mem_orig+gpu_peak_mem_orig=}MB")
        # print(f" {gpu_alloc_mem_bnb=}MB  {gpu_peak_mem_bnb=}MB  {gpu_alloc_mem_bnb+gpu_peak_mem_bnb=}MB")
        # print(f"{gpu_alloc_mem_diff=}MB")
        # print(f"{gpu_peak_mem_diff=}MB")
        # print(f"{gpu_total_mem_orig=}MB, {gpu_total_mem_bnb=}MB")
        # print(f"{gpu_total_mem_diff=}MB, {gpu_total_mem_diff=}MB")

        self.assertGreater(
            lowerCAmelCase							,   lowerCAmelCase							,   'should use ~150MB less alloc gpu memory with BNB, compared to without it for this model but got'
            F''' a difference of {gpu_alloc_mem_diff}MB, with gpu_alloc_mem_orig={gpu_alloc_mem_orig}MB and'''
            F''' gpu_alloc_mem_bnb={gpu_alloc_mem_bnb}MB'''							,   )

        self.assertGreater(
            lowerCAmelCase							,   lowerCAmelCase							,   'should use ~150MB less total gpu memory with BNB, compared to without it for this model but got'
            F''' a difference of {gpu_total_mem_diff}MB, with gpu_total_mem_orig={gpu_total_mem_orig}MB and'''
            F''' gpu_total_mem_bnb={gpu_total_mem_bnb}MB'''							,   )

        self.assertEqual(
            lowerCAmelCase							,   lowerCAmelCase							,   F'''loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}'''					)




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase = 3e-3							,   lowerCAmelCase = "adafactor"							,   lowerCAmelCase = False							,   lowerCAmelCase = None							,   lowerCAmelCase = 0							,   lowerCAmelCase = True							,   lowerCAmelCase = True							,   lowerCAmelCase = True							,   lowerCAmelCase = True							,   lowerCAmelCase = None							,   )  -> Optional[Any]:
        '''simple docstring'''






        _lowercase	       =self.test_file_dir / '../fixtures/tests_samples/wmt_en_ro'
        _lowercase	       =self.get_auto_remove_tmp_dir()
        _lowercase	       =F'''
            --model_name_or_path {model_name}
            --train_file {data_dir}/train.json
            --validation_file {data_dir}/val.json
            --test_file {data_dir}/test.json
            --output_dir {output_dir}
            --overwrite_output_dir
            --max_train_samples 8
            --max_source_length {max_len}
            --max_target_length {max_len}
            --do_train
            --num_train_epochs {str(lowerCAmelCase					)}
            --per_device_train_batch_size 4
            --learning_rate {learning_rate}
            --warmup_steps 8
            --logging_steps 0
            --logging_strategy no
            --save_steps {str(lowerCAmelCase					)}
            --group_by_length
            --label_smoothing_factor 0.1
            --target_lang ro_RO
            --source_lang en_XX
        '''.split()

        _lowercase	       =F'''
            --do_eval
            --per_device_eval_batch_size 4
            --max_eval_samples 8
            --val_max_target_length {max_len}
            --evaluation_strategy steps
            --eval_steps {str(lowerCAmelCase					)}
        '''.split()

        _lowercase	       ='\n            --do_predict\n        '.split()

        _lowercase	       =[]
        if do_train:
              args += args_train

        if do_eval:
              args += args_eval

        if do_predict:
              args += args_predict

        if predict_with_generate:
              args += "--predict_with_generate".split()

        if do_train:
              if optim == "adafactor":
                    args += "--adafactor".split()
              else:
                    args += F'''--optim {optim}'''.split()

        if extra_args_str is not None:
              args += extra_args_str.split()

        if distributed:
              if n_gpus_to_use is None:
                    _lowercase	       =get_gpu_count()
              _lowercase	       =get_torch_dist_unique_port()
              _lowercase	       =F'''
                -m torch.distributed.run
                --nproc_per_node={n_gpus_to_use}
                --master_port={master_port}
                {self.examples_dir_str}/pytorch/translation/run_translation.py
            '''.split()
              _lowercase	       =[sys.executable] + distributed_args + args
              # keep for quick debug
              # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
              execute_subprocess_async(lowerCAmelCase							,   env=self.get_env()					)
        else:
              _lowercase	       =['run_translation.py'] + args
              with patch.object(lowerCAmelCase							,   'argv'							,   lowerCAmelCase					):
                    main()

        return output_dir


