[AutoParallel] Support auto parallel distributeed branch code generation (#58819)

GhostScreaming · web-flow · commit 64882b888720 · 2023-11-09T21:00:21.000-06:00
for intermediate dygraph api, like layer_norm_intermediate.
diff --git a/paddle/phi/api/lib/CMakeLists.txt b/paddle/phi/api/lib/CMakeLists.txt
@@ -200,12 +200,22 @@ execute_process(
     --api_source_path ${strings_api_source_file_tmp})
 
 # generate dygraph(intermediate) api
-execute_process(
-  COMMAND
-    ${PYTHON_EXECUTABLE} ${im_api_gen_file} --api_yaml_path ${api_yaml_file}
-    ${legacy_api_yaml_file} --sparse_api_yaml_path ${sparse_api_yaml_file}
-    --dygraph_api_header_path ${dygraph_api_header_file_tmp}
-    --dygraph_api_source_path ${dygraph_api_source_file_tmp})
+if(WITH_DISTRIBUTE)
+  execute_process(
+    COMMAND
+      ${PYTHON_EXECUTABLE} ${im_api_gen_file} --api_yaml_path ${api_yaml_file}
+      ${legacy_api_yaml_file} --sparse_api_yaml_path ${sparse_api_yaml_file}
+      --dygraph_api_header_path ${dygraph_api_header_file_tmp}
+      --dygraph_api_source_path ${dygraph_api_source_file_tmp}
+      --gen_dist_branch)
+else()
+  execute_process(
+    COMMAND
+      ${PYTHON_EXECUTABLE} ${im_api_gen_file} --api_yaml_path ${api_yaml_file}
+      ${legacy_api_yaml_file} --sparse_api_yaml_path ${sparse_api_yaml_file}
+      --dygraph_api_header_path ${dygraph_api_header_file_tmp}
+      --dygraph_api_source_path ${dygraph_api_source_file_tmp})
+endif()
 
 # generate tensor and tensor operants file
 message("create or copy auto-geneated tensor files")
diff --git a/paddle/phi/api/yaml/generator/intermediate_api_gen.py b/paddle/phi/api/yaml/generator/intermediate_api_gen.py
@@ -16,6 +16,7 @@
 
 import yaml
 from api_gen import ForwardAPI
+from dist_api_gen import DistForwardAPI
 from sparse_api_gen import SparseAPI
 
 
@@ -56,6 +57,11 @@ def source_include(header_file_path):
 #include "paddle/phi/api/profiler/event_tracing.h"
 #include "paddle/phi/api/profiler/supplement_tracing.h"
 
+#ifdef PADDLE_WITH_DISTRIBUTE
+#include "paddle/phi/infermeta/spmd_rules/rules.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#endif
+
 PD_DECLARE_int32(low_precision_op_list);
 """
 
@@ -91,6 +97,7 @@ def generate_intermediate_api(
     sparse_api_yaml_path,
     dygraph_header_file_path,
     dygraph_source_file_path,
+    gen_dist_branch,
 ):
     dygraph_header_file = open(dygraph_header_file_path, 'w')
     dygraph_source_file = open(dygraph_source_file_path, 'w')
@@ -114,7 +121,7 @@ def generate_intermediate_api(
                 apis.extend(api_list)
 
     for api in apis:
-        foward_api = ForwardAPI(api)
+        foward_api = DistForwardAPI(api) if gen_dist_branch else ForwardAPI(api)
         if foward_api.is_dygraph_api:
             dygraph_header_file.write(foward_api.gene_api_declaration())
             dygraph_source_file.write(foward_api.gene_api_code())
@@ -170,18 +177,27 @@ def main():
         default='paddle/phi/api/lib/dygraph_api.cc',
     )
 
+    parser.add_argument(
+        '--gen_dist_branch',
+        help='whether generate distributed branch code',
+        dest='gen_dist_branch',
+        action='store_true',
+    )
+
     options = parser.parse_args()
 
     api_yaml_path = options.api_yaml_path
     sparse_api_yaml_path = options.sparse_api_yaml_path
     dygraph_header_file_path = options.dygraph_api_header_path
     dygraph_source_file_path = options.dygraph_api_source_path
+    gen_dist_branch = options.gen_dist_branch
 
     generate_intermediate_api(
         api_yaml_path,
         sparse_api_yaml_path,
         dygraph_header_file_path,
         dygraph_source_file_path,
+        gen_dist_branch,
     )
 
 
diff --git a/test/auto_parallel/test_api_dist_branch.py b/test/auto_parallel/test_api_dist_branch.py
@@ -397,6 +397,32 @@ def test_merged_adam_for_dist_tensor(self):
                 local_master_param_out[i], dist_master_param_out[i]
             )
 
+    # intermediate dygraph api test
+    def test_layer_norm_for_intermediate_dist_tensor(self):
+        x = np.random.random((2, 3, 10, 10)).astype("float32")
+        weight = np.random.random(300).astype("float32")
+        bias = np.random.random(300).astype("float32")
+
+        local_x, dist_x = self.create_local_and_dist_tensor_pair(x)
+        local_weight, dist_weight = self.create_local_and_dist_tensor_pair(
+            weight
+        )
+        local_bias, dist_bias = self.create_local_and_dist_tensor_pair(bias)
+
+        local_out = paddle.nn.functional.layer_norm(
+            local_x,
+            local_x.shape[1:],
+            local_weight,
+            local_bias,
+        )
+        dist_out = paddle.nn.functional.layer_norm(
+            dist_x,
+            dist_x.shape[1:],
+            dist_weight,
+            dist_bias,
+        )
+        self.check_tensor_eq(local_out, dist_out)
+
 
 if __name__ == "__main__":
     unittest.main()
