【SPMD】 fix bug in layernorm spmd in dy2stat (#65835)

liym27 · web-flow · commit 018035b1037e · 2024-07-09T20:18:18.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/operators/dist_layer_norm.py b/python/paddle/distributed/auto_parallel/static/operators/dist_layer_norm.py
@@ -18,6 +18,7 @@
 from paddle.base.log_helper import get_logger
 
 from ..completion import get_phi_spmd_rule
+from ..dist_attribute import DistTensorSpec, TensorDistAttr
 from ..utils import get_dist_tensor_spec, is_dim_shard
 from .common import (
     DistributedOperatorImplContainer,
@@ -41,16 +42,30 @@ def update_dims_mapping(dist_op):
         op_desc = dist_op.serial_op.desc
 
         x_name = op_desc.input('X')[0]
-        scale_name = op_desc.input('Scale')[0]
-        bias_name = op_desc.input('Bias')[0]
+        scale_name = (
+            op_desc.input('Scale')[0]
+            if len(op_desc.input('Scale')) > 0
+            else None
+        )
+        bias_name = (
+            op_desc.input('Bias')[0] if len(op_desc.input('Bias')) > 0 else None
+        )
         y_name = op_desc.output('Y')[0]
         var_name = op_desc.output('Variance')[0]
         mean_name = op_desc.output('Mean')[0]
         begin_norm_axis = op_desc.attr('begin_norm_axis')
 
         x_spec = get_dist_tensor_spec(dist_op, x_name)
-        scale_spec = get_dist_tensor_spec(dist_op, scale_name)
-        bias_spec = get_dist_tensor_spec(dist_op, bias_name)
+        scale_spec = (
+            DistTensorSpec([0], TensorDistAttr())
+            if scale_name is None
+            else get_dist_tensor_spec(dist_op, scale_name)
+        )
+        bias_spec = (
+            DistTensorSpec([0], TensorDistAttr())
+            if bias_name is None
+            else get_dist_tensor_spec(dist_op, bias_name)
+        )
         y_spec = get_dist_tensor_spec(dist_op, y_name, False)
         var_spec = get_dist_tensor_spec(dist_op, var_name, False)
         mean_spec = get_dist_tensor_spec(dist_op, mean_name, False)
@@ -74,9 +89,14 @@ def update_dims_mapping(dist_op):
 
         # step3: update dist_attr
         # tensor order following order in PHI definition
+        input_arg_names = [x_name]
+        if scale_name is not None:
+            input_arg_names.append(scale_name)
+        if bias_name is not None:
+            input_arg_names.append(bias_name)
         changed = update_op_dims_mapping(
             dist_op,
-            [x_name, scale_name, bias_name],
+            input_arg_names,
             [y_name, var_name, mean_name],
             fw_results,
             bw_results,
diff --git a/test/auto_parallel/spmd_rules/test_layer_norm_rule.py b/test/auto_parallel/spmd_rules/test_layer_norm_rule.py
@@ -150,6 +150,64 @@ def test_infer_forward(self):
         self.assertEqual(infered_output_dist_attrs[1].dims_mapping, [0])
         self.assertEqual(infered_output_dist_attrs[2].dims_mapping, [0])
 
+    def test_infer_forward_without_bias(self):
+        # ijk[1, -1, -1], k[-1], k[-1] -->
+        # ijk[1, -1, -1], k[-1], k[-1], (inputs)
+        # ijk[1, -1, -1], ij[1, -1], ij[1, -1],(outputs)
+        # begin_norm_axis=2
+        self.x_spec.set_dims_mapping([1, -1, -1])
+        self.bias_spec = DistTensorSpec([0], TensorDistAttr())
+        self.scale_spec.set_dims_mapping([-1])
+
+        result_dist_attrs = self.rule.infer_forward(
+            self.x_spec,
+            self.scale_spec,
+            self.bias_spec,
+            self.attrs['epsilon'],
+            self.attrs['begin_norm_axis'],
+        )
+        infered_input_dist_attrs = result_dist_attrs[0]
+        infered_output_dist_attrs = result_dist_attrs[1]
+
+        self.assertEqual(len(result_dist_attrs), 2)
+        self.assertEqual(len(infered_input_dist_attrs), 3)
+        self.assertEqual(len(infered_output_dist_attrs), 3)
+
+        self.assertEqual(infered_input_dist_attrs[0].dims_mapping, [1, -1, -1])
+        self.assertEqual(infered_input_dist_attrs[1].dims_mapping, [-1])
+        self.assertEqual(infered_input_dist_attrs[2].dims_mapping, [-1])
+        self.assertEqual(infered_output_dist_attrs[0].dims_mapping, [1, -1, -1])
+        self.assertEqual(infered_output_dist_attrs[1].dims_mapping, [1, -1])
+        self.assertEqual(infered_output_dist_attrs[2].dims_mapping, [1, -1])
+
+        # ijk[1, 0, -1],k[0],k[0] -->
+        # [1, 0, -1], [-1], [-1] (inputs)
+        # [1, 0, -1], [1, 0], [1, 0] (outputs)
+        self.x_spec.set_dims_mapping([1, 0, -1])
+        self.scale_spec = DistTensorSpec([0], TensorDistAttr())
+        self.bias_spec.set_dims_mapping([0])
+
+        result_dist_attrs = self.rule.infer_forward(
+            self.x_spec,
+            self.scale_spec,
+            self.bias_spec,
+            self.attrs['epsilon'],
+            self.attrs['begin_norm_axis'],
+        )
+        infered_input_dist_attrs = result_dist_attrs[0]
+        infered_output_dist_attrs = result_dist_attrs[1]
+
+        self.assertEqual(len(result_dist_attrs), 2)
+        self.assertEqual(len(infered_input_dist_attrs), 3)
+        self.assertEqual(len(infered_output_dist_attrs), 3)
+
+        self.assertEqual(infered_input_dist_attrs[0].dims_mapping, [1, 0, -1])
+        self.assertEqual(infered_input_dist_attrs[1].dims_mapping, [-1])
+        self.assertEqual(infered_input_dist_attrs[2].dims_mapping, [-1])
+        self.assertEqual(infered_output_dist_attrs[0].dims_mapping, [1, 0, -1])
+        self.assertEqual(infered_output_dist_attrs[1].dims_mapping, [1, 0])
+        self.assertEqual(infered_output_dist_attrs[2].dims_mapping, [1, 0])
+
     def test_infer_backward(self):
         # [1, -1, -1], [1, -1], [1, -1] (outputs) -->
         # [1, -1, -1], [-1], [-1], (inputs)
