Fix: Add optimizer initialization validation for backward and ZeRO-3

Signed-off-by: amadhan882 <amadhan882@gmail.com>

Author: amadhan882

deepspeed/runtime/engine.py

@@ -419,7 +419,7 @@ def __init__(self,
                 raise RuntimeError(
                     "DeepSpeedEngine: Optimizer initialization failed. Check for JIT compilation errors.")
 
-            optimizer_methods = ['step', 'load_state_dict']
+            optimizer_methods = ['step', 'load_state_dict','backward']
 
             if self.zero_optimization_partition_gradients():
                 optimizer_methods.append('overlapping_partition_gradients_reduce_epilogue')
@@ -432,13 +432,6 @@ def __init__(self,
                         "This indicates incomplete initialization (e.g., JIT/toolchain failure)."
                     )
 
-            # Validate engine separately
-            if not hasattr(self, "backward") or not callable(getattr(self, "backward")):
-                raise RuntimeError(
-                    "DeepSpeedEngine initialization failed: missing callable `backward`. "
-                    "Engine may be partially initialized."
-                )
-
         if self.global_rank == 0:
             self._config.print("DeepSpeedEngine configuration")
             if self.dump_state():
@@ -2438,8 +2431,7 @@ def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE):
         self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary()
         # ZeRO stage >= 2 communicates during non gradient accumulation boundaries as well
         if self.zero_optimization_partition_gradients():
-            if hasattr(self.optimizer, 'overlapping_partition_gradients_reduce_epilogue'):
-                self.optimizer.overlapping_partition_gradients_reduce_epilogue()
+            self.optimizer.overlapping_partition_gradients_reduce_epilogue()
 
         # Communicate only at gradient accumulation boundaries
         elif self.is_gradient_accumulation_boundary():