[Auto Parallel] handle partial status propagate for add_n (#69439)

waliwali777 · web-flow · commit 17f157ed827e · 2024-11-26T10:31:59.000+08:00
diff --git a/paddle/phi/infermeta/spmd_rules/add_n.cc b/paddle/phi/infermeta/spmd_rules/add_n.cc
@@ -54,14 +54,36 @@ SpmdInfo AddNInferSpmd(
   std::vector<int64_t> dims_mapping =
       GetDimsMappingForAxes(axes, axis_to_dim_map);
 
+  // Infer partial status from inputs
+  // Note: Now only supports cases where all inputs have same
+  // partial status, then output will have the same one.
+  bool all_input_partial_same = true;
+  paddle::flat_hash_map<int64_t, ReduceType> partial_status;
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    TensorDistAttr input_dist_attr_src = inputs[i].dist_attr();
+    if (input_dist_attr_src.is_partial()) {
+      auto input_partial_status = input_dist_attr_src.partial_status();
+      if (i == 0) {
+        partial_status = input_partial_status;
+      } else if (partial_status != input_partial_status) {
+        all_input_partial_same = false;
+        break;
+      }
+    } else {
+      all_input_partial_same = false;
+      break;
+    }
+  }
   std::vector<TensorDistAttr> inputs_spmd_info;
   for (const auto& input : inputs) {
     TensorDistAttr dist_attr_dst =
         CopyTensorDistAttrForOutput(input.dist_attr());
     dist_attr_dst.set_dims_mapping(dims_mapping);
+    if (all_input_partial_same) {
+      dist_attr_dst.set_partial_status(partial_status);
+    }
     inputs_spmd_info.push_back(dist_attr_dst);
   }
-  // Handle input tensor partial(TODO)
 
   return {{inputs_spmd_info}, {inputs_spmd_info[0]}};
 }
diff --git a/paddle/phi/infermeta/spmd_rules/rules.cc b/paddle/phi/infermeta/spmd_rules/rules.cc
@@ -726,5 +726,8 @@ PD_REGISTER_SPMD_RULE(pad,
 PD_REGISTER_SPMD_RULE(nonzero,
                       PD_INFER_SPMD(phi::distributed::NonZeroInferSpmd),
                       PD_INFER_SPMD(phi::distributed::NonZeroInferSpmdReverse));
+
+// add_n
+PD_REGISTER_SPMD_RULE(add_n, PD_INFER_SPMD(phi::distributed::AddNInferSpmd));
 }  // namespace distributed
 }  // namespace phi
diff --git a/test/auto_parallel/hybrid_strategy/semi_auto_llama_acc_align.py b/test/auto_parallel/hybrid_strategy/semi_auto_llama_acc_align.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import hashlib
 import os
 
 os.environ["FLAGS_enable_pir_api"] = str(1)
@@ -175,7 +174,7 @@ def run_llama(self, to_static=0):
                 model, dist_loader, criterion, optimizer, strategy=self.strategy
             )
         model.train()
-        md5_list = []
+        losses = []
         for step, inputs in enumerate(dist_loader()):
             if step >= self.run_step:
                 break
@@ -186,17 +185,17 @@ def run_llama(self, to_static=0):
                     numpy_array = np.array([])
                 else:
                     numpy_array = np.array(loss)
+                losses.append(numpy_array)
                 array_bytes = numpy_array.tobytes()
-                md5_list.append(hashlib.md5(array_bytes).hexdigest())
             else:
                 logits = model(input_ids)
                 loss = criterion(logits, labels)
+                losses.append(np.array(loss))
                 loss.backward()
                 optimizer.step()
                 optimizer.clear_grad()
-                md5_list.append(loss._local_value()._md5sum())
             lr_scheduler.step()
-        return md5_list
+        return losses
 
     def init_dist_env(self):
         order = ["dp", "pp", "mp"]
@@ -274,11 +273,7 @@ def run_dynamic(self):
                 lr_scheduler.step()
 
                 tr_loss_add = 0
-
-        numpy_array = np.array(tr_loss)
-        array_bytes = numpy_array.tobytes()
-        loss_md5 = hashlib.md5(array_bytes).hexdigest()
-        return loss_md5
+        return np.array(tr_loss)
 
     def run_dy2static(self):
         model = LlamaForCausalLMAuto(self.config)
@@ -336,27 +331,26 @@ def run_dy2static(self):
 
         if loss is not None:
             loss = np.average(loss)
-
-        numpy_array = np.array(loss)
-        array_bytes = numpy_array.tobytes()
-        loss_md5 = hashlib.md5(array_bytes).hexdigest()
-        return loss_md5
+        return np.array(loss)
 
     def run_test_cases(self):
         self.init_dist_env()
         if self.gradient_accumulation_steps > 1:
