Repository: microsoft/DeepSpeed

Issue:
The change of largest_partitioned_param_numel looks to be fixed, regarding of ds_numel, not numel().
The change uses tensor.numel() for retrieving the sizes. However, if we already built the model in the context of deepspeed.zero.Init(), tensor.numel() or tensor.data get pointing only the placeholder, after partitioning and offloading the original tensor into ds_tensor and ds_numel, accordingly. In some cases, this can cause some problems, including the initialization of a deepspeed model engine after deepspeed.zero.Init(). Please check out this situation and confirm the change is valid.

The original version, self.fp16_partitioned_groups_flat_numel [seems that it got retrieved from ds_numel. ](The change uses tensor.numel() for retrieving the sizes. However, if we already built the model in the context of deepspeed.zero.Init(), tensor.numel() or tensor.data get pointing only the placeholder, after partitioning and offloading the original tensor into ds_tensor and ds_numel, accordingly. In some cases, this can cause some problems, including the initialization of a deepspeed model engine after deepspeed.zero.Init(). Please check out this situation and confirm the change is valid.

The original version, self.fp16_partitioned_groups_flat_numel seems that it got retrieved from ds_numel. https://github.com/microsoft/DeepSpeed/blob/7567c76c05626c5acd8b5700bedfc412c55d5354/deepspeed/runtime/zero/stage3.py#L1151)

_Originally posted by @mrgomdev in https://github.com/microsoft/DeepSpeed/pull/1150#issuecomment-962951822_


Discussion:
@mrgomdev, thanks for raising this issue. I will take a closer look. Are you currently seeing a problem with this code?
Of course.

I can report a case caused by this issue.
When training massive minGPT model with zero-infinity for nvme as the remote device to offload, the code deciding the size of buffer refers `largest_partitioned_param_numel`. However, the change from PR #1150  , the buffer size gets restricted only by the total size of placeholders for offloading, not the actual size. This cuased crash using the buffer. Replacing `tensor.numel()` with `tensor.ds_numel` fixed our case.
Since you have fix already, do you mind sharing both the repro and a PR? We will greatly appreciation the contribution. 
The previous version (with `tensor.numel()`) seems to be working for the scenario using CPU Offload, regardless of whether `deepspeed.zero.Init()` is used or not. The latter version (with `tensor.ds_numel` fixes the error occured in the scenario using both NVMe Offload and `deepspeed.zero.Init()`. My colleague reported these two cases to me, but I'm not sure which change is the correct one for the intention of the code. (1) Revert the PR #1150  (2) Replace `tensor.numel()` with `tensor.ds_numel` if `ds_numel` exists (3) Replace it only when using both NVMe Offload and `deepspeed.zero.Init()`.
@mrgomdev, can you please try this fix for `largest_partitioned_param_numel` computation?
```python
        largest_partitioned_param_numel = max([
            max([max(tensor.numel(), tensor.ds_numel) for tensor in fp16_partitioned_group])
            for fp16_partitioned_group in self.fp16_partitioned_groups
        ])
```
Recently I have been out of the project using deepspeed and huge NN models bacause of some external reasons. I asked my colleague to try that fix and check out if the change makes expected difference compared to the original one. He would answer to it soon.
In my opinion, it looks reasonable, as long as AttributeError does not happen. If happen, my colleague would suggest and test some defensive changes for that.

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
@@ -848,8 +848,10 @@ def __init__(self,
 
         #Largest partitioned param
         largest_partitioned_param_numel = max([
-            max([tensor.numel() for tensor in fp16_partitioned_group])
-            for fp16_partitioned_group in self.fp16_partitioned_groups
+            max([
+                max(tensor.numel(),
+                    tensor.ds_numel) for tensor in fp16_partitioned_group
+            ]) for fp16_partitioned_group in self.fp16_partitioned_groups
         ])
         print_rank_0(
             f'Largest partitioned param numel = {largest_partitioned_param_numel}',
@@ -982,76 +984,6 @@ def _configure_tensor_swapping(self, offload_optimizer_config, aio_config):
             dtype=torch.float32,
             timers=self.timers)
 
-    def _create_fp16_partitions(self):
-        dist.barrier()
-        partition_id = dist.get_rank(group=self.dp_process_group)
-
-        # loop to deal with groups
-        for j, param_group in enumerate(self.optimizer.param_groups):
-
-            sub_groups = self._create_fp16_sub_groups(param_group['params'])
-            for sub_group in sub_groups:
-                i = len(self.fp16_groups)
-
-                # push this group to list before modify
-                self.fp16_groups.append(sub_group)
-                self.sub_group_to_group_id[i] = j
-
-                #These are the list of the partitioned parameters
-                self.fp16_partitioned_groups.append(
-                    [param.ds_tensor for param in self.fp16_groups[i]])
-
-                print_rank_0(
-                    f"fp16 group {i} partitioned_param norms : {[param.ds_tensor.norm().item() for param in self.fp16_groups[i]]}"
-                )
-
-                # Record padding required to align group to world size (only applies to last rank)
-                if partition_id == dist.get_world_size(group=self.dp_process_group) - 1:
-                    padding = [p.padding_size() for p in self.fp16_groups[i]]
-                else:
-                    padding = [0] * len(self.fp16_groups[i])
-                self.groups_padding.append(padding)
-
-                #not sure why apex was cloning the weights before flattening
-                #removing cloning here
-                see_memory_usage(f"Before Flattening param group {i}", force=False)
-
-                if not self.offload_param:
-                    see_memory_usage(f"Before moving param group {i} to CPU",
-                                     force=False)
-                    #move all the parameters to cpu to free up GPU space for creating flat buffer
-                    move_to_cpu(self.fp16_partitioned_groups[i])
-                    see_memory_usage(f"After moving param group {i} to CPU", force=False)
-
-                    #create flat buffer in CPU and move to GPU
-                    self.fp16_partitioned_groups_flat.append(
-                        self.flatten_dense_tensors_aligned(
-                            self.fp16_partitioned_groups[i],
-                            dist.get_world_size(group=self.dp_process_group)).cuda(
-                                torch.cuda.current_device()))
-                    see_memory_usage(
-                        f"After flattening and moving param group {i} to GPU",
-                        force=False)
-                else:
-                    #Without the detach, seems like the flattening becomes part of the
-                    #model graph causing errors downstream
-                    self.fp16_partitioned_groups_flat.append(
-                        self.flatten_dense_tensors_aligned(
-                            self.fp16_partitioned_groups[i],
-                            dist.get_world_size(
-                                group=self.dp_process_group)).detach().pin_memory())
-
-                see_memory_usage(f"After Flattening param group {i}", force=False)
-
-                see_memory_usage(f"After Flattening param group {i}", force=False)
-
-                #set model fp16 weight to slices of flattened buffer
-                updated_params = self.unflatten(self.fp16_partitioned_groups_flat[i],
-                                                self.fp16_partitioned_groups[i])
-
-                for partitioned_param, q in zip(self.fp16_partitioned_groups[i], updated_params):
-                    partitioned_param.data = q.data
-
     def _move_to_flat_buffer(self, param_list, flat_buffer, avoid_copy=False):
         '''If flat buffer is None then the parameters in the param_list are
         not copied to the flat buffer. This is because they excede the number of max_params_in_cpu


Test patch: