[AutoParallel]:fix fused_dropout_add bug (#68736)

blacksheep-Aristotle · web-flow · commit 90db9c67e6fa · 2024-10-18T14:10:00.000+08:00
* [AutoParallel]:fix fused_dropout_add bug

* [AutoParallel]:fix fused_dropout_add bug
diff --git a/paddle/phi/infermeta/spmd_rules/fused_dropout_add.cc b/paddle/phi/infermeta/spmd_rules/fused_dropout_add.cc
@@ -77,7 +77,9 @@ SpmdInfo FusedDropoutAddSpmd(const DistMetaTensor& x,
                              const std::string& mode,
                              int seed,
                              bool fix_seed) {
-  return FusedDropoutAddSpmdBase(x, y);
+  auto dropout_info = FusedDropoutAddSpmdBase(x, y);
+  dropout_info.first.push_back(seed_tensor.dist_attr());
+  return dropout_info;
 }
 
 SpmdInfo FusedDropoutAddSpmdReverse(const DistMetaTensor& x,
diff --git a/paddle/phi/kernels/fusion/gpu/fused_dropout_add_grad_kernel.cu b/paddle/phi/kernels/fusion/gpu/fused_dropout_add_grad_kernel.cu
@@ -263,5 +263,5 @@ PD_REGISTER_KERNEL(fused_dropout_add_grad,
                    double,
                    phi::dtype::bfloat16,
                    phi::dtype::float16) {
-  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);  // seed_offset
+  kernel->InputAt(0).SetBackend(phi::Backend::CPU);  // seed_offset
 }

Original file line number	Diff line number	Diff line change
`@@ -263,5 +263,5 @@ PD_REGISTER_KERNEL(fused_dropout_add_grad,`
`263`	`263`	`double,`
`264`	`264`	`phi::dtype::bfloat16,`
`265`	`265`	`phi::dtype::float16) {`
`266`		`- kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND); // seed_offset`
	`266`	`+ kernel->InputAt(0).SetBackend(phi::Backend::CPU); // seed_offset`
`267`	`267`	`}`