wehub-resource-sync
122 lines
4.3 KiB
Plaintext
122 lines
4.3 KiB
Plaintext
// Copyright (c) 2023 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/dense_tensor.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/tensor_utils.h"
|
|
|
|
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
|
#include "paddle/phi/kernels/prune_gate_by_capacity_kernel.h"
|
|
|
|
namespace phi {
|
|
|
|
static constexpr int kNumCUDAThreads = 512;
|
|
static constexpr int kNumMaximumNumBlocks = 4096;
|
|
|
|
static inline int NumBlocks(const int N) {
|
|
return std::min((N + kNumCUDAThreads - 1) / kNumCUDAThreads,
|
|
kNumMaximumNumBlocks);
|
|
}
|
|
|
|
template <typename T1, typename T2>
|
|
__global__ void prune_gate_by_capacity_kernel(const T1* gate_idx_data,
|
|
T1* new_gate_idx_data,
|
|
T2* expert_count_data,
|
|
const int64_t batch_size) {
|
|
CUDA_KERNEL_LOOP(i, batch_size) {
|
|
auto orig_cap = CudaAtomicAdd(expert_count_data + gate_idx_data[i], -1);
|
|
if (orig_cap <= 0) {
|
|
new_gate_idx_data[i] = -1;
|
|
} else {
|
|
new_gate_idx_data[i] = gate_idx_data[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename Context, typename T1>
|
|
class PruneGateByCapacityFunctor {
|
|
public:
|
|
PruneGateByCapacityFunctor(const Context& dev_ctx,
|
|
const DenseTensor* gate_idx,
|
|
DenseTensor* expert_count_out,
|
|
T1* new_gate_idx_data)
|
|
: dev_ctx_(dev_ctx),
|
|
gate_idx_(gate_idx),
|
|
expert_count_out_(expert_count_out),
|
|
new_gate_idx_data_(new_gate_idx_data) {}
|
|
|
|
template <typename T2>
|
|
void apply() {
|
|
auto batch_size = gate_idx_->numel();
|
|
auto* gate_idx_data = gate_idx_->data<T1>();
|
|
|
|
auto* expert_count_out_data = expert_count_out_->data<T2>();
|
|
|
|
int blocks = NumBlocks(batch_size);
|
|
int threads = kNumCUDAThreads;
|
|
|
|
prune_gate_by_capacity_kernel<T1, T2>
|
|
<<<blocks, threads, 0, dev_ctx_.stream()>>>(gate_idx_data,
|
|
new_gate_idx_data_,
|
|
expert_count_out_data,
|
|
batch_size);
|
|
}
|
|
|
|
private:
|
|
const Context& dev_ctx_;
|
|
const DenseTensor* gate_idx_;
|
|
DenseTensor* expert_count_out_;
|
|
T1* new_gate_idx_data_;
|
|
};
|
|
|
|
template <typename Visitor>
|
|
static void VisitType(DataType type, Visitor visitor) {
|
|
if (type == DataType::INT64) {
|
|
visitor.template apply<int64_t>();
|
|
} else {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"The received values gate_id type %s can not meet input requirements. "
|
|
"Because the given gate_id data type of operators must be "
|
|
"int64. Please input appropriate gate_id again! ",
|
|
"framework::DataTypeToString(type)"));
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void PruneGateByCapacityKernel(const Context& dev_ctx,
|
|
const DenseTensor& gate_idx,
|
|
const DenseTensor& expert_count,
|
|
int64_t n_expert,
|
|
int64_t n_worker,
|
|
DenseTensor* new_gate_idx) {
|
|
auto* gate_idx_ptr = &gate_idx;
|
|
// auto* expert_count_out =
|
|
// context.Output<DenseTensor>("ExpertCountOut");
|
|
auto* new_gate_idx_data = dev_ctx.template Alloc<T>(new_gate_idx);
|
|
|
|
DenseTensor expert_count_out;
|
|
Copy(dev_ctx, expert_count, dev_ctx.GetPlace(), false, &expert_count_out);
|
|
PruneGateByCapacityFunctor<Context, T> functor(
|
|
dev_ctx, gate_idx_ptr, &expert_count_out, new_gate_idx_data);
|
|
VisitType(expert_count.type(), functor);
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(prune_gate_by_capacity,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::PruneGateByCapacityKernel,
|
|
int64_t) {}
|