From 2de1b76d8100ffdfbc85419265c684012bedd546 Mon Sep 17 00:00:00 2001 From: Ali Abdelkader Date: Thu, 14 Mar 2024 17:25:05 +0200 Subject: [PATCH] set "use_cpu" to true in cpu pytorch training benchmark tests (#159) - set training_arguments.use_cpu to true so that the Trainer class runs training on the cpu instead of utilizing the available gpu --- optimum_benchmark/backends/pytorch/backend.py | 6 ++++++ tests/configs/cpu_training_pytorch_bert_sweep.yaml | 4 ++++ tests/configs/cpu_training_pytorch_gpt_sweep.yaml | 4 ++++ 3 files changed, 14 insertions(+) diff --git a/optimum_benchmark/backends/pytorch/backend.py b/optimum_benchmark/backends/pytorch/backend.py index 7dd2d7c8..4e2a853b 100644 --- a/optimum_benchmark/backends/pytorch/backend.py +++ b/optimum_benchmark/backends/pytorch/backend.py @@ -345,6 +345,12 @@ def train( ) -> TrainerState: LOGGER.info(f"\t+ Wrapping training arguments with {TrainingArguments.__name__}") training_arguments = TrainingArguments(**training_arguments) + + if self.config.device == "cpu": + assert training_arguments.use_cpu, f"""backend.device is set to {self.config.device} while use_cpu in training_arguments is set to false (default) + which will make the training run on gpu. Please set benchmark.training.training_arguments.use_cpu to true if the benchmark is intended to run on cpu + """ + LOGGER.info(f"\t+ Wrapping model with {Trainer.__name__}") trainer = Trainer( args=training_arguments, diff --git a/tests/configs/cpu_training_pytorch_bert_sweep.yaml b/tests/configs/cpu_training_pytorch_bert_sweep.yaml index 70651ba2..4779e335 100644 --- a/tests/configs/cpu_training_pytorch_bert_sweep.yaml +++ b/tests/configs/cpu_training_pytorch_bert_sweep.yaml @@ -8,3 +8,7 @@ defaults: - _self_ # hydra 1.1 compatibility experiment_name: cpu_training_pytorch_bert_sweep + +benchmark: + training_arguments: + use_cpu: true diff --git a/tests/configs/cpu_training_pytorch_gpt_sweep.yaml b/tests/configs/cpu_training_pytorch_gpt_sweep.yaml index 5f8987b6..76d97b41 100644 --- a/tests/configs/cpu_training_pytorch_gpt_sweep.yaml +++ b/tests/configs/cpu_training_pytorch_gpt_sweep.yaml @@ -8,3 +8,7 @@ defaults: - _self_ # hydra 1.1 compatibility experiment_name: cpu_training_pytorch_gpt_sweep + +benchmark: + training_arguments: + use_cpu: true