Samyamr/largest partitioned params calculation fix

[samyam]

largest_partiitoned_params is calculated incorrectly making it much larger than it has to be. The PR fixes the calculation.

[tjruwase]

Does self.fp16_partitioned_groups_flat_numel need to be updated?

[tjruwase]

Do references like this need updating?

[mrgomdev]

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)