






'''simple docstring'''

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


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


set_seed(42)
UpperCAmelCase__						:str          =			"""sshleifer/student_marian_en_ro_6_1"""
UpperCAmelCase__						:Optional[int]          =			"""sshleifer/tiny-mbart"""







@require_torch
class 	SCREAMING_SNAKE_CASE	(      lowerCAmelCase_  ):




       def 					a_					(					self      : List[str]	,    A__      : int=False	,    A__      : Tuple=None	,    A__      : Dict=True	,    A__      : Optional[int]=True	,    A__      : str=True	,    A__      : Union[str, Any]=True	,    ):


            """simple docstring"""







            __lowerCamelCase						:       str				   =				self.run_trainer(
                eval_steps=1	,    max_len=12	,    model_name=A__	,    num_train_epochs=1	,    distributed=A__	,    extra_args_str=A__	,    predict_with_generate=A__	,    do_train=A__	,    do_eval=A__	,    do_predict=A__	,    )
            __lowerCamelCase						:       Optional[int]				   =				TrainerState.load_from_json(os.path.join(A__	,    """trainer_state.json"""					)					).log_history

            if not do_eval:
                 return

            __lowerCamelCase						:       Optional[Any]				   =				[log for log in logs if """eval_loss""" in log.keys()]

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

                 __lowerCamelCase						:       Optional[Any]				   =				eval_metrics[-1]
                 assert isinstance(last_step_stats["""eval_bleu"""]	,    A__					)
                 assert not math.isnan(float(last_step_stats["""eval_loss"""]					)					), "eval_loss must not be `nan`"




       @require_torch_non_multi_gpu
       def 					a_					(					self      : Dict					):


            """simple docstring"""







            self.run_seqaseq_quick()




       @require_torch_multi_gpu
       def 					a_					(					self      : int					):


            """simple docstring"""







            self.run_seqaseq_quick(distributed=A__					)




       @require_torch_multi_gpu
       def 					a_					(					self      : List[Any]					):


            """simple docstring"""







            self.run_seqaseq_quick(distributed=A__					)




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


            """simple docstring"""







            self.run_seqaseq_quick(distributed=A__	,    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      : int					):


            """simple docstring"""







            self.run_seqaseq_quick(distributed=A__	,    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      : Optional[Any]					):


            """simple docstring"""







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




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


            """simple docstring"""







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




       @require_apex
       @require_torch_gpu
       def 					a_					(					self      : Optional[int]					):


            """simple docstring"""







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




       @parameterized.expand(["""base""", """low""", """high""", """mixed"""]					)
       @require_torch_multi_gpu
       def 					a_					(					self      : List[Any]	,    A__      : List[str]					):


            """simple docstring"""







            __lowerCamelCase						:       Union[str, Any]				   =				{
                # 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},
            }

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




       @slow
       def 					a_					(					self      : Optional[int]					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[int]				   =				self.run_trainer(
                eval_steps=2	,    max_len=128	,    model_name=A__	,    learning_rate=3e-4	,    num_train_epochs=10	,    distributed=A__	,    )

            # Check metrics
            __lowerCamelCase						:       str				   =				TrainerState.load_from_json(os.path.join(A__	,    """trainer_state.json"""					)					).log_history
            __lowerCamelCase						:       int				   =				[log for log in logs if """eval_loss""" in log.keys()]
            __lowerCamelCase						:       str				   =				eval_metrics[0]
            __lowerCamelCase						:       List[Any]				   =				eval_metrics[-1]

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

            # test if do_predict saves generations and metrics
            __lowerCamelCase						:       Union[str, Any]				   =				os.listdir(A__					)
            __lowerCamelCase						:       List[Any]				   =				{os.path.basename(A__					) for p in contents}
            assert "generated_predictions.txt" in contents
            assert "predict_results.json" in contents




       @slow
       @require_bitsandbytes
       def 					a_					(					self      : List[Any]					):


            """simple docstring"""







            from transformers.training_args import OptimizerNames

            def train_and_return_metrics(A__      : str					) -> Tuple[int, float]:
                 __lowerCamelCase						:       Optional[Any]				   =				"""--skip_memory_metrics 0"""

                 __lowerCamelCase						:       Optional[Any]				   =				self.run_trainer(
                     max_len=128	,    model_name=A__	,    learning_rate=3e-4	,    num_train_epochs=1	,    optim=A__	,    distributed=A__	,    extra_args_str=A__	,    do_eval=A__	,    do_predict=A__	,    n_gpus_to_use=1	,    )

                 # Check metrics
                 __lowerCamelCase						:       Optional[int]				   =				TrainerState.load_from_json(Path(A__	,    """trainer_state.json"""					)					).log_history
                 __lowerCamelCase						:       Tuple				   =				int(logs[0]["""train_mem_gpu_peaked_delta"""] / 2**20					)
                 __lowerCamelCase						:       List[str]				   =				int(logs[0]["""train_mem_gpu_alloc_delta"""] / 2**20					)

                 __lowerCamelCase						:       str				   =				logs[0]["""train_loss"""]
                 return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

            __lowerCamelCase   ,      __lowerCamelCase   ,      __lowerCamelCase						:       Dict				   =				train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value					)
            __lowerCamelCase   ,      __lowerCamelCase   ,      __lowerCamelCase						:       List[Any]				   =				train_and_return_metrics(OptimizerNames.ADAMW_BNB.value					)

            __lowerCamelCase						:       Any				   =				gpu_alloc_mem_orig - gpu_alloc_mem_bnb

            __lowerCamelCase						:       str				   =				gpu_peak_mem_orig + gpu_alloc_mem_orig
            __lowerCamelCase						:       Tuple				   =				gpu_peak_mem_bnb + gpu_alloc_mem_bnb
            __lowerCamelCase						:       str				   =				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
            __lowerCamelCase						:       Union[str, Any]				   =				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(
                A__	,    A__	,    """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(
                A__	,    A__	,    """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(
                A__	,    A__	,    f"loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"					)




       def 					a_					(					self      : int	,    A__      : int	,    A__      : str	,    A__      : int	,    A__      : float = 3e-3	,    A__      : str = "adafactor"	,    A__      : bool = False	,    A__      : str = None	,    A__      : int = 0	,    A__      : bool = True	,    A__      : bool = True	,    A__      : bool = True	,    A__      : bool = True	,    A__      : int = None	,    ):


            """simple docstring"""







            __lowerCamelCase						:       int				   =				self.test_file_dir / """../fixtures/tests_samples/wmt_en_ro"""
            __lowerCamelCase						:       Any				   =				self.get_auto_remove_tmp_dir()
            __lowerCamelCase						:       List[str]				   =				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(A__					)}\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(A__					)}\n            --group_by_length\n            --label_smoothing_factor 0.1\n            --target_lang ro_RO\n            --source_lang en_XX\n        ".split()

            __lowerCamelCase						:       List[Any]				   =				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(A__					)}\n        ".split()

            __lowerCamelCase						:       Union[str, Any]				   =				"""
            --do_predict
        """.split()

            __lowerCamelCase						:       Optional[int]				   =				[]
            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:
                      __lowerCamelCase						:       Any				   =				get_gpu_count()
                 __lowerCamelCase						:       Dict				   =				get_torch_dist_unique_port()
                 __lowerCamelCase						:       Union[str, Any]				   =				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()
                 __lowerCamelCase						:       int				   =				[sys.executable] + distributed_args + args
                 # keep for quick debug
                 # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
                 execute_subprocess_async(A__	,    env=self.get_env()					)
            else:
                 __lowerCamelCase						:       int				   =				["""run_translation.py"""] + args
                 with patch.object(A__	,    """argv"""	,    A__					):
                      main()

            return output_dir











