Skip to content

[XPU] swiglu support xpu && xhpc date update && set value support bf16 #64070

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
houj04 merged 1 commit into from
May 8, 2024
Merged

[XPU] swiglu support xpu && xhpc date update && set value support bf16 #64070

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cmake/external/xpu.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ if(NOT DEFINED XPU_XDNN_BASE_DATE)
set(XPU_XDNN_BASE_DATE "20240327")
endif()
if(NOT DEFINED XPU_XHPC_BASE_DATE)
set(XPU_XHPC_BASE_DATE "20240422")
set(XPU_XHPC_BASE_DATE "20240506")
endif()
set(XPU_XCCL_BASE_VERSION "1.2.0.5")
if(NOT DEFINED XPU_XFT_BASE_VERSION)
Expand Down
14 changes: 12 additions & 2 deletions paddle/phi/backends/xpu/xpu3_op_list.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -845,13 +845,15 @@ XPUOpMap& get_kl3_ops() {
phi::DataType::INT32,
phi::DataType::INT64,
phi::DataType::FLOAT16,
phi::DataType::BOOL})},
phi::DataType::BOOL,
phi::DataType::BFLOAT16})},
{"set_value_with_tensor",
XPUKernelSet({phi::DataType::FLOAT32,
phi::DataType::FLOAT16,
phi::DataType::INT32,
phi::DataType::INT64,
phi::DataType::BOOL})},
phi::DataType::BOOL,
phi::DataType::BFLOAT16})},
{"set_value_grad",
XPUKernelSet({phi::DataType::FLOAT32,
phi::DataType::INT32,
Expand Down Expand Up @@ -1030,6 +1032,14 @@ XPUOpMap& get_kl3_ops() {
phi::DataType::INT16,
phi::DataType::INT32})},
{"sum", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
{"swiglu",
XPUKernelSet({phi::DataType::FLOAT32,
phi::DataType::FLOAT16,
phi::DataType::BFLOAT16})},
{"swiglu_grad",
XPUKernelSet({phi::DataType::FLOAT32,
phi::DataType::FLOAT16,
phi::DataType::BFLOAT16})},
{"swish",
XPUKernelSet({phi::DataType::FLOAT32,
phi::DataType::FLOAT16,
Expand Down
2 changes: 2 additions & 0 deletions paddle/phi/kernels/xpu/set_value_kernel.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -430,6 +430,7 @@ PD_REGISTER_KERNEL(set_value,
phi::SetValueKernel,
float,
phi::dtype::float16,
phi::dtype::bfloat16,
int,
int64_t,
bool) {}
Expand All @@ -440,6 +441,7 @@ PD_REGISTER_KERNEL(set_value_with_tensor,
phi::SetTensorValueKernel,
float,
phi::dtype::float16,
phi::dtype::bfloat16,
int,
int64_t,
bool) {}
66 changes: 66 additions & 0 deletions paddle/phi/kernels/xpu/swiglu_kernel.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
// 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/kernels/swiglu_kernel.h"
#include "paddle/phi/backends/xpu/enforce_xpu.h"
#include "paddle/phi/core/kernel_registry.h"

namespace phi {
template <typename T, typename Context>
void SwiGluKernel(const Context& ctx,
const DenseTensor& x,
const paddle::optional<DenseTensor>& y,
DenseTensor* z) {
using XPUType = typename XPUTypeTrait<T>::Type;
using XPUTypefp32 = typename XPUTypeTrait<float>::Type;
const auto* x_data = x.data<T>();
auto* z_data = ctx.template Alloc<T>(z);
const auto& dims = x.dims();
int64_t axis = dims.size() - 1;
auto dims_vec = common::vectorize<int64_t>(dims);
const XPUTypefp32* const_nullptr = nullptr;
const XPUType* y_ptr = nullptr;

if (y) {
const auto& y_tensor = y.get();
const auto& y_dims = y_tensor.dims();
const auto* y_data = y_tensor.data<T>();
y_ptr = reinterpret_cast<const XPUType*>(y_data);
PADDLE_ENFORCE_EQ(
y_dims,
dims,
phi::errors::InvalidArgument("The shape of Input(Y):[%s] must be equal "
"to the shape of Input(X):[%s].",
y_dims,
dims));
}
int ret = xpu::swiglu(ctx.x_context(),
reinterpret_cast<const XPUType*>(x_data),
reinterpret_cast<XPUType*>(z_data),
dims_vec,
axis,
false,
const_nullptr,
nullptr,
y_ptr);
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "swiglu");
}
} // namespace phi
PD_REGISTER_KERNEL(swiglu,
XPU,
ALL_LAYOUT,
phi::SwiGluKernel,
float,
phi::dtype::float16,
phi::dtype::bfloat16){};
79 changes: 79 additions & 0 deletions paddle/phi/kernels/xpu/swiglu_kernel_grad.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
// 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.

// 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/backends/xpu/enforce_xpu.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/swiglu_grad_kernel.h"

namespace phi {

template <typename T, typename Context>
void SwiGluGradKernel(const Context& ctx,
const DenseTensor& x,
const paddle::optional<DenseTensor>& y,
const DenseTensor& dz,
DenseTensor* dx,
DenseTensor* dy) {
using XPUType = typename XPUTypeTrait<T>::Type;
const auto* x_data = x.data<T>();
const auto* dz_data = dz.data<T>();
auto* dx_data = ctx.template Alloc<T>(dx);
const auto& dims = x.dims();
int64_t axis = dims.size() - 1;
auto dims_vec = common::vectorize<int64_t>(dims);
const XPUType* y_ptr = nullptr;
XPUType* dy_ptr = nullptr;

if (y) {
const auto& y_tensor = y.get();
const auto& y_dims = y_tensor.dims();
const auto* y_data = y_tensor.data<T>();
auto* dy_data = ctx.template Alloc<T>(dy);
y_ptr = reinterpret_cast<const XPUType*>(y_data);
dy_ptr = reinterpret_cast<XPUType*>(dy_data);
PADDLE_ENFORCE_EQ(
y_dims,
dims,
phi::errors::InvalidArgument("The shape of Input(Y):[%s] must be equal "
"to the shape of Input(X):[%s].",
y_dims,
dims));
}
int ret = xpu::swiglu_grad(ctx.x_context(),
reinterpret_cast<const XPUType*>(x_data),
reinterpret_cast<const XPUType*>(dz_data),
reinterpret_cast<XPUType*>(dx_data),
dims_vec,
axis,
false,
y_ptr,
dy_ptr);
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "swiglu_grad");
}
} // namespace phi
PD_REGISTER_KERNEL(swiglu_grad,
XPU,
ALL_LAYOUT,
phi::SwiGluGradKernel,
float,
phi::dtype::float16,
phi::dtype::bfloat16){};