[AutoParallel] support spmd for squared_l2_norm (#61460)

liym27 · web-flow · commit 94fffec08b69 · 2024-02-05T11:25:21.000+08:00
* [AutoParallel] support spmd for squared_l2_norm
diff --git a/paddle/phi/infermeta/spmd_rules/rules.cc b/paddle/phi/infermeta/spmd_rules/rules.cc
@@ -490,6 +490,10 @@ PD_REGISTER_SPMD_RULE(
     reduce_sum,  // static
     PD_INFER_SPMD(phi::distributed::ReductionInferSpmd),
     PD_INFER_SPMD(phi::distributed::ReductionInferSpmdReverse));
+PD_REGISTER_SPMD_RULE(
+    squared_l2_norm,
+    PD_INFER_SPMD(phi::distributed::ReductionInferSpmd),
+    PD_INFER_SPMD(phi::distributed::ReductionInferSpmdReverse));
 
 // layer_norm
 PD_REGISTER_SPMD_RULE(
diff --git a/test/auto_parallel/spmd_rules/test_reduction_rule.py b/test/auto_parallel/spmd_rules/test_reduction_rule.py
@@ -28,8 +28,12 @@ class TestReductionSPMDRule(unittest.TestCase):
     Unit tests for reduction spmd rule.
     """
 
+    def config(self):
+        self.kernel = "max"
+
     def setUp(self):
-        self.rule = core.get_phi_spmd_rule("max")
+        self.config()
+        self.rule = core.get_phi_spmd_rule(self.kernel)
 
         x_shape = [64, 32]
         process_mesh = auto.ProcessMesh(mesh=[0, 1, 2, 3])
@@ -461,5 +465,14 @@ def test_backward_multi_mesh_dim_parital(self):
         self.assertEqual(infered_input_dist_attrs[0]._is_partial(), False)
 
 
+class TestSquaredL2NormSPMDRule(TestReductionSPMDRule):
+    """
+    Unit tests for squared_l2_norm spmd rule.
+    """
+
+    def config(self):
+        self.kernel = "squared_l2_norm"
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/legacy_test/test_squared_l2_norm_op.py b/test/legacy_test/test_squared_l2_norm_op.py
@@ -19,6 +19,7 @@
 from op_test import OpTest
 
 import paddle
+import paddle.distributed as dist
 from paddle import _C_ops, _legacy_C_ops
 from paddle.framework import in_dynamic_mode
 
@@ -73,12 +74,17 @@ def init_test_case(self):
 class TestL2LossOp(OpTest):
     """Test squared_l2_norm"""
 
+    def config(self):
+        self.x_shape = (13, 19)
+        self.check_auto_parallel = False
+
     def setUp(self):
+        self.config()
         self.python_api = test_squared_l2_norm
         self.op_type = "squared_l2_norm"
         self.max_relative_error = 0.05
 
-        X = np.random.uniform(-1, 1, (13, 19)).astype("float32")
+        X = np.random.uniform(-1, 1, self.x_shape).astype("float32")
         X[np.abs(X) < self.max_relative_error] = 0.1
         self.inputs = {'X': X}
         self.outputs = {'Out': np.array([np.square(LA.norm(X))])}
@@ -91,9 +97,37 @@ def test_check_grad(self):
             ['X'],
             'Out',
             max_relative_error=self.max_relative_error,
+            check_auto_parallel=self.check_auto_parallel,
         )
 
 
+class TestSquaredL2NormAutoParallel_1(TestL2LossOp):
+    def config(self):
+        self.x_shape = (14, 18)
+        self.check_auto_parallel = True
+        self.placements = {
+            'X': [dist.Replicate()],
+        }
+
+
+class TestSquaredL2NormAutoParallel_2(TestL2LossOp):
+    def config(self):
+        self.x_shape = (14, 18)
+        self.check_auto_parallel = True
+        self.placements = {
+            'X': [dist.Shard(0)],
+        }
+
+
+class TestSquaredL2NormAutoParallel_3(TestL2LossOp):
+    def config(self):
+        self.x_shape = (14, 18)
+        self.check_auto_parallel = True
+        self.placements = {
+            'X': [dist.Shard(1)],
+        }
+
+
 class TestL2LossDeterministic(unittest.TestCase):
     def check_place(self, place):
         with paddle.base.dygraph.guard(place):
