【Auto-Parallel】optimize nd_mesh_reshard_function in dynamic-auto-parallel and fix transpose (#71969)

Function-Samuel · web-flow · commit 715aab2907fa · 2025-04-07T22:53:35.000+08:00
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/nd_mesh_reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/nd_mesh_reshard_function.cc
@@ -148,7 +148,8 @@ void SameNdMeshReshardFunction::Eval(phi::DeviceContext* dev_ctx,
   // 2. change all the shard status to replicated status
   for (int64_t i = first_diff_axis; i >= 0; --i) {
     int64_t in_mesh_axis = out->dist_attr().dims_mapping()[i];
-    if (in_mesh_axis != -1) {
+    int64_t out_mesh_axis = out_dist_attr_orig.dims_mapping()[i];
+    if (in_mesh_axis != -1 && in_mesh_axis != out_mesh_axis) {
       VLOG(3) << "Step2: in_mesh axis " << in_mesh_axis;
       // 2.1 Calculate the dist_attr after this transform
       TensorDistAttr real_out_dist_attr(out->dist_attr());
@@ -223,8 +224,9 @@ void SameNdMeshReshardFunction::Eval(phi::DeviceContext* dev_ctx,
 
   // 4. Change replicated/partial to shard
   for (int64_t i = first_diff_axis; i >= 0; --i) {
+    int64_t in_mesh_axis = out->dist_attr().dims_mapping()[i];
     int64_t out_mesh_axis = out_dist_attr_orig.dims_mapping()[i];
-    if (out_mesh_axis != -1) {
+    if (in_mesh_axis != out_mesh_axis) {
       const auto& in_partial_status = out->dist_attr().partial_status();
       bool is_partial = in_partial_status.count(out_mesh_axis) != 0;
 
diff --git a/paddle/phi/kernels/transpose_kernel.h b/paddle/phi/kernels/transpose_kernel.h
@@ -44,7 +44,7 @@ void Transpose(const Context& dev_ctx,
 
   // do not call TransposeStridedKernel, because some other kernels call
   // Transpose directly
-  if (x.has_allocation()) {
+  if (x.has_allocation() && x.capacity() > 0) {
     TransposeKernel<T, Context>(dev_ctx, x, axis, dense_out);
   }
 }

Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,7 @@ void Transpose(const Context& dev_ctx,`
`44`	`44`
`45`	`45`	`// do not call TransposeStridedKernel, because some other kernels call`
`46`	`46`	`// Transpose directly`
`47`		`- if (x.has_allocation()) {`
	`47`	`+ if (x.has_allocation() && x.capacity() > 0) {`
`48`	`48`	`TransposeKernel<T, Context>(dev_ctx, x, axis, dense_out);`
`49`	`49`	`}`
`50`	`50`	`}`