Assert on a missing gradient in the ZenFlow parallel offload copy

ZenFlowZeroOptimizerParallel.async_inplace_copy_grad_to_fp32_buffer_from_gpu
branches on `grad_accum is None` and then calls `.view(-1)` on it in both
branches, so the None branch raises AttributeError instead of handling None.

DeepSpeedZeroOptimizer's own copy of this method asserts the attribute is
present rather than branching. Do the same here: identical behaviour when the
gradient exists, and a clear assertion instead of an AttributeError when it
does not.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>

Author: vineethsaivs

deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py

@@ -685,10 +685,9 @@ def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param):
             0, dest_offset, num_elements)
 
         grad_accum = self.get_param_gradient_attribute(param)
-        if grad_accum is None:
-            src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements)
-        else:
-            src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements)
+        assert grad_accum is not None
+
+        src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements)
         if src_tensor.dtype != self.master_weights_and_grads_dtype:
             src_tensor = src_tensor.to(self.master_weights_and_grads_dtype)