Refactor: Generalize ZeROOptimizer support and extend core APIs for BF16/ZenFlow integration

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

Author: amadhan882

deepspeed/runtime/engine.py

@@ -415,13 +415,20 @@ def __init__(self,
         self.engine_timers_cache = {}
 
         if self.optimizer_name() or self.client_optimizer is not None:
-          if self.optimizer is None:
-            raise RuntimeError("DeepSpeedEngine: Optimizer initialization failed. Check for JIT compilation errors.")
-          # ZeRO-3 specific check to prevent step 0 deadlocks
-          if self.zero_optimization_stage() == 3:
-            if not hasattr(self.optimizer, 'step'):
-              raise AttributeError("DeepSpeedEngine: ZeRO-3 optimizer is missing core functional attributes (.step). "
-                               "This usually indicates a toolchain mismatch or failed JIT kernels.")
+            if self.optimizer is None:
+                raise RuntimeError(
+                    "DeepSpeedEngine: Optimizer initialization failed. Check for JIT compilation errors.")
+
+            required_methods = ['step', 'backward', 'load_state_dict']
+
+            if self.zero_optimization_partition_gradients():
+                required_methods.append('overlapping_partition_gradients_reduce_epilogue')
+
+            for method in required_methods:
+                if not hasattr(self.optimizer, method):
+                    raise AttributeError(
+                        f"DeepSpeedEngine: Optimizer is missing core functional attribute (.{method}). "
+                        "This usually indicates a toolchain mismatch or failed JIT kernels.")
 
         if self.global_rank == 0:
             self._config.print("DeepSpeedEngine configuration")
@@ -2422,9 +2429,8 @@ 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'):
+            if hasattr(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():