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 sys16from typing import Dict17 18from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available19from transformers.testing_utils import (20    TestCasePlus,21    execute_subprocess_async,22    get_torch_dist_unique_port,23    require_torch_multi_gpu,24    require_torch_neuroncore,25)26from transformers.utils import logging27 28 29logger = logging.get_logger(__name__)30 31 32if is_torch_available():33    import torch34    from torch import nn35    from torch.utils.data import Dataset36 37    from transformers import Trainer38 39    class DummyDataset(Dataset):40        def __init__(self, length: int = 101):41            self.length = length42 43        def __len__(self):44            return self.length45 46        def __getitem__(self, i) -> int:47            return i48 49    class DummyDataCollator:50        def __call__(self, features):51            return {"input_ids": torch.tensor(features), "labels": torch.tensor(features)}52 53    class DummyModel(nn.Module):54        def __init__(self):55            super().__init__()56            # Add some (unused) params otherwise DDP will complain.57            self.fc = nn.Linear(120, 80)58 59        def forward(self, input_ids, labels=None):60            if labels is not None:61                return torch.tensor(0.0, device=input_ids.device), input_ids62            else:63                return input_ids64 65 66class TestTrainerDistributedNeuronCore(TestCasePlus):67    @require_torch_neuroncore68    def test_trainer(self):69        distributed_args = f"""70            -m torch.distributed.run71            --nproc_per_node=272            --master_port={get_torch_dist_unique_port()}73            {self.test_file_dir}/test_trainer_distributed.py74        """.split()75        output_dir = self.get_auto_remove_tmp_dir()76        args = f"--output_dir {output_dir}".split()77        cmd = [sys.executable] + distributed_args + args78        execute_subprocess_async(cmd, env=self.get_env())79        # successful return here == success - any errors would have caused an error in the sub-call80 81 82class TestTrainerDistributed(TestCasePlus):83    @require_torch_multi_gpu84    def test_trainer(self):85        distributed_args = f"""86            -m torch.distributed.run87            --nproc_per_node={torch.cuda.device_count()}88            --master_port={get_torch_dist_unique_port()}89            {self.test_file_dir}/test_trainer_distributed.py90        """.split()91        output_dir = self.get_auto_remove_tmp_dir()92        args = f"--output_dir {output_dir}".split()93        cmd = [sys.executable] + distributed_args + args94        execute_subprocess_async(cmd, env=self.get_env())95        # successful return here == success - any errors would have caused an error in the sub-call96 97 98if __name__ == "__main__":99    # The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:100    #101    # PYTHONPATH="src" python -m torch.distributed.run --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py102 103    parser = HfArgumentParser((TrainingArguments,))104    training_args = parser.parse_args_into_dataclasses()[0]105 106    logger.warning(107        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "108        f"distributed training: {training_args.local_rank != -1}"109    )110 111    # Essentially, what we want to verify in the distributed case is that we get all samples back,112    # in the right order. (this is crucial for prediction for instance)113    for dataset_length in [101, 40, 7]:114        dataset = DummyDataset(dataset_length)115 116        def compute_metrics(p: EvalPrediction) -> Dict:117            sequential = list(range(len(dataset)))118            success = p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential119            if not success and training_args.local_rank == 0:120                logger.warning(121                    "Predictions and/or labels do not match expected results:\n  - predictions: "122                    f"{p.predictions.tolist()}\n  - labels: {p.label_ids.tolist()}\n  - expected: {sequential}"123                )124            return {"success": success}125 126        trainer = Trainer(127            model=DummyModel(),128            args=training_args,129            data_collator=DummyDataCollator(),130            eval_dataset=dataset,131            compute_metrics=compute_metrics,132        )133        metrics = trainer.evaluate()134        logger.info(metrics)135        if metrics["eval_success"] is not True:136            logger.error(metrics)137            exit(1)138 139        p = trainer.predict(dataset)140        logger.info(p.metrics)141        if p.metrics["test_success"] is not True:142            logger.error(p.metrics)143            exit(1)144 145        trainer.args.eval_accumulation_steps = 2146 147        metrics = trainer.evaluate()148        logger.info(metrics)149        if metrics["eval_success"] is not True:150            logger.error(metrics)151            exit(1)152 153        p = trainer.predict(dataset)154        logger.info(p.metrics)155        if p.metrics["test_success"] is not True:156            logger.error(p.metrics)157            exit(1)158 159        trainer.args.eval_accumulation_steps = None160