[AutoParallel] Add Global mesh and sub mesh reshard function #61796

deepllz · 2024-02-18T10:29:32Z

PR types

New features

PR changes

Others

Description

Pcard-73145
新增全局mesh到pp子mesh的reshard function，支持动半场景下的pp全局算子输入

paddle-bot · 2024-02-18T10:29:36Z

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

zhiqiu · 2024-02-19T03:06:15Z

paddle/phi/core/distributed/auto_parallel/reshard/global_and_sub_pp_mesh_reshard_function.h

+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "GlobalToSubPPMeshReshardFunction"; }


sub to global

zhiqiu · 2024-02-19T03:11:52Z

paddle/phi/core/distributed/auto_parallel/reshard/global_and_sub_pp_mesh_reshard_function.cc

+  const ProcessMesh& in_process_mesh = in_dist_attr.process_mesh();
+  const ProcessMesh& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() ==


Is this condition needed?

计划暂时只支持global mesh比sub mesh维度多1的情况，目前看能够cover当前遇到的case，如果后续有更多需求的话，再补充新功能。

zhiqiu · 2024-02-19T03:13:10Z

paddle/phi/core/distributed/auto_parallel/reshard/global_and_sub_pp_mesh_reshard_function.cc

+
+bool SubPPMeshToGlobalReshardFunction::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_replicated(0));


Is this condition needed?

zhiqiu · 2024-02-19T03:28:31Z

paddle/phi/core/distributed/auto_parallel/reshard/global_and_sub_pp_mesh_reshard_function.cc

+  if (IsCurRankInMesh(in_process_mesh)) {
+    const DenseTensor& in_dense_value = in.value();
+    std::vector<int64_t>& recv_vec = send2recv_map[cur_global_rank];
+    for (int64_t recv_id : recv_vec) {


use broadcast to reduce the number of communication?

好的，我试试，感谢

经过确认，无法使用broadcast，原因是：recv节点上，没法拿到send tensor的shape（因为send tensor在另一个rank上），要使用broadcast，需要先把shape信息也发送过去。所以这里暂时保持使用send recv算子，先跑通流程。

思考了一下，后续也许可以用global shape + 切分信息推导出local shape，来避免发送shape信息

zhiqiu · 2024-02-19T03:29:47Z

test/auto_parallel/hybrid_strategy/semi_auto_parallel_3d_global_mesh_reshard.py

+        )  # 1.0
+        global_input.stop_gradient = False
+        # forward on mesh0
+        input_mesh0 = dist.reshard(


What if resshard is not called?

目前的话不调用会报错，是因为动半在计算的时候，要求算子输入tensor的mesh必须相同；后续打算针对sub mesh的情况，把这个reshard的调用隐藏起来，具体做法是在ConvertAllInputsToDistTensor调用reshard_function。

LiYuRio · 2024-02-21T03:28:41Z

paddle/phi/core/distributed/auto_parallel/reshard/CMakeLists.txt

@@ -14,4 +14,5 @@ collect_srcs(
  r_to_x_reshard_function.cc
  nd_mesh_reshard_function.cc
  same_status_reshard_function.cc
+  global_and_sub_pp_mesh_reshard_function.cc


这个reshard名字这里，是不是不引入pp概念会好一点，因为reshard是通用的，没有并行策略的概念

LiYuRio · 2024-02-21T03:29:22Z

paddle/phi/core/distributed/auto_parallel/reshard/global_and_sub_pp_mesh_reshard_function.cc

+std::vector<ProcessMesh> GetSubPPMesh(const ProcessMesh& process_mesh) {
+  const std::vector<int64_t>& shape = process_mesh.shape();
+  const std::vector<int64_t>& process_ids = process_mesh.process_ids();
+  const std::vector<std::string>& dim_names = process_mesh.dim_names();
+  int64_t total_process_num = process_ids.size();
+  int64_t sub_process_num = total_process_num / shape[0];
+  std::vector<int64_t> sub_process_mesh_shape(shape.begin() + 1, shape.end());
+  std::vector<std::string> sub_process_mesh_dim_names(dim_names.begin() + 1,
+                                                      dim_names.end());
+
+  std::vector<ProcessMesh> sub_process_meshes;
+  for (int i = 0; i < shape[0]; ++i) {
+    int64_t start_position = i * sub_process_num;
+    int64_t end_position = start_position + sub_process_num;
+    std::vector<int64_t> sub_process_ids(process_ids.begin() + start_position,
+                                         process_ids.begin() + end_position);
+
+    sub_process_meshes.emplace_back(ProcessMesh(
+        sub_process_mesh_shape, sub_process_ids, sub_process_mesh_dim_names));
+  }
+  return sub_process_meshes;
+}
+
+bool is_sub_mesh(const ProcessMesh& global_mesh, const ProcessMesh& sub_mesh) {
+  std::vector<ProcessMesh> sub_process_meshes = GetSubPPMesh(global_mesh);
+  for (const ProcessMesh& mesh : sub_process_meshes) {
+    if (mesh == sub_mesh) {
+      return true;
+    }
+  }
+  return false;
+}


这种通用的方法，或许可以放到reshard_utils里

LiYuRio · 2024-02-21T03:32:54Z

paddle/phi/core/distributed/auto_parallel/reshard/global_and_sub_pp_mesh_reshard_function.cc

+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() ==
+                            out_process_mesh.ndim() + 1);


这种能通用的支持从三维mesh到一维mesh的转换吗

暂时只支持升降一维的情况，目前能覆盖遇到的case，后面可以根据实际需求再补充功能。

LiYuRio · 2024-02-21T03:38:06Z

paddle/phi/core/distributed/auto_parallel/reshard/s_to_s_reshard_function.cc

+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.dims_mapping() !=
+                            out_dist_attr.dims_mapping());


这里是不是应该判断mesh，不应该是dims_mapping，本质上是不想让跨mesh的same status reshard进入s_to_s这个函数吗

我确认一下

经过确认，未修复之前：s_to_s_cross_mesh reshard_function已经判断了mesh不相同，所以需要加一个切分状态不相同的判断，避免进到这个函数内。

zhiqiu

LGTM

deepllz added 2 commits February 18, 2024 16:41

add global and sub mesh reshard function

fb4056e

add unittest

86e091e

zhiqiu reviewed Feb 19, 2024

View reviewed changes

LiYuRio reviewed Feb 21, 2024

View reviewed changes

deepllz added 8 commits February 23, 2024 14:44

use broadcast kernel

995d813

Merge remote-tracking branch 'upstream/develop' into global_mesh

735108c

Merge remote-tracking branch 'upstream/develop' into global_mesh

7882b4d

rm unused header

acb82f5

Merge remote-tracking branch 'upstream/develop' into global_mesh

b5abcd2

revert broadcast

69ae3b3

merge upstream

8122a71

polish code

8dd4b9d

zhiqiu approved these changes Feb 23, 2024

View reviewed changes

XieYunshen approved these changes Feb 24, 2024

View reviewed changes

GhostScreaming merged commit 735fd08 into PaddlePaddle:develop Feb 26, 2024

		RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() ==
		out_process_mesh.ndim() + 1);

		RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.dims_mapping() !=
		out_dist_attr.dims_mapping());

[AutoParallel] Add Global mesh and sub mesh reshard function #61796

[AutoParallel] Add Global mesh and sub mesh reshard function #61796

Uh oh!

Conversation

deepllz commented Feb 18, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR types

PR changes

Description

Uh oh!

paddle-bot bot commented Feb 18, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhiqiu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

deepllz commented Feb 18, 2024 •

edited

Loading