[Auto Parallel] Add spmd rule No.4、13 for (batch_norm,sync_batch_norm) and their backward ops. #72918

Glencsa · 2025-05-24T12:49:09Z

PR Category

Auto Parallel

PR Types

New features

Description

【开源任务】算子切分推导规则开发，支持更多模型使用自动并行，简化更多用户的分布式开发成本。
No.4 batch_norm
No.13 sync_batch_norm
将除了做batch_norm以外的维度全部强制为Replicated

… develop

paddle-bot · 2025-05-24T12:49:14Z

你的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.

jeff41404 · 2025-06-10T02:35:41Z

test/cpp/auto_parallel/spmd_rule_test.cc

@@ -2614,7 +2785,7 @@ TEST(Topk, Ctor) {

  // test forward
  // axis = 1
-  // [0, 1, -1] -> [0, -1, -1], [0, -1, -1]
+  // [0, -1, -1, 1],[-1],[-1],[-1],[-1] ->[-1 , -1, -1, 1],[1],[1],[1],[1],[1]


should not modify this annotation?

Thanks for your notice, I will revert my changes in next commit.

jeff41404 · 2025-06-10T02:46:37Z

paddle/phi/infermeta/spmd_rules/batch_norm.cc

+                            const std::string data_format,
+                            const bool use_global_stats,
+                            const bool trainable_statistics) {
+  return BatchNormInferSpmdBase(x, mean, variance, scale, bias);


do we need parameter of data_format in BatchNormInferSpmdBase?
if user pass data_format="NHWC" or "NLC", will right?

Thanks, I have fit all the data_format status in the new commit.

jeff41404 · 2025-06-10T02:48:57Z

paddle/phi/infermeta/spmd_rules/batch_norm.cc

+                                const bool is_test,
+                                const bool use_global_stats,
+                                const bool trainable_statistics) {
+  return BatchNormGradInferSpmdBase(x,


same issue of data_format as in BatchNormInferSpmdBase

Thanks, I have fit all the data_format status in the new commit.

jeff41404 · 2025-06-10T02:56:58Z

paddle/phi/ops/yaml/ops.yaml

@@ -5056,6 +5056,7 @@
  output : Tensor(out), Tensor(mean_out), Tensor(variance_out), Tensor(saved_mean), Tensor(saved_variance), Tensor(reserve_space)
  infer_meta :
    func : BatchNormInferMeta
+    spmd_rule : SyncBatchNormInferSpmd


the operator of sync_batch_norm_ is used for manual parallelism, and the implementation of operator includes communication, not just a calculation operator. should have spmd rule?

Thanks, I think you are right, the operator of sync_batch_norm_ cause different GPUs have different batch, and their mean and variance on device should be communication, the tensor can not be sharded. I will remove the spmd_rule for sync_batch_norm_ in next commit.

jeff41404 · 2025-06-10T03:30:20Z

paddle/phi/infermeta/spmd_rules/batch_norm.cc

+  VLOG(4) << "Einsum Notation: " << x_axes << "," << mean_axes << ","
+          << variance_axes << "," << scale_axes << "," << bias_axes << "-->"
+          << out_axes << "," << mean_axes << "," << variance_axes;
+  VLOG(4) << "X"
+          << " shape: [" << str_join(x_shape) << "] "
+          << "src_dims_mapping: [" << str_join(x_dist_attr_src.dims_mapping())
+          << "] "
+          << "dst_dims_mapping: [" << str_join(x_dims_mapping) << "]";
+  VLOG(4) << "Mean"
+          << " shape: [" << str_join(mean_shape) << "] "
+          << "src_dims_mapping: [" << str_join(mean_dims_mapping) << "] "
+          << "dst_dims_mapping: ["
+          << str_join(mean_dist_attr_dst.dims_mapping()) << "]";
+  VLOG(4) << "Variance"
+          << " shape: [" << str_join(variance_shape) << "] "
+          << "src_dims_mapping: [" << str_join(variance_dims_mapping) << "] "
+          << "dst_dims_mapping: ["
+          << str_join(variance_dist_attr_dst.dims_mapping()) << "]";
+  VLOG(4) << "Scale"
+          << " shape: [" << str_join(scale_shape) << "] "
+          << "src_dims_mapping: [" << str_join(scale_dims_mapping) << "] "
+          << "dst_dims_mapping: ["
+          << str_join(scale_dist_attr_dst.dims_mapping()) << "]";
+  VLOG(4) << "Bias"
+          << " shape: [" << str_join(bias_shape) << "] "
+          << "src_dims_mapping: [" << str_join(bias_dims_mapping) << "] "
+          << "dst_dims_mapping: ["
+          << str_join(bias_dist_attr_dst.dims_mapping()) << "]";
+  VLOG(4) << "Out dims mapping: [" << str_join(out_dist_attr.dims_mapping())
+          << "]";
+  VLOG(4) << "Mean_out dims mapping: ["
+          << str_join(mean_dist_attr.dims_mapping()) << "]";
+  VLOG(4) << "Variance_out dims mapping: ["
+          << str_join(variance_dist_attr.dims_mapping()) << "]";
+  VLOG(4) << "Saved_mean dims mapping: ["
+          << str_join(mean_dist_attr.dims_mapping()) << "]";
+  VLOG(4) << "Saved_variance dims mapping: ["
+          << str_join(variance_dist_attr.dims_mapping()) << "]";
+  VLOG(4) << "Reserve_space dims mapping: ["
+          << str_join(reserve_space_dist_attr.dims_mapping()) << "]";
+  VLOG(4) << std::endl;


shall we use macro LOG_SPMD_INPUT or LOG_SPMD_OUTPUT to simplify log code

Thanks, I will use it to simplify log code in next commit.

jeff41404 · 2025-06-10T03:31:35Z

paddle/phi/infermeta/spmd_rules/batch_norm.cc

+  VLOG(4) << "Einsum Notation: " << x_axes << scale_axes << "," << bias_axes
+          << "," << mean_out_axes << "," << variance_out_axes << ","
+          << saved_mean_axes << "," << saved_variance_axes << ","
+          << "-->" << reserve_space_axes << "," << out_grad_axes;
+  VLOG(4) << "X"
+          << " shape: [" << str_join(x_shape) << "] "
+          << "src_dims_mapping: [" << str_join(x_dist_attr_src.dims_mapping())
+          << "] "
+          << "dst_dims_mapping: [" << str_join(x_dims_mapping) << "]";
+  VLOG(4) << "Mean_out"
+          << " shape: [" << str_join(mean_out_shape) << "] "
+          << "src_dims_mapping: ["
+          << str_join(mean_out.dist_attr().dims_mapping()) << "] "
+          << "dst_dims_mapping: [" << str_join(mean_out_attr_dst.dims_mapping())
+          << "]";
+  VLOG(4) << "Variance_out"
+          << " shape: [" << str_join(variance_out_shape) << "] "
+          << "src_dims_mapping: ["
+          << str_join(variance_out.dist_attr().dims_mapping()) << "] "
+          << "dst_dims_mapping: ["
+          << str_join(variance_out_attr_dst.dims_mapping()) << "]";
+  VLOG(4) << "Scale"
+          << " shape: [" << str_join(scale_shape) << "] "
+          << "src_dims_mapping: [" << str_join(scale.dist_attr().dims_mapping())
+          << "] "
+          << "dst_dims_mapping: [" << str_join(scale_attr_dst.dims_mapping())
+          << "]";
+  VLOG(4) << "Bias"
+          << " shape: [" << str_join(bias_shape) << "] "
+          << "src_dims_mapping: [" << str_join(bias.dist_attr().dims_mapping())
+          << "] "
+          << "dst_dims_mapping: [" << str_join(bias_attr_dst.dims_mapping())
+          << "]";


shall we use macro LOG_SPMD_INPUT or LOG_SPMD_OUTPUT to simplify log code

Thanks, I will use it to simplify log code in next commit.

jeff41404

LGTM

Glencsa added 10 commits April 10, 2025 19:31

add unary ops which have spmd_rule but not add in yaml file.

8dfce3a

Merge branch 'spmd_test' into develop

1d129c2

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

b9c9e6a

… develop

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

f24c883

… develop

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

746356c

… develop

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

efc91c5

… develop

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

2109cf9

… develop

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

773fda6

… develop

Add spmd_rule for batch_norm ops.

044b45b

Add spmd_rule for batch_norm and batch_norm_grad.

bc53066

Glencsa requested review from LiYuRio, ForFishes and zhiqiu as code owners May 24, 2025 12:49

paddle-bot bot added the contributor External developers label May 24, 2025

Glencsa added 3 commits May 24, 2025 21:54

fix bug

99b2c70

fix bug.

e8a47a6

fix bug.

da1c541

luotao1 changed the title ~~[Auto Parallel] Add spmd rule for batch_norm and batch_norm_grad ops.~~ [Auto Parallel] Add spmd rule No.4 for batch_norm and batch_norm_grad ops. May 26, 2025

luotao1 added the HappyOpenSource Pro 进阶版快乐开源活动，更具挑战性的任务 label May 26, 2025

luotao1 assigned luotao1 and jeff41404 May 26, 2025

luotao1 mentioned this pull request May 26, 2025

【开源任务】算子切分推导规则开发，支持更多模型使用自动并行，简化更多用户的分布式开发成本 #72415

Closed

Glencsa added 4 commits May 31, 2025 16:14

add spmd_rule for sync_natch_norm

19683c0

add spmd_rule for sync_natch_norm

7b8aa5f

fix bug.

134d4f8

fix bug.

7fa09bd

Glencsa changed the title ~~[Auto Parallel] Add spmd rule No.4 for batch_norm and batch_norm_grad ops.~~ [Auto Parallel] Add spmd rule for No.4(batch_norm, batch_norm_grad) and No.13(sync_batch_norm,sync_batch_norm_grad) ops. Jun 1, 2025

Glencsa added 2 commits June 4, 2025 10:13

Add partial status.

ccc360e

fix ci bug.

4d220eb

fix CI bug.

874f5ae

Glencsa changed the title ~~[Auto Parallel] Add spmd rule for No.4(batch_norm, batch_norm_grad) and No.13(sync_batch_norm,sync_batch_norm_grad) ops.~~ [Auto Parallel] Add spmd rule for No.4、13 for (batch_norm,sync_batch_norm) and their backward ops. Jun 9, 2025

Glencsa changed the title ~~[Auto Parallel] Add spmd rule for No.4、13 for (batch_norm,sync_batch_norm) and their backward ops.~~ [Auto Parallel] Add spmd rule No.4、13 for (batch_norm,sync_batch_norm) and their backward ops. Jun 9, 2025

jeff41404 reviewed Jun 10, 2025

View reviewed changes

Glencsa added 2 commits June 13, 2025 20:05

apply review.

6ebf375

fix bug.

8cea62d

jeff41404 approved these changes Jun 16, 2025

View reviewed changes

luotao1 merged commit 4aee08b into PaddlePaddle:develop Jun 16, 2025
49 of 50 checks passed

Glencsa deleted the batch_norm_spmd_rule branch June 30, 2025 07:45

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Auto Parallel] Add spmd rule No.4、13 for (batch_norm,sync_batch_norm) and their backward ops. #72918

[Auto Parallel] Add spmd rule No.4、13 for (batch_norm,sync_batch_norm) and their backward ops. #72918

Uh oh!

Glencsa commented May 24, 2025 •

edited

Loading

Uh oh!

paddle-bot bot commented May 24, 2025

Uh oh!

jeff41404 Jun 10, 2025

Uh oh!

Glencsa Jun 12, 2025

Uh oh!

jeff41404 Jun 10, 2025

Uh oh!

Glencsa Jun 13, 2025

Uh oh!

jeff41404 Jun 10, 2025

Uh oh!

Glencsa Jun 13, 2025

Uh oh!

jeff41404 Jun 10, 2025 •

edited

Loading

Uh oh!

Glencsa Jun 13, 2025

Uh oh!

jeff41404 Jun 10, 2025

Uh oh!

Glencsa Jun 13, 2025

Uh oh!

jeff41404 Jun 10, 2025

Uh oh!

Glencsa Jun 13, 2025

Uh oh!

jeff41404 left a comment

Uh oh!

Uh oh!

Uh oh!

[Auto Parallel] Add spmd rule No.4、13 for (batch_norm,sync_batch_norm) and their backward ops. #72918

[Auto Parallel] Add spmd rule No.4、13 for (batch_norm,sync_batch_norm) and their backward ops. #72918

Uh oh!

Conversation

Glencsa commented May 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Category

PR Types

Description

Uh oh!

paddle-bot bot commented May 24, 2025

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!

jeff41404 Jun 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

jeff41404 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Glencsa commented May 24, 2025 •

edited

Loading

jeff41404 Jun 10, 2025 •

edited

Loading