Fix (PaddlePaddle#67012)

Author: co63oc

paddle/fluid/operators/collective/partial_allgather_op.cc

@@ -12,7 +12,9 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
 
-#include "paddle/fluid/operators/collective/partial_allgather_op.h"
+#include "paddle/fluid/framework/data_type.h"
+#include "paddle/fluid/framework/lod_tensor.h"
+#include "paddle/fluid/framework/op_registry.h"
 
 namespace paddle {
 namespace operators {
@@ -83,13 +85,3 @@ REGISTER_OPERATOR(
     paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
     paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>,
     ops::PartialAllGatherOpInplaceInferer)
-
-PD_REGISTER_STRUCT_KERNEL(partial_allgather,
-                          CPU,
-                          ALL_LAYOUT,
-                          ops::PartialAllGatherOpCPUKernel,
-                          float,
-                          double,
-                          int,
-                          int64_t,
-                          phi::dtype::float16) {}

paddle/fluid/operators/collective/partial_allgather_op.cu.cc

@@ -0,0 +1,41 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void PartialAllGatherOpCPUKernel(const Context &dev_ctx,
+                                 const DenseTensor &x,
+                                 int nranks,
+                                 int rank,
+                                 int ring_id,
+                                 bool use_calc_stream,
+                                 DenseTensor *out) {
+  PADDLE_THROW(common::errors::Unavailable(
+      "Do not support partial_allgather for cpu kernel now."));
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(partial_allgather,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::PartialAllGatherOpCPUKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::dtype::float16) {}