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


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


set_seed(4_2)
_SCREAMING_SNAKE_CASE      						=	"""sshleifer/student_marian_en_ro_6_1"""
_SCREAMING_SNAKE_CASE      						=	"""sshleifer/tiny-mbart"""
@require_torch
class SCREAMING_SNAKE_CASE_						(	lowercase_						):







   def 		lowerCamelCase_     (			self		:		int					,					lowerCamelCase_		:		str=False					,					lowerCamelCase_		:		List[str]=None					,					lowerCamelCase_		:		str=True					,					lowerCamelCase_		:		Tuple=True					,					lowerCamelCase_		:		List[Any]=True					,					lowerCamelCase_		:		Any=True					,					):





    """simple docstring"""


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

    if not do_eval:
     return

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

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

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







   @require_torch_non_multi_gpu
   def 		lowerCamelCase_     (			self		:		Union[str, Any]							):





    """simple docstring"""


    self.run_seqaseq_quick()







   @require_torch_multi_gpu
   def 		lowerCamelCase_     (			self		:		List[str]							):





    """simple docstring"""


    self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE_							)







   @require_torch_multi_gpu
   def 		lowerCamelCase_     (			self		:		int							):





    """simple docstring"""


    self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE_							)







   @unittest.skip("""Requires an update of the env running those tests"""							)
   @require_torch_multi_gpu
   @require_fairscale
   def 		lowerCamelCase_     (			self		:		str							):





    """simple docstring"""


    self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE_					,					extra_args_str="""--sharded_ddp simple"""							)







   @unittest.skip("""Requires an update of the env running those tests"""							)
   @require_torch_multi_gpu
   @require_fairscale
   def 		lowerCamelCase_     (			self		:		List[Any]							):





    """simple docstring"""


    self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE_					,					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 		lowerCamelCase_     (			self		:		Dict							):





    """simple docstring"""


    self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE_					,					extra_args_str="""--sharded_ddp zero_dp_2"""					,					predict_with_generate=SCREAMING_SNAKE_CASE_							)







   @unittest.skip("""Requires an update of the env running those tests"""							)
   @require_torch_multi_gpu
   @require_fairscale
   def 		lowerCamelCase_     (			self		:		str							):





    """simple docstring"""


    self.run_seqaseq_quick(
        distributed=SCREAMING_SNAKE_CASE_					,					extra_args_str="""--sharded_ddp zero_dp_2 --fp16"""					,					predict_with_generate=SCREAMING_SNAKE_CASE_							)







   @require_apex
   @require_torch_gpu
   def 		lowerCamelCase_     (			self		:		Dict							):





    """simple docstring"""


    self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE_					,					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=SCREAMING_SNAKE_CASE_					,					extra_args_str="""--fp16 --fp16_backend=apex"""							)







   @parameterized.expand(["""base""", """low""", """high""", """mixed"""]							)
   @require_torch_multi_gpu
   def 		lowerCamelCase_     (			self		:		Dict					,					lowerCamelCase_		:		Optional[Any]							):





    """simple docstring"""


    UpperCamelCase          =	{
        # 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},
    }

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







   @slow
   def 		lowerCamelCase_     (			self		:		List[Any]							):





    """simple docstring"""


    UpperCamelCase          =	self.run_trainer(
        eval_steps=2					,					max_len=128					,					model_name=SCREAMING_SNAKE_CASE_					,					learning_rate=3E-4					,					num_train_epochs=10					,					distributed=SCREAMING_SNAKE_CASE_					,					)

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

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

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







   @slow
   @require_bitsandbytes
   def 		lowerCamelCase_     (			self		:		int							):





    """simple docstring"""


    from transformers.training_args import OptimizerNames

    def train_and_return_metrics(lowerCamelCase_		:		str							) -> Tuple[int, float]:
     UpperCamelCase          =	'--skip_memory_metrics 0'

     UpperCamelCase          =	self.run_trainer(
         max_len=128					,					model_name=SCREAMING_SNAKE_CASE_					,					learning_rate=3E-4					,					num_train_epochs=1					,					optim=SCREAMING_SNAKE_CASE_					,					distributed=SCREAMING_SNAKE_CASE_					,					extra_args_str=SCREAMING_SNAKE_CASE_					,					do_eval=SCREAMING_SNAKE_CASE_					,					do_predict=SCREAMING_SNAKE_CASE_					,					n_gpus_to_use=1					,					)

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

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

    UpperCamelCase          =	train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value							)
    UpperCamelCase          =	train_and_return_metrics(OptimizerNames.ADAMW_BNB.value							)

    UpperCamelCase          =	gpu_alloc_mem_orig - gpu_alloc_mem_bnb

    UpperCamelCase          =	gpu_peak_mem_orig + gpu_alloc_mem_orig
    UpperCamelCase          =	gpu_peak_mem_bnb + gpu_alloc_mem_bnb
    UpperCamelCase          =	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
    UpperCamelCase          =	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(
        SCREAMING_SNAKE_CASE_					,					SCREAMING_SNAKE_CASE_					,					"""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(
        SCREAMING_SNAKE_CASE_					,					SCREAMING_SNAKE_CASE_					,					"""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(
        SCREAMING_SNAKE_CASE_					,					SCREAMING_SNAKE_CASE_					,					f"""loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"""							)







   def 		lowerCamelCase_     (			self		:		str					,					lowerCamelCase_		:		int					,					lowerCamelCase_		:		str					,					lowerCamelCase_		:		int					,					lowerCamelCase_		:		float = 3E-3					,					lowerCamelCase_		:		str = "adafactor"					,					lowerCamelCase_		:		bool = False					,					lowerCamelCase_		:		str = None					,					lowerCamelCase_		:		int = 0					,					lowerCamelCase_		:		bool = True					,					lowerCamelCase_		:		bool = True					,					lowerCamelCase_		:		bool = True					,					lowerCamelCase_		:		bool = True					,					lowerCamelCase_		:		int = None					,					):





    """simple docstring"""


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

    UpperCamelCase          =	f"""\n            --do_eval\n            --per_device_eval_batch_size 4\n            --max_eval_samples 8\n            --val_max_target_length {max_len}\n            --evaluation_strategy steps\n            --eval_steps {str(SCREAMING_SNAKE_CASE_							)}\n        """.split()

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

    UpperCamelCase          =	[]
    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:
      UpperCamelCase          =	get_gpu_count()
     UpperCamelCase          =	get_torch_dist_unique_port()
     UpperCamelCase          =	f"""\n                -m torch.distributed.run\n                --nproc_per_node={n_gpus_to_use}\n                --master_port={master_port}\n                {self.examples_dir_str}/pytorch/translation/run_translation.py\n            """.split()
     UpperCamelCase          =	[sys.executable] + distributed_args + args
     # keep for quick debug
     # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
     execute_subprocess_async(SCREAMING_SNAKE_CASE_					,					env=self.get_env()							)
    else:
     UpperCamelCase          =	['run_translation.py'] + args
     with patch.object(SCREAMING_SNAKE_CASE_					,					"""argv"""					,					SCREAMING_SNAKE_CASE_							):
      main()

    return output_dir





"""simple docstring"""







import inspect
import os
import unittest

import torch

import accelerate
from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_multi_gpu
from accelerate.utils import patch_environment







class       UpperCamelCase__   (	unittest.TestCase  ):




      """simple docstring"""




      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]       ):
        lowerCAmelCase_   :				Union[str, Any]								= inspect.getfile(accelerate.test_utils       )
        lowerCAmelCase_   :				Dict								= os.path.sep.join(mod_file.split(os.path.sep       )[:-1] + ['scripts', 'test_script.py']       )
        lowerCAmelCase_   :				int								= os.path.sep.join(
            mod_file.split(os.path.sep       )[:-1] + ['scripts', 'test_distributed_data_loop.py']       )
        lowerCAmelCase_   :				Union[str, Any]								= os.path.sep.join(mod_file.split(os.path.sep       )[:-1] + ['scripts', 'test_ops.py']       )



      @require_multi_gpu
      def      SCREAMING_SNAKE_CASE__ (      self				:      Any       ):
        print(F"Found {torch.cuda.device_count()} devices."       )
        lowerCAmelCase_   :				Optional[int]								= ['torchrun', F"--nproc_per_node={torch.cuda.device_count()}", self.test_file_path]
        with patch_environment(omp_num_threads=1       ):
          execute_subprocess_async(SCREAMING_SNAKE_CASE_    ,			env=os.environ.copy()       )



      @require_multi_gpu
      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[Any]       ):
        print(F"Found {torch.cuda.device_count()} devices."       )
        lowerCAmelCase_   :				int								= ['torchrun', F"--nproc_per_node={torch.cuda.device_count()}", self.operation_file_path]
        print(F"Command: {cmd}"       )
        with patch_environment(omp_num_threads=1       ):
          execute_subprocess_async(SCREAMING_SNAKE_CASE_    ,			env=os.environ.copy()       )



      @require_multi_gpu
      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]       ):
        lowerCAmelCase_   :				List[Any]								= ['torchrun', F"--nproc_per_node={torch.cuda.device_count()}", inspect.getfile(self.__class__       )]
        with patch_environment(omp_num_threads=1       ):
          execute_subprocess_async(SCREAMING_SNAKE_CASE_    ,			env=os.environ.copy()       )





      @require_multi_gpu
      def      SCREAMING_SNAKE_CASE__ (      self				:      Union[str, Any]       ):
        print(F"Found {torch.cuda.device_count()} devices, using 2 devices only"       )
        lowerCAmelCase_   :				Any								= ['torchrun', F"--nproc_per_node={torch.cuda.device_count()}", self.data_loop_file_path]
        with patch_environment(omp_num_threads=1    ,			cuda_visible_devices='0,1'       ):
          execute_subprocess_async(SCREAMING_SNAKE_CASE_    ,			env=os.environ.copy()       )


if __name__ == "__main__":
     lowercase__		:				Dict          =	Accelerator()
     lowercase__		:				List[Any]          =	(accelerator.state.process_index + 2, 1_0)
     lowercase__		:				Any          =	torch.randint(0, 1_0, shape).to(accelerator.device)

     lowercase__		:				List[Any]          =	""""""

     lowercase__		:				Union[str, Any]          =	accelerator.pad_across_processes(tensor)
     if tensora.shape[0] != accelerator.state.num_processes + 1:
          error_msg += f"Found shape {tensora.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
     if not torch.equal(tensora[: accelerator.state.process_index + 2], tensor):
          error_msg += "Tensors have different values."
     if not torch.all(tensora[accelerator.state.process_index + 2 :] == 0):
          error_msg += "Padding was not done with the right value (0)."

     lowercase__		:				Optional[Any]          =	accelerator.pad_across_processes(tensor, pad_first=True)
     if tensora.shape[0] != accelerator.state.num_processes + 1:
          error_msg += f"Found shape {tensora.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
     lowercase__		:				Any          =	accelerator.state.num_processes - accelerator.state.process_index - 1
     if not torch.equal(tensora[index:], tensor):
          error_msg += "Tensors have different values."
     if not torch.all(tensora[:index] == 0):
          error_msg += "Padding was not done with the right value (0)."

     # Raise error at the end to make sure we don't stop at the first failure.
     if len(error_msg) > 0:
          raise ValueError(error_msg)





