* 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: NOlatunji Ruwase <olruwase@microsoft.com>