Fix step logic in DeepSpeedZeroOptimizer for calling ZenFlow

- Updated gradient tensor copying logic to accommodate ZenFlow's dimensionality and contiguity requirements.
- Added conditional handling for gradient updates based on the zenflow flag, ensuring compatibility with both standard and ZenFlow optimizers.

Signed-off-by: Tingfeng Lan <erc8gx@virginia.edu>
Co-authored-by: Yusen Wu <xrn4ub@virginia.edu>

Author: Antlera

deepspeed/runtime/zero/stage_1_and_2.py

@@ -1941,6 +1941,8 @@ def _optimizer_step(self, group_no):
         elif self.torch_autocast_gradscaler:
             self.torch_autocast_gradscaler.step(self.optimizer)
             self.torch_autocast_gradscaler.update()
+        elif self.zenflow:
+            self.zenflow_cpu_optimizer_step(group_no)
         else:
             self.optimizer.step()