From 5bb4ec6233d6414a922ad2818f0bcf879de81c28 Mon Sep 17 00:00:00 2001 From: Zachary Mueller Date: Tue, 18 Apr 2023 14:25:02 -0400 Subject: [PATCH] Raise err if minimum Accelerate version isn't available (#22841) * Add warning about accelerate * Version block Accelerate * Include parse * Apply suggestions from code review Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> * Check partial state * Update param --------- Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> --- src/transformers/training_args.py | 4 ++++ src/transformers/utils/import_utils.py | 11 +++++++++-- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py index e1aee1b8d5..fc2ba42786 100644 --- a/src/transformers/training_args.py +++ b/src/transformers/training_args.py @@ -1531,6 +1531,10 @@ class TrainingArguments: def _setup_devices(self) -> "torch.device": requires_backends(self, ["torch"]) logger.info("PyTorch: setting up devices") + if not is_sagemaker_mp_enabled() and not is_accelerate_available(check_partial_state=True): + raise ImportError( + "Using the `Trainer` with `PyTorch` requires `accelerate`: Run `pip install --upgrade accelerate`" + ) if self.no_cuda: self.distributed_state = PartialState(cpu=True) device = self.distributed_state.device diff --git a/src/transformers/utils/import_utils.py b/src/transformers/utils/import_utils.py index f250a670a9..e1fc406bb6 100644 --- a/src/transformers/utils/import_utils.py +++ b/src/transformers/utils/import_utils.py @@ -575,8 +575,15 @@ def is_protobuf_available(): return importlib.util.find_spec("google.protobuf") is not None -def is_accelerate_available(): - return importlib.util.find_spec("accelerate") is not None +def is_accelerate_available(check_partial_state=False): + accelerate_available = importlib.util.find_spec("accelerate") is not None + if accelerate_available: + if check_partial_state: + return version.parse(importlib_metadata.version("accelerate")) >= version.parse("0.17.0") + else: + return True + else: + return False def is_optimum_available():