From d53e567a242419e5e08db67a4def15cd05855cc4 Mon Sep 17 00:00:00 2001
From: ShenLiang <1422485404@qq.com>
Date: Fri, 10 Sep 2021 11:25:18 +0800
Subject: [PATCH] fix bug of recompute in hybridparallel (#35588)

---
 paddle/fluid/operators/flatten_op.cu.cc                      | 5 +++++
 .../paddle/distributed/fleet/meta_parallel/pp_utils/utils.py | 1 +
 python/paddle/fluid/contrib/mixed_precision/fp16_lists.py    | 1 +
 python/paddle/fluid/dygraph/amp/auto_cast.py                 | 1 +
 4 files changed, 8 insertions(+)

diff --git a/paddle/fluid/operators/flatten_op.cu.cc b/paddle/fluid/operators/flatten_op.cu.cc
index 223cfc6ac667d..e0987288abdd7 100644
--- a/paddle/fluid/operators/flatten_op.cu.cc
+++ b/paddle/fluid/operators/flatten_op.cu.cc
@@ -15,6 +15,7 @@ limitations under the License. */
 #include "paddle/fluid/operators/flatten_op.h"
 
 namespace ops = paddle::operators;
+namespace plat = paddle::platform;
 
 REGISTER_OP_CUDA_KERNEL(
     flatten, ops::FlattenKernel<paddle::platform::CUDADeviceContext, float>,
@@ -50,6 +51,8 @@ REGISTER_OP_CUDA_KERNEL(
     flatten_contiguous_range,
     ops::FlattenContiguousRangeKernel<paddle::platform::CUDADeviceContext,
                                       float>,
+    ops::FlattenContiguousRangeKernel<paddle::platform::CUDADeviceContext,
+                                      plat::float16>,
     ops::FlattenContiguousRangeKernel<paddle::platform::CUDADeviceContext,
                                       double>,
     ops::FlattenContiguousRangeKernel<paddle::platform::CUDADeviceContext,
@@ -63,6 +66,8 @@ REGISTER_OP_CUDA_KERNEL(
     flatten_contiguous_range_grad,
     ops::FlattenContiguousRangeGradKernel<paddle::platform::CUDADeviceContext,
                                           float>,
+    ops::FlattenContiguousRangeGradKernel<paddle::platform::CUDADeviceContext,
+                                          plat::float16>,
     ops::FlattenContiguousRangeGradKernel<paddle::platform::CUDADeviceContext,
                                           double>,
     ops::FlattenContiguousRangeGradKernel<paddle::platform::CUDADeviceContext,
diff --git a/python/paddle/distributed/fleet/meta_parallel/pp_utils/utils.py b/python/paddle/distributed/fleet/meta_parallel/pp_utils/utils.py
index fc1fc4f992e36..598c4b26423fd 100644
--- a/python/paddle/distributed/fleet/meta_parallel/pp_utils/utils.py
+++ b/python/paddle/distributed/fleet/meta_parallel/pp_utils/utils.py
@@ -133,6 +133,7 @@ def _split_activation(tensor):
 
     # use inplace operation to save memory
     data = tensor.flatten_()
+
     part_size = tensor_numel // mp_degree
     start = part_size * mp_rank
     end = start + part_size
diff --git a/python/paddle/fluid/contrib/mixed_precision/fp16_lists.py b/python/paddle/fluid/contrib/mixed_precision/fp16_lists.py
index 703146736e3c1..5b662b09f1cf6 100644
--- a/python/paddle/fluid/contrib/mixed_precision/fp16_lists.py
+++ b/python/paddle/fluid/contrib/mixed_precision/fp16_lists.py
@@ -94,6 +94,7 @@ def _update_list(self):
     'softmax',
     'softmax_with_cross_entropy',
     'sigmoid_cross_entropy_with_logits',
+    'c_softmax_with_cross_entropy',
     'cross_entropy',
     'cross_entropy2',
     # fp16 is slower than fp32, though fp16 is supported.
diff --git a/python/paddle/fluid/dygraph/amp/auto_cast.py b/python/paddle/fluid/dygraph/amp/auto_cast.py
index a7eb0d31b7f85..759ce3d16a46a 100644
--- a/python/paddle/fluid/dygraph/amp/auto_cast.py
+++ b/python/paddle/fluid/dygraph/amp/auto_cast.py
@@ -45,6 +45,7 @@
     'softmax',
     'softmax_with_cross_entropy',
     'sigmoid_cross_entropy_with_logits',
+    'c_softmax_with_cross_entropy',
     'cross_entropy',
     'cross_entropy2',
     # default fp32 can avoid return inf when the sum value large than 65504