fix bugs

PaddlePaddle · Oct 5, 2023 · 9e9140f · 9e9140f
1 parent 4aff5eb
commit 9e9140f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/paddle/phi/infermeta/spmd_rules/rules.h b/paddle/phi/infermeta/spmd_rules/rules.h
@@ -65,6 +65,12 @@ PD_REGISTER_SPMD_RULE(
     PD_INFER_SPMD(phi::distributed::ElementwiseBinaryInferSpmd),
     PD_INFER_SPMD(phi::distributed::ElementwiseBinaryInferSpmdReverse));
 
+// default_data_parallel rule
+PD_REGISTER_SPMD_RULE(
+    default_data_parallel,
+    PD_INFER_SPMD(phi::distributed::DefaultDataParallelInferSpmd),
+    PD_INFER_SPMD(phi::distributed::DefaultDataParallelInferSpmdReverse));
+
 // unsqueeze rule
 PD_REGISTER_SPMD_RULE(
     unsqueeze,

diff --git a/test/auto_parallel/spmd_rules/test_default_data_parallel_rule.py b/test/auto_parallel/spmd_rules/test_default_data_parallel_rule.py
@@ -26,7 +26,7 @@ class TestDefaultDataParallelSPMDRule(unittest.TestCase):
     def setUp(self):
         # After replaced all spmd rules by phi impl, we can recover the
         # api name to `get_spmd_rule`
-        self.rule = core.get_phi_spmd_rule("unsqueeze")
+        self.rule = core.get_phi_spmd_rule("default_data_parallel")
 
         x_shape = [10, 10, 32, 48]
         y_shape = [32, 48]