chendl/compositional_test

1# Copyright 2020 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import math16import os17import re18import sys19import unittest20from pathlib import Path21from typing import Tuple22from unittest.mock import patch23 24from parameterized import parameterized25 26from transformers.testing_utils import (27    CaptureStderr,28    ExtendSysPath,29    TestCasePlus,30    execute_subprocess_async,31    get_gpu_count,32    get_torch_dist_unique_port,33    require_apex,34    require_bitsandbytes,35    require_fairscale,36    require_torch,37    require_torch_gpu,38    require_torch_multi_gpu,39    require_torch_non_multi_gpu,40    slow,41)42from transformers.trainer_callback import TrainerState43from transformers.trainer_utils import set_seed44 45 46bindir = os.path.abspath(os.path.dirname(__file__))47with ExtendSysPath(f"{bindir}/../../examples/pytorch/translation"):48    from run_translation import main  # noqa49 50 51set_seed(42)52MARIAN_MODEL = "sshleifer/student_marian_en_ro_6_1"53MBART_TINY = "sshleifer/tiny-mbart"54 55 56@require_torch57class TestTrainerExt(TestCasePlus):58    def run_seq2seq_quick(59        self,60        distributed=False,61        extra_args_str=None,62        predict_with_generate=True,63        do_train=True,64        do_eval=True,65        do_predict=True,66    ):67        output_dir = self.run_trainer(68            eval_steps=1,69            max_len=12,70            model_name=MBART_TINY,71            num_train_epochs=1,72            distributed=distributed,73            extra_args_str=extra_args_str,74            predict_with_generate=predict_with_generate,75            do_train=do_train,76            do_eval=do_eval,77            do_predict=do_predict,78        )79        logs = TrainerState.load_from_json(os.path.join(output_dir, "trainer_state.json")).log_history80 81        if not do_eval:82            return83 84        eval_metrics = [log for log in logs if "eval_loss" in log.keys()]85 86        first_step_stats = eval_metrics[0]87        if predict_with_generate:88            assert "eval_bleu" in first_step_stats89 90            last_step_stats = eval_metrics[-1]91            assert isinstance(last_step_stats["eval_bleu"], float)92            assert not math.isnan(float(last_step_stats["eval_loss"])), "eval_loss must not be `nan`"93 94    @require_torch_non_multi_gpu95    def test_run_seq2seq_no_dist(self):96        self.run_seq2seq_quick()97 98    # verify that the trainer can handle non-distributed with n_gpu > 199    @require_torch_multi_gpu100    def test_run_seq2seq_dp(self):101        self.run_seq2seq_quick(distributed=False)102 103    # verify that the trainer can handle distributed with n_gpu > 1104    @require_torch_multi_gpu105    def test_run_seq2seq_ddp(self):106        self.run_seq2seq_quick(distributed=True)107 108    # test --sharded_ddp w/o --fp16109    @unittest.skip("Requires an update of the env running those tests")110    @require_torch_multi_gpu111    @require_fairscale112    def test_run_seq2seq_sharded_ddp(self):113        self.run_seq2seq_quick(distributed=True, extra_args_str="--sharded_ddp simple")114 115    # test --sharded_ddp w/ --fp16116    @unittest.skip("Requires an update of the env running those tests")117    @require_torch_multi_gpu118    @require_fairscale119    def test_run_seq2seq_sharded_ddp_fp16(self):120        self.run_seq2seq_quick(distributed=True, extra_args_str="--sharded_ddp simple --fp16")121 122    # test --sharded_ddp zero_dp_2 w/o --fp16123    @unittest.skip("Requires an update of the env running those tests")124    @require_torch_multi_gpu125    @require_fairscale126    def test_run_seq2seq_fully_sharded_ddp(self):127        self.run_seq2seq_quick(distributed=True, extra_args_str="--sharded_ddp zero_dp_2", predict_with_generate=False)128 129    # test --sharded_ddp zero_dp_2 w/ --fp16130    @unittest.skip("Requires an update of the env running those tests")131    @require_torch_multi_gpu132    @require_fairscale133    def test_run_seq2seq_fully_sharded_ddp_fp16(self):134        self.run_seq2seq_quick(135            distributed=True, extra_args_str="--sharded_ddp zero_dp_2 --fp16", predict_with_generate=False136        )137 138    @require_apex139    @require_torch_gpu140    def test_run_seq2seq_apex(self):141        # XXX: apex breaks the trainer if it's run twice e.g. run_seq2seq.main() from the same142        # program and it breaks other tests that run from the same pytest worker, therefore until this is143        # sorted out it must be run only in an external program, that is distributed=True in this144        # test and only under one or more gpus - if we want cpu will need to make a special test145        #146        # specifically to the problem traced it to self.optimizer.step() - if it's run 2nd time via147        # 2nd main() call it botches the future eval.148        #149        self.run_seq2seq_quick(distributed=True, extra_args_str="--fp16 --fp16_backend=apex")150        # test 2nd time - was getting eval_loss': nan'151        # to reproduce the problem set distributed=False152        self.run_seq2seq_quick(distributed=True, extra_args_str="--fp16 --fp16_backend=apex")153 154    @parameterized.expand(["base", "low", "high", "mixed"])155    @require_torch_multi_gpu156    def test_trainer_log_level_replica(self, experiment_id):157        # as each sub-test is slow-ish split into multiple sub-tests to avoid CI timeout158        experiments = {159            # test with the default log_level - should be info and thus log info once160            "base": {"extra_args_str": "", "n_matches": 1},161            # test with low log_level and log_level_replica - should be noisy on all processes162            # now the info string should appear twice on 2 processes163            "low": {"extra_args_str": "--log_level debug --log_level_replica debug", "n_matches": 2},164            # test with high log_level and low log_level_replica165            # now the info string should appear once only on the replica166            "high": {"extra_args_str": "--log_level error --log_level_replica debug", "n_matches": 1},167            # test with high log_level and log_level_replica - should be quiet on all processes168            "mixed": {"extra_args_str": "--log_level error --log_level_replica error", "n_matches": 0},169        }170 171        data = experiments[experiment_id]172        kwargs = {"distributed": True, "predict_with_generate": False, "do_eval": False, "do_predict": False}173        log_info_string = "Running training"174        with CaptureStderr() as cl:175            self.run_seq2seq_quick(**kwargs, extra_args_str=data["extra_args_str"])176        n_matches = len(re.findall(log_info_string, cl.err))177        self.assertEqual(n_matches, data["n_matches"])178 179    @slow180    def test_run_seq2seq(self):181        output_dir = self.run_trainer(182            eval_steps=2,183            max_len=128,184            model_name=MARIAN_MODEL,185            learning_rate=3e-4,186            num_train_epochs=10,187            distributed=False,188        )189 190        # Check metrics191        logs = TrainerState.load_from_json(os.path.join(output_dir, "trainer_state.json")).log_history192        eval_metrics = [log for log in logs if "eval_loss" in log.keys()]193        first_step_stats = eval_metrics[0]194        last_step_stats = eval_metrics[-1]195 196        assert first_step_stats["eval_loss"] > last_step_stats["eval_loss"], "model learned nothing"197        assert isinstance(last_step_stats["eval_bleu"], float)198 199        # test if do_predict saves generations and metrics200        contents = os.listdir(output_dir)201        contents = {os.path.basename(p) for p in contents}202        assert "generated_predictions.txt" in contents203        assert "predict_results.json" in contents204 205    @slow206    @require_bitsandbytes207    def test_run_seq2seq_bnb(self):208        from transformers.training_args import OptimizerNames209 210        def train_and_return_metrics(optim: str) -> Tuple[int, float]:211            extra_args = "--skip_memory_metrics 0"212 213            output_dir = self.run_trainer(214                max_len=128,215                model_name=MARIAN_MODEL,216                learning_rate=3e-4,217                num_train_epochs=1,218                optim=optim,219                distributed=True,  # force run in a new process220                extra_args_str=extra_args,221                do_eval=False,222                do_predict=False,223                n_gpus_to_use=1,  # to allow deterministic fixed memory usage224            )225 226            # Check metrics227            logs = TrainerState.load_from_json(Path(output_dir, "trainer_state.json")).log_history228            gpu_peak_mem_mb = int(logs[0]["train_mem_gpu_peaked_delta"] / 2**20)229            gpu_alloc_mem_mb = int(logs[0]["train_mem_gpu_alloc_delta"] / 2**20)230 231            loss = logs[0]["train_loss"]232            return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss233 234        gpu_peak_mem_orig, gpu_alloc_mem_orig, loss_orig = train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value)235        gpu_peak_mem_bnb, gpu_alloc_mem_bnb, loss_bnb = train_and_return_metrics(OptimizerNames.ADAMW_BNB.value)236 237        gpu_alloc_mem_diff = gpu_alloc_mem_orig - gpu_alloc_mem_bnb238 239        gpu_total_mem_orig = gpu_peak_mem_orig + gpu_alloc_mem_orig240        gpu_total_mem_bnb = gpu_peak_mem_bnb + gpu_alloc_mem_bnb241        gpu_total_mem_diff = gpu_total_mem_orig - gpu_total_mem_bnb242 243        # sshleifer/student_marian_en_ro_6_1 has 54M parameter, 29M of which is `nn.Embedding` which244        # doesn't get quantized and remains in fp32. Therefore we only have 25M parameters quantized245        # in 2 bytes and the diff in optim memory usage is derived as so:246        #247        # - normal 25*8=~200MB (8 bytes per param)248        # - bnb    25*2= ~50MB (2 bytes per param)249        #250        # Thus we should expect ~150MB total memory saved.251        #252        # Peak memory should be the same - the total should be different by about that same margin253        #254        # After leaving a small margin to accommodate for differences between gpus let's check255        # that we have at least 120MB in savings256        expected_savings = 120257 258        # uncomment the following if this test starts failing - requires py38 for a new print feature259        # gpu_peak_mem_diff = gpu_peak_mem_orig - gpu_peak_mem_bnb260        # print(f"{gpu_alloc_mem_orig=}MB {gpu_peak_mem_orig=}MB {gpu_alloc_mem_orig+gpu_peak_mem_orig=}MB")261        # print(f" {gpu_alloc_mem_bnb=}MB  {gpu_peak_mem_bnb=}MB  {gpu_alloc_mem_bnb+gpu_peak_mem_bnb=}MB")262        # print(f"{gpu_alloc_mem_diff=}MB")263        # print(f"{gpu_peak_mem_diff=}MB")264        # print(f"{gpu_total_mem_orig=}MB, {gpu_total_mem_bnb=}MB")265        # print(f"{gpu_total_mem_diff=}MB, {gpu_total_mem_diff=}MB")266 267        self.assertGreater(268            gpu_alloc_mem_diff,269            expected_savings,270            "should use ~150MB less alloc gpu memory with BNB, compared to without it for this model but got"271            f" a difference of {gpu_alloc_mem_diff}MB, with gpu_alloc_mem_orig={gpu_alloc_mem_orig}MB and"272            f" gpu_alloc_mem_bnb={gpu_alloc_mem_bnb}MB",273        )274 275        self.assertGreater(276            gpu_total_mem_diff,277            expected_savings,278            "should use ~150MB less total gpu memory with BNB, compared to without it for this model but got"279            f" a difference of {gpu_total_mem_diff}MB, with gpu_total_mem_orig={gpu_total_mem_orig}MB and"280            f" gpu_total_mem_bnb={gpu_total_mem_bnb}MB",281        )282 283        self.assertEqual(284            loss_orig, loss_bnb, f"loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"285        )286 287    def run_trainer(288        self,289        max_len: int,290        model_name: str,291        num_train_epochs: int,292        learning_rate: float = 3e-3,293        optim: str = "adafactor",294        distributed: bool = False,295        extra_args_str: str = None,296        eval_steps: int = 0,297        predict_with_generate: bool = True,298        do_train: bool = True,299        do_eval: bool = True,300        do_predict: bool = True,301        n_gpus_to_use: int = None,302    ):303        data_dir = self.test_file_dir / "../fixtures/tests_samples/wmt_en_ro"304        output_dir = self.get_auto_remove_tmp_dir()305        args_train = f"""306            --model_name_or_path {model_name}307            --train_file {data_dir}/train.json308            --validation_file {data_dir}/val.json309            --test_file {data_dir}/test.json310            --output_dir {output_dir}311            --overwrite_output_dir312            --max_train_samples 8313            --max_source_length {max_len}314            --max_target_length {max_len}315            --do_train316            --num_train_epochs {str(num_train_epochs)}317            --per_device_train_batch_size 4318            --learning_rate {learning_rate}319            --warmup_steps 8320            --logging_steps 0321            --logging_strategy no322            --save_steps {str(eval_steps)}323            --group_by_length324            --label_smoothing_factor 0.1325            --target_lang ro_RO326            --source_lang en_XX327        """.split()328 329        args_eval = f"""330            --do_eval331            --per_device_eval_batch_size 4332            --max_eval_samples 8333            --val_max_target_length {max_len}334            --evaluation_strategy steps335            --eval_steps {str(eval_steps)}336        """.split()337 338        args_predict = """339            --do_predict340        """.split()341 342        args = []343        if do_train:344            args += args_train345 346        if do_eval:347            args += args_eval348 349        if do_predict:350            args += args_predict351 352        if predict_with_generate:353            args += "--predict_with_generate".split()354 355        if do_train:356            if optim == "adafactor":357                args += "--adafactor".split()358            else:359                args += f"--optim {optim}".split()360 361        if extra_args_str is not None:362            args += extra_args_str.split()363 364        if distributed:365            if n_gpus_to_use is None:366                n_gpus_to_use = get_gpu_count()367            master_port = get_torch_dist_unique_port()368            distributed_args = f"""369                -m torch.distributed.run370                --nproc_per_node={n_gpus_to_use}371                --master_port={master_port}372                {self.examples_dir_str}/pytorch/translation/run_translation.py373            """.split()374            cmd = [sys.executable] + distributed_args + args375            # keep for quick debug376            # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die377            execute_subprocess_async(cmd, env=self.get_env())378        else:379            testargs = ["run_translation.py"] + args380            with patch.object(sys, "argv", testargs):381                main()382 383        return output_dir384