Repository: microsoft/DeepSpeed

Issue:
[BUG] Stage3 optimizer offload at partition_grads function
**Describe the bug**
These code will have logical bug follow the [1], [2], [3] comment.
[1] The dictionary mapping `offload_fp32_gradients = {}` and `offload_fp32_offsets = {}` have been initialized at the begining of function.
[2] The dictionary mapping `offload_fp32_gradients = {}` and `offload_fp32_offsets = {}`  are re-initialized at the **inner-loop**
[3] The optimizer swap out is at the outer of the loop. It will only swap out the last iteration result of for-loop .

```    python
@instrument_w_nvtx
def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None:
        offload_fp32_gradients = {}    # [1] Here is the globa initalize 
        offload_fp32_offsets = {}       # [1] Here is the globa initalize 
        buffers = []
        for param, grad_partition in zip(params_to_release, grad_partitions):

            contains_real_data = param.partition_numel() * dist.get_rank(self.dp_process_group) < param.ds_numel
            if not contains_real_data:
                # this grad partition is empty - don't need to do anything
                param.grad = None
                continue

            # move or accumulate gradient partition to target buffer
            grad_buffer = self.__param_id_to_grad_partition[param.ds_id].narrow(0, 0, grad_partition.numel())
            buffers.append(grad_buffer)
            if self.micro_step_id == 0:  # don't accumulate
                grad_buffer.copy_(grad_partition, non_blocking=True)
                # ensure grad buffer is a CUDA buffer to speed up the next few
                # operations and so it can be used asynchronously
                grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True)
            elif get_accelerator().on_accelerator(grad_buffer):
                grad_buffer.add_(grad_partition.to(self.gradient_accumulation_dtype).view(grad_buffer.shape))
            else:
                # if dst is CPU, copy first to src device, do the addition
                # there, then move back to dst. adding directly to cpu is very slow
                cuda_grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True)
                cuda_grad_buffer.add_(grad_partition.to(self.gradient_accumulation_dtype).view(cuda_grad_buffer.shape))
                grad_buffer.copy_(cuda_grad_buffer, non_blocking=True)
                # ensure grad buffer is a CUDA buffer to speed up the next few
                # operations and so it can be used asynchronously
                grad_buffer = cuda_grad_buffer

            # offload the gradient partition if applicable
            if self.offload_optimizer:
                i, dest_offset, _ = self.grad_position[self.get_param_id(param)]
                offload_fp32_gradients = {}     # [2] FIXED IT: Reinialize at the inner loop
                offload_fp32_offsets = {}         # [2] FIXED IT: Reinialize at the inner loop

                if self.is_gradient_accumulation_boundary:
                    self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer)

                    if self._swappable_optimizer_subgroup(i):
                        if not i in offload_fp32_gradients.keys():
                            offload_fp32_gradients[i] = []
                            offload_fp32_offsets[i] = []

                        offload_fp32_gradients[i].append(grad_buffer.float())
                        offload_fp32_offsets[i].append(dest_offset)
                    else:
                        fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow(
                            0, dest_offset, grad_buffer.numel())
                        fp32_grad_tensor.copy_(grad_buffer)

            # free the gradient
            if not get_accelerator().is_synchronized_device():
                param.grad.record_stream(get_accelerator().current_stream())
            param.grad = None

        if self.offload_optimizer and self.swap_optimizer:
            for i in offload_fp32_gradients.keys():   # [3] Here will only contain the last for-loop result
                self.optimizer_swapper.swap_out_gradients(parameter=self.fp32_partitioned_groups_flat[i],
                                                          gradient_offsets=offload_fp32_offsets[i],        # [3]
                                                          gradient_tensors=offload_fp32_gradients[i])  # [3]
        return buffers
```




Discussion:
@PhisonAI, thanks for this issue and analysis. Your analysis is correct. This is the same issue as #4565. 

Resolving patch:
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -1376,8 +1376,6 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L
             # offload the gradient partition if applicable
             if self.offload_optimizer:
                 i, dest_offset, _ = self.grad_position[self.get_param_id(param)]
-                offload_fp32_gradients = {}
-                offload_fp32_offsets = {}
 
                 if self.is_gradient_accumulation_boundary:
                     self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer)


Test patch: