[REQUEST] DeepSpeed Zero3 swap off gradients unnecessarily when swap_optimizer is True

** DeepSpeed Zero3 swap off gradients unnecessarily when swap_optimizer is True

According to the function DeepSpeedZeroOptimizer_Stage3.partition_grads in
https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage3.py#L1194, gradients are accumulated into tensors holding by self.__param_id_to_grad_partition, so why bother swap out gradients in gradient accumulation boundary?
https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage3.py#L1258

        if self.offload_optimizer and self.swap_optimizer:
            for i in offload_fp32_gradients.keys():
                self.optimizer_swapper.swap_out_gradients(parameter=self.fp32_partitioned_groups_flat[i],
                                                          gradient_offsets=offload_fp32_offsets[i],
                                                          gradient_tensors=offload_fp32_gradients[i])

I'm quite confused by this, please help me out. thanks!

@tjruwase

Start in deepspeed/runtime/zero/stage3.py at DeepSpeedZeroOptimizer_Stage3.partition_grads, especially the gradient accumulation boundary around the cited lines. Trace self.__param_id_to_grad_partition, offload_fp32_gradients, and the optimizer_swapper.swap_out_gradients call to determine whether the swap is necessary when swap_optimizer is enabled. Done means explaining the behavior and identifying whether a code change is warranted.