[Auto Parallel] Sharding Optimization：Partition Algorithm & Stage2 Parameter Bucket communication

python/paddle/distributed/auto_parallel/constants.py

@@ -82,7 +82,9 @@ def set_field_default_config(category, field, default_value):
 set_field_default_config(SHARDING, "enable", False)
 set_field_default_config(SHARDING, "stage", 1)
 set_field_default_config(SHARDING, "degree", 8)
-set_field_default_config(SHARDING, "segment_broadcast_MB", 32.0)
+set_field_default_config(SHARDING, "overlap_grad_comm", False)
+set_field_default_config(SHARDING, "bucket_size_numel", -1)
+set_field_default_config(SHARDING, "partition_algor", "greedy_even")
 set_field_default_config(SHARDING, "enable_tuning", False)
 set_field_default_config(SHARDING, "tuning_range", [])
 

python/paddle/distributed/auto_parallel/utils.py

@@ -22,6 +22,7 @@
 from functools import reduce
 
 import paddle.fluid.core as core
+from paddle.fluid.framework import Variable
 from paddle.distributed.fleet.meta_optimizers.common import OpRole
 from paddle.distributed.auto_parallel.process_group import (
     get_all_process_groups,
@@ -1790,6 +1791,18 @@ def find_higher_order_backward_op(program):
     return False
 
 
+def get_var_numel(var):
+    """
+    input:
+        - var: variable
+    return:
+        number of elemnet in var
+    """
+    assert isinstance(var, Variable)
+    assert -1 not in var.shape
+    return reduce(lambda x, y: x * y, var.shape)
+
+
 def get_lr(optimizer):
     if isinstance(optimizer, paddle.optimizer.Optimizer):
         return optimizer.get_lr()

python/paddle/distributed/passes/auto_parallel_data_parallel_optimization.py

@@ -13,12 +13,12 @@
 # limitations under the License.
 
 from collections import OrderedDict
-import numpy as np
 
 import paddle
 from paddle.fluid import unique_name
 from paddle.fluid.framework import default_main_program
 from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
+from .pass_base import PassBase, PassType, register_pass
 from paddle.distributed.auto_parallel.operators.common import (
     is_data_parallel_scale_op,
     is_data_parallel_reduce_op,
@@ -28,8 +28,8 @@
     is_loss_grad_op,
     is_optimize_op,
     ring_id_to_process_group,
+    get_var_numel,
 )
-from .pass_base import PassBase, PassType, register_pass
 
 # add new optimizers supporting rescale_grad here
 __rescale_grad_supported_opts__ = [
@@ -44,10 +44,6 @@
 __max_stream_num_allow__ = 16
 
 
-def numel(var):
-    return np.prod(list(var.shape))
-
-
 @register_pass("auto_parallel_data_parallel_optimization")
 class DataParallelOptimizationPass(PassBase):
     """
@@ -430,7 +426,7 @@ def op_depend_on_group(op, group):
                 ring_id = op.attr("ring_id")
                 grad_name = op.output_arg_names[0]
                 grad_var = block.var(grad_name)
-                grad_numel = numel(grad_var)
+                grad_numel = get_var_numel(grad_var)
 
                 if cur_group.acceptable(grad_var, ring_id):
                     assert grad_name not in grouped_grad_names
@@ -594,7 +590,7 @@ def acceptable(self, grad_var, ring_id):
             return True
         if ring_id != self.ring_id:
             return False
-        if numel(grad_var) + self.numel > self.max_group_size:
+        if get_var_numel(grad_var) + self.numel > self.max_group_size:
             return False
         if grad_var.dtype != self.dtype:
             return False
@@ -605,7 +601,7 @@ def add(self, grad_var, ring_id, i):
         self.gradients.append(grad_var)
         self.ring_id = ring_id
         self.dtype = grad_var.dtype
-        self.numel += numel(grad_var)
+        self.numel += get_var_numel(grad_var)
 
         # remove auxiliary ops in non-fuse dp allreduce
         self.remove_allreduce_op_indices.append(i)