regression: deepspeed check fix

bghira · Oct 11, 2024 · 66d288f · 66d288f
1 parent d785fb5
commit 66d288f
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/helpers/training/trainer.py b/helpers/training/trainer.py
@@ -1343,9 +1343,11 @@ def init_unload_vae(self):
 
     def init_validations(self):
         if (
-            self.accelerator.state.deepspeed_plugin.deepspeed_config[
-                "zero_optimization"
-            ].get("stage")
+            hasattr(self.accelerator, "state")
+            and hasattr(self.accelerator.state, "deepspeed_plugin")
+            and getattr(self.accelerator.state.deepspeed_plugin, "deepspeed_config", {})
+            .get("zero_optimization", {})
+            .get("stage")
             == 3
         ):
             logger.error("Cannot run validations with DeepSpeed ZeRO stage 3.")