-            dy_loss_md5 = self.run_dynamic()
+            dy_losses = self.run_dynamic()
             self.init_dist_env()
-            st_loss_md5 = self.run_dy2static()
+            st_losses = self.run_dy2static()
             if int(dist.get_rank()) in [2, 3, 6, 7]:
-                assert dy_loss_md5 == st_loss_md5
+                np.testing.assert_allclose(dy_losses, st_losses, atol=1e-7)
+
         else:
-            dy_loss_md5 = self.run_llama(to_static=0)
+            dy_losses = self.run_llama(to_static=0)
             self.init_dist_env()
-            st_loss_md5 = self.run_llama(to_static=1)
-            assert len(dy_loss_md5) == len(st_loss_md5)
-            for idx in range(len(dy_loss_md5)):
-                assert dy_loss_md5[idx] == st_loss_md5[idx]
+            st_losses = self.run_llama(to_static=1)
+            assert len(dy_losses) == len(st_losses)
+            for idx in range(len(dy_losses)):
+                np.testing.assert_allclose(
+                    dy_losses[idx], st_losses[idx], atol=1e-7
+                )
 
 
 if __name__ == '__main__':
diff --git a/test/auto_parallel/spmd_rules/CMakeLists.txt b/test/auto_parallel/spmd_rules/CMakeLists.txt
@@ -43,6 +43,9 @@ if(WITH_DISTRIBUTE)
                   test_fused_dropout_add_rule)
   py_test_modules(test_logsumexp_rule MODULES test_logsumexp_rule)
   py_test_modules(test_nonzero_rule MODULES test_nonzero_rule)
+  if(NOT WITH_ROCM)
+    py_test_modules(test_add_n_rule MODULES test_add_n_rule)
+  endif()
   # End of unittests WITH single card WITHOUT timeout
 
 endif()