'''simple docstring'''

from transformers import HfArgumentParser, TensorFlowBenchmark, TensorFlowBenchmarkArguments
def     __lowercase  ()   ->     str:


     """simple docstring"""
     __lowerCamelCase						:       Any				   =				HfArgumentParser(_lowercase     )
     __lowerCamelCase						:       List[str]				   =				parser.parse_args_into_dataclasses()[0]
     __lowerCamelCase						:       Dict				   =				TensorFlowBenchmark(args=_lowercase     )
     try:
          __lowerCamelCase						:       List[str]				   =				parser.parse_args_into_dataclasses()[0]
     except ValueError as e:
          __lowerCamelCase						:       Tuple				   =				"""Arg --no_{0} is no longer used, please use --no-{0} instead."""
          __lowerCamelCase						:       List[Any]				   =				""" """.join(str(_lowercase     ).split(""" """     )[:-1]     )
          __lowerCamelCase						:       Tuple				   =				""""""
          __lowerCamelCase						:       List[str]				   =				eval(str(_lowercase     ).split(""" """     )[-1]     )
          __lowerCamelCase						:       Tuple				   =				[]
          for arg in depreciated_args:
               # arg[2:] removes '--'
               if arg[2:] in TensorFlowBenchmark.deprecated_args:
                    # arg[5:] removes '--no_'
                    full_error_msg += arg_error_msg.format(arg[5:]     )
               else:
                    wrong_args.append(_lowercase     )
          if len(_lowercase     ) > 0:
               __lowerCamelCase						:       Tuple				   =				full_error_msg + begin_error_msg + str(_lowercase     )
          raise ValueError(_lowercase     )
     benchmark.run()


if __name__ == "__main__":
   main()



