diff --git a/deepspeed/runtime/zero/constants.py b/deepspeed/runtime/zero/constants.py
index eaeb2a95ccd9..7beebe00e717 100755
--- a/deepspeed/runtime/zero/constants.py
+++ b/deepspeed/runtime/zero/constants.py
@@ -91,7 +91,7 @@
 ZERO_OPTIMIZATION_OFFLOAD_OPTIMIZER_DEFAULT = None
 
 ZERO_OPTIMIZATION_SUB_GROUP_SIZE = 'sub_group_size'
-ZERO_OPTIMIZATION_SUB_GROUP_SIZE_DEFAULT = 1000000000000
+ZERO_OPTIMIZATION_SUB_GROUP_SIZE_DEFAULT = 1000000000
 
 #maximum number of parameters per GPU before releasing them
 ZERO_OPTIMIZATION_MAX_LIVE_PARAMETERS = 'stage3_max_live_parameters'
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
index cac206ed759b..548c38a072c3 100755
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -852,7 +852,10 @@ def __init__(self,
                 count = count + 1
 
         #Largest partitioned param
-        largest_partitioned_param_numel = max(self.fp16_partitioned_groups_flat_numel)
+        largest_partitioned_param_numel = max([
+            max([tensor.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}',
             force=True)