Samyamr/largest partitioned params calculation fix (#1150)
* largest_partitioned_params calculation fix
largest partitioned params was getting calculated incorrectly
* Update stage3.py
* Update stage3.py
* formatting fix
* changing sub-group size default to 1e9
Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