diff --git a/test/auto_parallel/spmd_rules/test_add_n_rule.py b/test/auto_parallel/spmd_rules/test_add_n_rule.py
@@ -0,0 +1,164 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+from paddle.distributed.auto_parallel.static.dist_attribute import (
+    DistTensorSpec,
+    TensorDistAttr,
+)
+from paddle.distributed.fleet import auto
+from paddle.framework import core
+
+
+class TestAddNSPMDRule(unittest.TestCase):
+    """
+    Unit tests for add_n spmd rule.
+    """
+
+    def setUp(self):
+        self.rule1 = core.get_phi_spmd_rule("add_n")
+        process_mesh = auto.ProcessMesh(mesh=[[0, 1], [2, 3]])
+        self.x_shape = [16, 16, 16]
+        self.x_tensor_dist_attr = TensorDistAttr()
+        self.x_tensor_dist_attr.process_mesh = process_mesh
+
+        self.y_shape = [16, 16, 16]
+        self.y_tensor_dist_attr = TensorDistAttr()
+        self.y_tensor_dist_attr.process_mesh = process_mesh
+
+        self.output_shape = [16, 16, 16]
+        self.output_tensor_dist_attr = TensorDistAttr()
+        self.output_tensor_dist_attr.process_mesh = process_mesh
+
+    def test_infer_forward(self):
+        # [0, -1, -1], [-1, -1, -1] (x,y) -->
+        # [0, -1, -1], [0, -1, -1] (x, y)
+        # [0, -1, -1] (output)
+        self.x_dist_tensor_spec = DistTensorSpec(
+            self.x_shape, self.x_tensor_dist_attr
+        )
+        self.x_dist_tensor_spec.set_dims_mapping([0, -1, -1])
+        self.y_dist_tensor_spec = DistTensorSpec(
+            self.y_shape, self.y_tensor_dist_attr
+        )
+        self.y_dist_tensor_spec.set_dims_mapping([0, -1, -1])
+
+        infered_dist_attr = self.rule1.infer_forward(
+            [self.x_dist_tensor_spec, self.y_dist_tensor_spec]
+        )
+
+        self.assertEqual(len(infered_dist_attr), 2)
+        infered_input_dist_attr = infered_dist_attr[0]
+        infered_output_dist_attr = infered_dist_attr[1]
+
+        self.assertEqual(len(infered_input_dist_attr), 1)
+        self.assertEqual(len(infered_input_dist_attr[0]), 2)
+        self.assertEqual(len(infered_output_dist_attr), 1)
+
+        self.assertEqual(
+            infered_input_dist_attr[0][0].dims_mapping, [0, -1, -1]
+        )
+        self.assertEqual(
+            infered_input_dist_attr[0][1].dims_mapping, [0, -1, -1]
+        )
+        self.assertEqual(infered_output_dist_attr[0].dims_mapping, [0, -1, -1])
+
+        # [0, -1, -1], [-1, -1, -1] (x, y) partial_dim=[1] -->
+        # [0, -1, -1], [0, -1, -1]  (x, y) partial_dim=[1]
+        # [0, -1, -1] (output) partial_dim=[1]
+        self.x_tensor_dist_attr._set_partial_dims([1])
+        self.x_dist_tensor_spec = DistTensorSpec(
+            self.x_shape, self.x_tensor_dist_attr
+        )
+        self.x_dist_tensor_spec.set_dims_mapping([0, -1, -1])
+        self.y_tensor_dist_attr._set_partial_dims([1])
+        self.y_dist_tensor_spec = DistTensorSpec(
+            self.y_shape, self.y_tensor_dist_attr
+        )
+        self.y_dist_tensor_spec.set_dims_mapping([-1, -1, -1])
+
+        infered_dist_attr = self.rule1.infer_forward(
+            [self.x_dist_tensor_spec, self.y_dist_tensor_spec]
+        )
+
+        self.assertEqual(len(infered_dist_attr), 2)
+        infered_input_dist_attr = infered_dist_attr[0]
+        infered_output_dist_attr = infered_dist_attr[1]
+
+        self.assertEqual(len(infered_input_dist_attr), 1)
+        self.assertEqual(len(infered_input_dist_attr[0]), 2)
+        self.assertEqual(len(infered_output_dist_attr), 1)
+
+        self.assertEqual(
+            infered_input_dist_attr[0][0].dims_mapping, [0, -1, -1]
+        )
+        self.assertEqual(infered_input_dist_attr[0][0]._is_partial(), True)
+        self.assertEqual(infered_input_dist_attr[0][0]._partial_dims(), {1})
+
+        self.assertEqual(
+            infered_input_dist_attr[0][1].dims_mapping, [0, -1, -1]
+        )
+        self.assertEqual(infered_input_dist_attr[0][1]._is_partial(), True)
+        self.assertEqual(infered_input_dist_attr[0][1]._partial_dims(), {1})
+
+        self.assertEqual(infered_output_dist_attr[0].dims_mapping, [0, -1, -1])
+        self.assertEqual(infered_output_dist_attr[0]._is_partial(), True)
+        self.assertEqual(infered_output_dist_attr[0]._partial_dims(), {1})
+
+        # [0, -1, -1] partial_dim=[0], [-1, -1, -1]partial_dim=[1] (x,y)  -->
+        # [0, -1, -1], [0, -1, -1] (x, y)
+        # [0, -1, -1] (output)
+        self.x_tensor_dist_attr._clean_partial_dims([1])
+        self.x_tensor_dist_attr._set_partial_dims([0])
+        self.x_dist_tensor_spec = DistTensorSpec(
+            self.x_shape, self.x_tensor_dist_attr
+        )
+        self.x_dist_tensor_spec.set_dims_mapping([0, -1, -1])
+
+        self.y_tensor_dist_attr._clean_partial_dims([1])
+        self.y_tensor_dist_attr._set_partial_dims([1])
+        self.y_dist_tensor_spec = DistTensorSpec(
+            self.y_shape, self.y_tensor_dist_attr
+        )
+        self.y_dist_tensor_spec.set_dims_mapping([-1, -1, -1])
+
+        infered_dist_attr = self.rule1.infer_forward(
+            [self.x_dist_tensor_spec, self.y_dist_tensor_spec]
+        )
+
+        self.assertEqual(len(infered_dist_attr), 2)
+        infered_input_dist_attr = infered_dist_attr[0]
+        infered_output_dist_attr = infered_dist_attr[1]
+
+        self.assertEqual(len(infered_input_dist_attr), 1)
+        self.assertEqual(len(infered_input_dist_attr[0]), 2)
+        self.assertEqual(len(infered_output_dist_attr), 1)
+
+        self.assertEqual(
+            infered_input_dist_attr[0][0].dims_mapping, [0, -1, -1]
+        )
+        self.assertEqual(infered_input_dist_attr[0][0]._is_partial(), False)
+
+        self.assertEqual(
+            infered_input_dist_attr[0][1].dims_mapping, [0, -1, -1]
+        )
+        self.assertEqual(infered_input_dist_attr[0][1]._is_partial(), False)
+
+        self.assertEqual(infered_output_dist_attr[0].dims_mapping, [0, -1, -1])
+        self.assertEqual(infered_output_dist_attr[0]._is_partial(), False)
+
+
+if __name__ == "__main__":
+    unittest.main()
