               << " , " << GetMinMaxStr(tensors[i]);
     }
   }
+}
+
+template <typename T>
+static bool HasNanInf(const phi::GPUContext &dev_ctx, const T *x, int numel) {
+  if (numel <= 0) return false;
+  cub::TransformInputIterator<bool, IsNanInfFunctor<T>, const T *> iter(
+      x, IsNanInfFunctor<T>());
+  memory::Buffer buffer(dev_ctx.GetPlace());
+  memory::Buffer out(dev_ctx.GetPlace());
+  CubDeviceReduce(iter,
+                  out.Alloc<bool>(1),
+                  numel,
+                  OrFunctor(),
+                  false,
+                  dev_ctx.stream(),
+                  &buffer);
+  bool flag;
+#ifdef PADDLE_WITH_HIP
+  PADDLE_ENFORCE_GPU_SUCCESS(hipMemcpyAsync(&flag,
+                                            out.Get<bool>(),
+                                            sizeof(flag),
+                                            hipMemcpyDeviceToHost,
+                                            dev_ctx.stream()));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpyAsync(&flag,
+                                             out.Get<bool>(),
+                                             sizeof(flag),
+                                             cudaMemcpyDeviceToHost,
+                                             dev_ctx.stream()));
+#endif
+  dev_ctx.Wait();
+  return flag;
 }
 
 static void CheckHasNanInfGrad(const float *fp32_grad,

                                     stream,
                                     dev_ctx,
                                     fp16_scale);
        }
        // (4) mark max_global_grad_norm as 0, meaning that clip has been
        // already performed

        } else {
          VLOG(1) << "Grad scale: " << FlattenToString(fp16_scale, 1, place);
        }
        // (3) Do ReduceScatter with scale
        VLOG(1) << "FP32 HasNanInf before all reduce: "
                << HasNanInf(dev_ctx, fp32_grad, fp32_numel);
        VLOG(1) << "FP16 HasNanInf before all reduce: "
                << HasNanInf(dev_ctx, fp16_grad, fp16_numel);
        if (local_shard) {
          if (use_hierarchical_allreduce) {
            NCCLReduceScatterWithScale(
