mirror of https://github.com/hpcaitech/ColossalAI
1002 lines
37 KiB
Plaintext
1002 lines
37 KiB
Plaintext
#include <cooperative_groups.h>
|
|
|
|
#include <chrono>
|
|
#include <ctime>
|
|
|
|
#include "kernels.h"
|
|
|
|
namespace cg = cooperative_groups;
|
|
|
|
curandStatePhilox4_32_10_t *curandstate;
|
|
|
|
/**
|
|
* @brief element-wise activation function on device, like Relu, Gelu
|
|
*
|
|
* @tparam enum class ActivationType, kRelu, kGelu
|
|
* @tparam input type
|
|
* @param any shape of float and __half2
|
|
* @return same shape and type with input
|
|
*/
|
|
template <ActivationType, typename T>
|
|
__forceinline__ __device__ T activation_kernel(T x);
|
|
|
|
template <>
|
|
__device__ float activation_kernel<ActivationType::kGelu, float>(float x) {
|
|
float cdf =
|
|
0.5f *
|
|
(1.0f + tanhf((0.7978845608028654f * (x + 0.044715f * x * x * x))));
|
|
return x * cdf;
|
|
}
|
|
|
|
template <>
|
|
__device__ __half2
|
|
activation_kernel<ActivationType::kGelu, __half2>(__half2 val) {
|
|
__half2 val_pow3 = __hmul2(val, __hmul2(val, val));
|
|
float2 tmp_pow = __half22float2(val_pow3);
|
|
float2 tmp = __half22float2(val);
|
|
|
|
tmp.x =
|
|
0.5f *
|
|
(1.0f + tanhf((0.7978845608028654f * (tmp.x + 0.044715f * tmp_pow.x))));
|
|
tmp.y =
|
|
0.5f *
|
|
(1.0f + tanhf((0.7978845608028654f * (tmp.y + 0.044715f * tmp_pow.y))));
|
|
return __hmul2(val, __float22half2_rn(tmp));
|
|
}
|
|
|
|
template <>
|
|
__device__ float activation_kernel<ActivationType::kRelu, float>(float x) {
|
|
return fmaxf(x, 0);
|
|
}
|
|
|
|
template <>
|
|
__device__ __half2
|
|
activation_kernel<ActivationType::kRelu, __half2>(__half2 x) {
|
|
return __floats2half2_rn(fmaxf(0.f, __half2float(x.x)),
|
|
fmaxf(0.f, __half2float(x.y)));
|
|
}
|
|
|
|
/**
|
|
* @brief element-wise activation backward function on device
|
|
*
|
|
* @tparam enum class ActivationType
|
|
* @tparam input type
|
|
* @param any shape of float and __half2
|
|
* @return same shape of input
|
|
*/
|
|
template <ActivationType, typename T>
|
|
__forceinline__ __device__ T activation_bwd_kernel(T grad, T x);
|
|
|
|
template <>
|
|
__device__ float activation_bwd_kernel<ActivationType::kGelu, float>(float grad,
|
|
float x) {
|
|
const float sqrt_param = 0.79788456080286535587989211986876f;
|
|
const float mul_param = 0.044715;
|
|
|
|
float x2mul = x * x * mul_param;
|
|
float tan_h = tanhf(sqrt_param * (x + x * x2mul));
|
|
float dg1 = 0.5f * (1.0f + tan_h);
|
|
float dg2 = x * 0.5f * sqrt_param * (1 - tan_h * tan_h);
|
|
float dg3 = dg2 * 3 * x2mul;
|
|
return grad * (dg1 + dg2 + dg3);
|
|
}
|
|
|
|
template <>
|
|
__device__ __half activation_bwd_kernel<ActivationType::kGelu, __half>(
|
|
__half grad, __half x_half) {
|
|
float x = __half2float(x_half);
|
|
const float sqrt_param = 0.79788456080286535587989211986876f;
|
|
const float mul_param = 0.044715;
|
|
|
|
float x2mul = x * x * mul_param;
|
|
float tan_h = tanhf(sqrt_param * (x + x * x2mul));
|
|
float dg1 = 0.5f * (1.0f + tan_h);
|
|
float dg2 = x * 0.5f * sqrt_param * (1 - tan_h * tan_h);
|
|
float dg3 = dg2 * 3 * x2mul;
|
|
return grad * __float2half(dg1 + dg2 + dg3);
|
|
}
|
|
|
|
template <>
|
|
__device__ float activation_bwd_kernel<ActivationType::kRelu, float>(float grad,
|
|
float x) {
|
|
return x > 0.f ? grad : 0.f;
|
|
}
|
|
|
|
template <>
|
|
__device__ __half
|
|
activation_bwd_kernel<ActivationType::kRelu, __half>(__half grad, __half x) {
|
|
const __half half_zero = __float2half(0.f);
|
|
return x > half_zero ? grad : half_zero;
|
|
}
|
|
|
|
template <>
|
|
__device__ __half2 activation_bwd_kernel<ActivationType::kRelu, __half2>(
|
|
__half2 grad2, __half2 x_half2) {
|
|
const __half half_zero = __float2half(0.f);
|
|
return __floats2half2_rn(x_half2.x > half_zero ? grad2.x : half_zero,
|
|
x_half2.y > half_zero ? grad2.y : half_zero);
|
|
}
|
|
|
|
/**
|
|
* @brief init curand states in global memory
|
|
*
|
|
* @thread grid_dim * block*dim to suuport any size of states
|
|
* @param state persistant curand states
|
|
* @param seed seed to init states
|
|
* @return void
|
|
*/
|
|
__global__ void curand_init_kernel(curandStatePhilox4_32_10_t *state,
|
|
int seed) {
|
|
/* Each thread gets same seed, a different sequence
|
|
number, no offset */
|
|
int id = threadIdx.x + blockIdx.x * blockDim.x;
|
|
curand_init(seed, id, 0, &state[id]);
|
|
}
|
|
|
|
void launch_curand_init(int total_count, int dim, cudaStream_t stream) {
|
|
cudaMalloc(&curandstate, total_count * sizeof(curandStatePhilox4_32_10_t));
|
|
int grid_dim = total_count >> 9;
|
|
curand_init_kernel<<<grid_dim, 512, 0, stream>>>(
|
|
curandstate, std::chrono::duration_cast<std::chrono::microseconds>(
|
|
std::chrono::system_clock::now().time_since_epoch())
|
|
.count());
|
|
}
|
|
|
|
/**
|
|
* @brief element-wise dropout, store dropped position in mask, it's not
|
|
* in-place
|
|
*
|
|
* @thread
|
|
* gridDim.x = total_count / 1024
|
|
* blockDim.x = 1024
|
|
*
|
|
* @param total_count total elements
|
|
* @param ratio drop ratio
|
|
* @param out any size of float and __half
|
|
* @param in same with out
|
|
* @param mask uint8 type, same size with out
|
|
* @param seed seed to curand
|
|
* @return void
|
|
*/
|
|
__global__ void ls_dropout_kernel(const int total_count, const float ratio,
|
|
float *__restrict__ out,
|
|
const float *__restrict__ in,
|
|
uint8_t *__restrict__ mask, const int seed) {
|
|
const float scale = 1.f / (1.f - ratio);
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (i * 4 >= total_count) return;
|
|
|
|
curandStatePhilox4_32_10_t state;
|
|
curand_init(seed, i, 0, &state);
|
|
uint8_t m[4];
|
|
|
|
float4 *out4 = reinterpret_cast<float4 *>(out);
|
|
const float4 *data4 = reinterpret_cast<const float4 *>(in);
|
|
uint32_t *mask4 = reinterpret_cast<uint32_t *>(mask);
|
|
float4 rand = curand_uniform4(&state);
|
|
|
|
m[0] = (uint8_t)(rand.x > ratio);
|
|
m[1] = (uint8_t)(rand.y > ratio);
|
|
m[2] = (uint8_t)(rand.z > ratio);
|
|
m[3] = (uint8_t)(rand.w > ratio);
|
|
|
|
uint32_t *m4 = reinterpret_cast<uint32_t *>(m);
|
|
mask4[i] = m4[0];
|
|
|
|
float4 input4 = data4[i];
|
|
float4 res4;
|
|
res4.x = input4.x * scale * m[0];
|
|
res4.y = input4.y * scale * m[1];
|
|
res4.z = input4.z * scale * m[2];
|
|
res4.w = input4.w * scale * m[3];
|
|
out4[i] = res4;
|
|
}
|
|
|
|
__global__ void ls_dropout_kernel(const int total_count, const float ratio,
|
|
__half *__restrict__ out,
|
|
const __half *__restrict__ in,
|
|
uint8_t *__restrict__ mask, const int seed) {
|
|
const float scale = 1.f / (1.f - ratio);
|
|
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (i * 8 >= total_count) return;
|
|
|
|
curandStatePhilox4_32_10_t state;
|
|
curand_init(seed, i, 0, &state);
|
|
|
|
const float4 *vals_float4 = reinterpret_cast<const float4 *>(in);
|
|
float4 *outs_float4 = reinterpret_cast<float4 *>(out);
|
|
uint64_t *mask8 = reinterpret_cast<uint64_t *>(mask);
|
|
|
|
uint8_t m[8];
|
|
float4 rand = curand_uniform4(&state);
|
|
m[0] = (uint8_t)(rand.x > ratio);
|
|
m[1] = (uint8_t)(rand.y > ratio);
|
|
m[2] = (uint8_t)(rand.z > ratio);
|
|
m[3] = (uint8_t)(rand.w > ratio);
|
|
rand = curand_uniform4(&state);
|
|
m[4] = (uint8_t)(rand.x > ratio);
|
|
m[5] = (uint8_t)(rand.y > ratio);
|
|
m[6] = (uint8_t)(rand.z > ratio);
|
|
m[7] = (uint8_t)(rand.w > ratio);
|
|
uint64_t *m8 = reinterpret_cast<uint64_t *>(m);
|
|
mask8[i] = *m8;
|
|
|
|
float4 val_float4 = vals_float4[i];
|
|
float4 out_float4;
|
|
__half2 *val_half2 = reinterpret_cast<__half2 *>(&val_float4);
|
|
__half2 *out_half2 = reinterpret_cast<__half2 *>(&out_float4);
|
|
__half2 scale_mask_1 = __floats2half2_rn(scale * m[0], scale * m[1]);
|
|
__half2 scale_mask_2 = __floats2half2_rn(scale * m[2], scale * m[3]);
|
|
__half2 scale_mask_3 = __floats2half2_rn(scale * m[4], scale * m[5]);
|
|
__half2 scale_mask_4 = __floats2half2_rn(scale * m[6], scale * m[7]);
|
|
out_half2[0] = __hmul2(val_half2[0], scale_mask_1);
|
|
out_half2[1] = __hmul2(val_half2[1], scale_mask_2);
|
|
out_half2[2] = __hmul2(val_half2[2], scale_mask_3);
|
|
out_half2[3] = __hmul2(val_half2[3], scale_mask_4);
|
|
outs_float4[i] = out_float4;
|
|
}
|
|
|
|
/**
|
|
* @brief element-wise dropout backward with dropout mask, it's
|
|
* not in-place
|
|
*
|
|
* @thread
|
|
* gridDim.x = total_count / 1024
|
|
* blockDim.x = 1024
|
|
*
|
|
* @param total_count total elements
|
|
* @param ratio drop ratio
|
|
* @param in any size of float and __half
|
|
* @param mask uint8 type, same size with in
|
|
* @return void
|
|
*/
|
|
__global__ void ls_dropout_bwd_kernel(const int total_count, const float ratio,
|
|
float *out, const float *in,
|
|
const uint8_t *__restrict__ mask) {
|
|
const float scale = 1.f / (1.f - ratio);
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (i * 4 >= total_count) return;
|
|
|
|
uint8_t m[4];
|
|
|
|
float4 *out4 = reinterpret_cast<float4 *>(out);
|
|
const float4 *in4 = reinterpret_cast<const float4 *>(in);
|
|
const uint32_t *mask4 = reinterpret_cast<const uint32_t *>(mask);
|
|
|
|
uint32_t *m4 = reinterpret_cast<uint32_t *>(m);
|
|
m4[0] = mask4[i];
|
|
|
|
float4 input4 = in4[i];
|
|
float4 res4;
|
|
res4.x = input4.x * scale * static_cast<float>(m[0]);
|
|
res4.y = input4.y * scale * static_cast<float>(m[1]);
|
|
res4.z = input4.z * scale * static_cast<float>(m[2]);
|
|
res4.w = input4.w * scale * static_cast<float>(m[3]);
|
|
out4[i] = res4;
|
|
}
|
|
|
|
__global__ void ls_dropout_bwd_kernel(const int total_count, const float ratio,
|
|
__half *out, const __half *in,
|
|
const uint8_t *__restrict__ mask) {
|
|
const __half scale = 1.f / (1.f - ratio);
|
|
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (i * 8 >= total_count) return;
|
|
|
|
float4 *out4 = reinterpret_cast<float4 *>(out);
|
|
const float4 *vals_float4 = reinterpret_cast<const float4 *>(in);
|
|
const uint64_t *mask8 = reinterpret_cast<const uint64_t *>(mask);
|
|
|
|
uint8_t m[8];
|
|
uint64_t *m8 = reinterpret_cast<uint64_t *>(m);
|
|
m8[0] = mask8[i];
|
|
|
|
float4 val_float4 = vals_float4[i];
|
|
float4 out_float4;
|
|
__half2 *val_half2 = reinterpret_cast<__half2 *>(&val_float4);
|
|
__half2 *out_half2 = reinterpret_cast<__half2 *>(&out_float4);
|
|
__half2 scale_mask_1 =
|
|
__halves2half2(scale * __float2half(m[0]), scale * __float2half(m[1]));
|
|
__half2 scale_mask_2 =
|
|
__halves2half2(scale * __float2half(m[2]), scale * __float2half(m[3]));
|
|
__half2 scale_mask_3 =
|
|
__halves2half2(scale * __float2half(m[4]), scale * __float2half(m[5]));
|
|
__half2 scale_mask_4 =
|
|
__halves2half2(scale * __float2half(m[6]), scale * __float2half(m[7]));
|
|
out_half2[0] = __hmul2(val_half2[0], scale_mask_1);
|
|
out_half2[1] = __hmul2(val_half2[1], scale_mask_2);
|
|
out_half2[2] = __hmul2(val_half2[2], scale_mask_3);
|
|
out_half2[3] = __hmul2(val_half2[3], scale_mask_4);
|
|
out4[i] = out_float4;
|
|
}
|
|
|
|
template <>
|
|
void launch_ls_dropout<float>(float *out, const float *vals, uint8_t *mask,
|
|
int total_count, float ratio, cudaStream_t stream,
|
|
bool backward) {
|
|
int grid_dim = total_count >> 12;
|
|
if (!backward) {
|
|
ls_dropout_kernel<<<grid_dim + 1, 1024, 0, stream>>>(
|
|
total_count, ratio, out, vals, mask,
|
|
std::chrono::duration_cast<std::chrono::microseconds>(
|
|
std::chrono::system_clock::now().time_since_epoch())
|
|
.count());
|
|
} else {
|
|
ls_dropout_bwd_kernel<<<grid_dim + 1, 1024, 0, stream>>>(total_count, ratio,
|
|
out, vals, mask);
|
|
}
|
|
}
|
|
|
|
template <>
|
|
void launch_ls_dropout<__half>(__half *out, const __half *vals, uint8_t *mask,
|
|
int total_count, float ratio,
|
|
cudaStream_t stream, bool backward) {
|
|
int grid_dim = total_count >> 13;
|
|
if (!backward) {
|
|
ls_dropout_kernel<<<grid_dim + 1, 1024, 0, stream>>>(
|
|
total_count, ratio, out, vals, mask,
|
|
std::chrono::duration_cast<std::chrono::microseconds>(
|
|
std::chrono::system_clock::now().time_since_epoch())
|
|
.count());
|
|
} else {
|
|
ls_dropout_bwd_kernel<<<grid_dim + 1, 1024, 0, stream>>>(total_count, ratio,
|
|
out, vals, mask);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* @brief fused bias, dropout, and residual at the end of Attention and FFN,
|
|
* store dropped position in mask, it's not in-place
|
|
*
|
|
* @thread
|
|
* gridDim.x = total_count / 1024
|
|
* blockDim.x = 1024
|
|
*
|
|
* @param total_count total elements
|
|
* @param ratio drop ratio
|
|
* @param out [batch_size, seq_len, hidden_size], float and __half
|
|
* @param in [batch_size, seq_len, hidden_size], float and __half
|
|
* @param mask [batch_size, seq_len, hidden_size], uint8 type
|
|
* @param bias [hidden_size], ffn bias
|
|
* @param residual [batch_size, seq_len, hidden_size], float and __half
|
|
* @param seed seed to curand
|
|
* @param hidden_size hidden size
|
|
* @return void
|
|
*/
|
|
__global__ void ls_dropout_res_bias_kernel(
|
|
const int total_count, const float ratio, float *__restrict__ out,
|
|
const float *__restrict__ in, uint8_t *__restrict__ mask,
|
|
const float *__restrict__ bias, const float *__restrict__ residual,
|
|
const int seed, const int hidden_size) {
|
|
const float scale = 1.f / (1.f - ratio);
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (i * 4 >= total_count) return;
|
|
|
|
curandStatePhilox4_32_10_t state;
|
|
curand_init(seed, i, 0, &state);
|
|
uint8_t m[4];
|
|
|
|
float4 *out4 = reinterpret_cast<float4 *>(out);
|
|
const float4 *data4 = reinterpret_cast<const float4 *>(in);
|
|
const float4 *residual4 = reinterpret_cast<const float4 *>(residual);
|
|
const float4 *bias4 = reinterpret_cast<const float4 *>(bias);
|
|
uint32_t *mask4 = reinterpret_cast<uint32_t *>(mask);
|
|
float4 rand = curand_uniform4(&state);
|
|
|
|
m[0] = static_cast<uint8_t>(rand.x > ratio);
|
|
m[1] = static_cast<uint8_t>(rand.y > ratio);
|
|
m[2] = static_cast<uint8_t>(rand.z > ratio);
|
|
m[3] = static_cast<uint8_t>(rand.w > ratio);
|
|
|
|
int bias_i = i % (hidden_size >> 2);
|
|
uint32_t *m4 = reinterpret_cast<uint32_t *>(m);
|
|
mask4[i] = m4[0];
|
|
const float4 input4 = data4[i];
|
|
const float4 b4 = __ldg(&bias4[bias_i]);
|
|
const float4 res4 = residual4[i];
|
|
float4 output4;
|
|
|
|
output4.x = (input4.x + b4.x) * scale * m[0] + res4.x;
|
|
output4.y = (input4.y + b4.y) * scale * m[1] + res4.y;
|
|
output4.z = (input4.z + b4.z) * scale * m[2] + res4.z;
|
|
output4.w = (input4.w + b4.w) * scale * m[3] + res4.w;
|
|
|
|
out4[i] = output4;
|
|
}
|
|
|
|
__global__ void ls_dropout_res_bias_kernel(
|
|
const int total_count, const float ratio, __half *__restrict__ out,
|
|
const __half *__restrict__ in, uint8_t *__restrict__ mask,
|
|
const __half *__restrict__ bias, const __half *__restrict__ residual,
|
|
const int seed, const int hidden_size) {
|
|
const __half scale = 1. / (1. - ratio);
|
|
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (i * 8 >= total_count) return;
|
|
|
|
curandStatePhilox4_32_10_t state;
|
|
curand_init(seed, i, 0, &state);
|
|
|
|
const float4 *vals_float4 = reinterpret_cast<const float4 *>(in);
|
|
float4 *outs_float4 = reinterpret_cast<float4 *>(out);
|
|
const float4 *residual4 = reinterpret_cast<const float4 *>(residual);
|
|
const float4 *bias4 = reinterpret_cast<const float4 *>(bias);
|
|
uint64_t *mask8 = reinterpret_cast<uint64_t *>(mask);
|
|
|
|
uint8_t m[8];
|
|
float4 rand = curand_uniform4(&state);
|
|
m[0] = static_cast<uint8_t>(rand.x > ratio);
|
|
m[1] = static_cast<uint8_t>(rand.y > ratio);
|
|
m[2] = static_cast<uint8_t>(rand.z > ratio);
|
|
m[3] = static_cast<uint8_t>(rand.w > ratio);
|
|
rand = curand_uniform4(&state);
|
|
m[4] = static_cast<uint8_t>(rand.x > ratio);
|
|
m[5] = static_cast<uint8_t>(rand.y > ratio);
|
|
m[6] = static_cast<uint8_t>(rand.z > ratio);
|
|
m[7] = static_cast<uint8_t>(rand.w > ratio);
|
|
uint64_t *m8 = reinterpret_cast<uint64_t *>(m);
|
|
mask8[i] = m8[0];
|
|
|
|
int bias_i = i % (hidden_size >> 3);
|
|
float4 val_float4 = vals_float4[i];
|
|
const float4 b4 = __ldg(&bias4[bias_i]);
|
|
const float4 res4 = residual4[i];
|
|
float4 out_float4;
|
|
|
|
__half2 *val_half2 = reinterpret_cast<__half2 *>(&val_float4);
|
|
__half2 *out_half2 = reinterpret_cast<__half2 *>(&out_float4);
|
|
const __half2 *b_half2 = reinterpret_cast<const __half2 *>(&b4);
|
|
const __half2 *res_half2 = reinterpret_cast<const __half2 *>(&res4);
|
|
__half2 scale_mask_1 =
|
|
__halves2half2(scale * __float2half(m[0]), scale * __float2half(m[1]));
|
|
__half2 scale_mask_2 =
|
|
__halves2half2(scale * __float2half(m[2]), scale * __float2half(m[3]));
|
|
__half2 scale_mask_3 =
|
|
__halves2half2(scale * __float2half(m[4]), scale * __float2half(m[5]));
|
|
__half2 scale_mask_4 =
|
|
__halves2half2(scale * __float2half(m[6]), scale * __float2half(m[7]));
|
|
out_half2[0] =
|
|
__hfma2(__hadd2(val_half2[0], b_half2[0]), scale_mask_1, res_half2[0]);
|
|
out_half2[1] =
|
|
__hfma2(__hadd2(val_half2[1], b_half2[1]), scale_mask_2, res_half2[1]);
|
|
out_half2[2] =
|
|
__hfma2(__hadd2(val_half2[2], b_half2[2]), scale_mask_3, res_half2[2]);
|
|
out_half2[3] =
|
|
__hfma2(__hadd2(val_half2[3], b_half2[3]), scale_mask_4, res_half2[3]);
|
|
outs_float4[i] = out_float4;
|
|
}
|
|
|
|
template <>
|
|
void launch_ls_dropout_res_bias<float>(float *out, const float *vals,
|
|
uint8_t *mask, const float *bias,
|
|
const float *residual, int total_count,
|
|
int dim, float ratio,
|
|
cudaStream_t stream) {
|
|
int grid_dim = total_count >> 12;
|
|
ls_dropout_res_bias_kernel<<<grid_dim + 1, 1024, 0, stream>>>(
|
|
total_count, ratio, out, vals, mask, bias, residual,
|
|
std::chrono::duration_cast<std::chrono::microseconds>(
|
|
std::chrono::system_clock::now().time_since_epoch())
|
|
.count(),
|
|
dim);
|
|
}
|
|
|
|
template <>
|
|
void launch_ls_dropout_res_bias<__half>(__half *out, const __half *vals,
|
|
uint8_t *mask, const __half *bias,
|
|
const __half *residual, int total_count,
|
|
int dim, float ratio,
|
|
cudaStream_t stream) {
|
|
int grid_dim = total_count >> 13;
|
|
ls_dropout_res_bias_kernel<<<grid_dim + 1, 1024, 0, stream>>>(
|
|
total_count, ratio, out, vals, mask, bias, residual,
|
|
std::chrono::duration_cast<std::chrono::microseconds>(
|
|
std::chrono::system_clock::now().time_since_epoch())
|
|
.count(),
|
|
dim);
|
|
}
|
|
|
|
/**
|
|
* @brief fused bias and dropout backward at the end of Attention and FFN
|
|
*
|
|
* @thread
|
|
* gridDim.x = hidden_size / 8
|
|
* blockDim.x = 8
|
|
* blockDim.y = 1024 / 8 = 128
|
|
*
|
|
* @param row_size batch_size * seq_len
|
|
* @param ratio dropout ratio
|
|
* @param in_grad [batch_size, seq_len, hidden_size], input grad
|
|
* @param bias_grad [hidden_size], bias grad
|
|
* @param out_grad [batch_size, seq_len, hidden_size], output grad
|
|
* @param mask [batch_size, seq_len, hidden_size], dropout mask
|
|
* @param hidden_size
|
|
* @return void
|
|
*/
|
|
__global__ void ls_dropout_bias_bwd_kernel(
|
|
const int row_size, const float ratio, float *__restrict__ in_grad,
|
|
float *__restrict__ bias_grad, const float *__restrict__ out_grad,
|
|
const uint8_t *__restrict__ mask, const int hidden_size) {
|
|
const float scale = 1.f / (1.f - ratio);
|
|
// every block generate 8 bias result
|
|
__shared__ float tile[8][129];
|
|
|
|
cg::thread_block b = cg::this_thread_block();
|
|
cg::thread_block_tile<WARP_SIZE> g = cg::tiled_partition<WARP_SIZE>(b);
|
|
|
|
int col_idx = flat_2dim(blockIdx.x, threadIdx.x, 8);
|
|
int stride = hidden_size * 128;
|
|
float local_sum = 0;
|
|
|
|
int idx = flat_2dim(threadIdx.y, col_idx, hidden_size);
|
|
for (int r = threadIdx.y; r < row_size; r += 128) {
|
|
float val = out_grad[idx];
|
|
val *= scale * static_cast<float>(mask[idx]);
|
|
local_sum += val;
|
|
in_grad[idx] = val;
|
|
idx += stride;
|
|
}
|
|
|
|
tile[threadIdx.x][threadIdx.y] = local_sum;
|
|
__syncthreads();
|
|
|
|
float sum = 0;
|
|
int tid = threadIdx.y * blockDim.x + threadIdx.x;
|
|
int x = tid >> 7;
|
|
int y = tid & (127);
|
|
if (y < 32) {
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; i++) {
|
|
sum += tile[x][y + i * 32];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
for (int i = 1; i < 32; i <<= 1) sum += g.shfl_down(sum, i);
|
|
|
|
if (y == 0) tile[0][x] = sum;
|
|
__syncthreads();
|
|
|
|
if (threadIdx.x < 8) {
|
|
int pos = flat_2dim(blockIdx.x, threadIdx.x, 8);
|
|
bias_grad[pos] = tile[0][threadIdx.x];
|
|
}
|
|
}
|
|
|
|
__global__ void ls_dropout_bias_bwd_kernel(
|
|
const int row_size, const float ratio, __half *__restrict__ in_grad,
|
|
__half *__restrict__ bias_grad, const __half *__restrict__ out_grad,
|
|
const uint8_t *__restrict__ mask, const int hidden_size) {
|
|
const __half2 scale = __float2half2_rn(1.f / (1.f - ratio));
|
|
__shared__ __half2 tile[8][129];
|
|
|
|
cg::thread_block b = cg::this_thread_block();
|
|
cg::thread_block_tile<WARP_SIZE> g = cg::tiled_partition<WARP_SIZE>(b);
|
|
|
|
__half2 *in_grad2 = reinterpret_cast<__half2 *>(in_grad);
|
|
const __half2 *out_grad2 = reinterpret_cast<const __half2 *>(out_grad);
|
|
__half2 *bias_grad2 = reinterpret_cast<__half2 *>(bias_grad);
|
|
|
|
int col_idx = flat_2dim(blockIdx.x, threadIdx.x, 8);
|
|
int stride = hidden_size * 128;
|
|
__half2 local_sum = __float2half2_rn(0.f);
|
|
|
|
int idx = flat_2dim(threadIdx.y, col_idx, hidden_size);
|
|
for (int r = threadIdx.y; r < row_size; r += 128) {
|
|
__half2 val = out_grad2[idx];
|
|
__half2 m2 = __floats2half2_rn(mask[2 * idx], mask[2 * idx + 1]);
|
|
val *= scale * m2;
|
|
local_sum += val;
|
|
in_grad2[idx] = val;
|
|
idx += stride;
|
|
}
|
|
|
|
tile[threadIdx.x][threadIdx.y] = local_sum;
|
|
__syncthreads();
|
|
|
|
__half2 sum = __float2half2_rn(0.f);
|
|
int tid = threadIdx.y * blockDim.x + threadIdx.x;
|
|
int x = tid >> 7;
|
|
int y = tid & (127);
|
|
if (y < 32) {
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; i++) {
|
|
sum += tile[x][y + i * 32];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
for (int i = 1; i < WARP_SIZE; i <<= 1) sum += g.shfl_down(sum, i);
|
|
|
|
if (y == 0) tile[0][x] = sum;
|
|
__syncthreads();
|
|
|
|
if (threadIdx.x < 8) {
|
|
int pos = flat_2dim(blockIdx.x, threadIdx.x, 8);
|
|
bias_grad2[pos] = tile[0][threadIdx.x];
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void launch_ls_dropout_bias_bwd(T *in_grad, T *bias_grad, const T *out_grad,
|
|
const uint8_t *mask, int row_size, int dim,
|
|
float ratio, cudaStream_t stream) {
|
|
dim3 grid_dim((dim - 1) / 8 + 1);
|
|
dim3 block_dim(8, 128);
|
|
ls_dropout_bias_bwd_kernel<<<grid_dim, block_dim, 0, stream>>>(
|
|
row_size, ratio, in_grad, bias_grad, out_grad, mask, dim);
|
|
}
|
|
|
|
template <>
|
|
void launch_ls_dropout_bias_bwd(__half *in_grad, __half *bias_grad,
|
|
const __half *out_grad, const uint8_t *mask,
|
|
int row_size, int dim, float ratio,
|
|
cudaStream_t stream) {
|
|
dim >>= 1;
|
|
dim3 grid_dim((dim - 1) / 8 + 1);
|
|
dim3 block_dim(8, 128);
|
|
ls_dropout_bias_bwd_kernel<<<grid_dim, block_dim, 0, stream>>>(
|
|
row_size, ratio, in_grad, bias_grad, out_grad, mask, dim);
|
|
}
|
|
|
|
template void launch_ls_dropout_bias_bwd(float *in_grad, float *bias_grad,
|
|
const float *out_grad,
|
|
const uint8_t *mask, int row_size,
|
|
int dim, float ratio,
|
|
cudaStream_t stream);
|
|
|
|
/**
|
|
* @brief fused bias, activation, and dropout at the end of first ffn
|
|
*
|
|
* @thread
|
|
* gridDim.x = hidden_size / 8
|
|
* blockDim.x = 8
|
|
* blockDim.y = 1024 / 8 = 128
|
|
*
|
|
* @tparam act_type activation function, like kRelu, kGelu
|
|
* @param total_count total elements
|
|
* @param ratio drop ratio
|
|
* @param out [batch_size, seq_len, hidden_size], float and __half
|
|
* @param in [batch_size, seq_len, hidden_size], float and __half
|
|
* @param mask [batch_size, seq_len, hidden_size], uint8 type
|
|
* @param bias [hidden_size], ffn bias
|
|
* @param seed seed to curand
|
|
* @param hidden_size
|
|
* @return void
|
|
*/
|
|
template <ActivationType act_type>
|
|
__global__ void ls_dropout_act_bias_kernel(
|
|
const int total_count, const float ratio, float *__restrict__ out,
|
|
const float *__restrict__ in, uint8_t *__restrict__ mask,
|
|
const float *__restrict__ bias, const int seed, const int hidden_size) {
|
|
const float scale = 1.f / (1.f - ratio);
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (i * 4 >= total_count) return;
|
|
|
|
curandStatePhilox4_32_10_t state;
|
|
curand_init(seed, i, 0, &state);
|
|
uint8_t m[4];
|
|
|
|
float4 *out4 = reinterpret_cast<float4 *>(out);
|
|
const float4 *data4 = reinterpret_cast<const float4 *>(in);
|
|
const float4 *bias4 = reinterpret_cast<const float4 *>(bias);
|
|
uint32_t *mask4 = reinterpret_cast<uint32_t *>(mask);
|
|
float4 rand = curand_uniform4(&state);
|
|
|
|
m[0] = (uint8_t)(rand.x > ratio);
|
|
m[1] = (uint8_t)(rand.y > ratio);
|
|
m[2] = (uint8_t)(rand.z > ratio);
|
|
m[3] = (uint8_t)(rand.w > ratio);
|
|
|
|
int bias_i = i % (hidden_size >> 2);
|
|
uint32_t *m4 = reinterpret_cast<uint32_t *>(m);
|
|
mask4[i] = m4[0];
|
|
const float4 input4 = data4[i];
|
|
const float4 b4 = __ldg(&bias4[bias_i]);
|
|
float4 output4;
|
|
|
|
output4.x =
|
|
activation_kernel<act_type, float>(input4.x + b4.x) * scale * m[0];
|
|
output4.y =
|
|
activation_kernel<act_type, float>(input4.y + b4.y) * scale * m[1];
|
|
output4.z =
|
|
activation_kernel<act_type, float>(input4.z + b4.z) * scale * m[2];
|
|
output4.w =
|
|
activation_kernel<act_type, float>(input4.w + b4.w) * scale * m[3];
|
|
|
|
out4[i] = output4;
|
|
}
|
|
|
|
template <ActivationType act_type>
|
|
__global__ void ls_dropout_act_bias_kernel(
|
|
const int total_count, const float ratio, __half *__restrict__ out,
|
|
const __half *__restrict__ in, uint8_t *__restrict__ mask,
|
|
const __half *__restrict__ bias, const int seed, const int hidden_size) {
|
|
const float scale = 1.f / (1.f - ratio);
|
|
|
|
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
if (i * 8 >= total_count) return;
|
|
|
|
curandStatePhilox4_32_10_t state;
|
|
curand_init(seed, i, 0, &state);
|
|
|
|
const float4 *vals_float4 = reinterpret_cast<const float4 *>(in);
|
|
float4 *outs_float4 = reinterpret_cast<float4 *>(out);
|
|
const float4 *bias4 = reinterpret_cast<const float4 *>(bias);
|
|
uint64_t *mask8 = reinterpret_cast<uint64_t *>(mask);
|
|
|
|
uint8_t m[8];
|
|
float4 rand = curand_uniform4(&state);
|
|
m[0] = (uint8_t)(rand.x > ratio);
|
|
m[1] = (uint8_t)(rand.y > ratio);
|
|
m[2] = (uint8_t)(rand.z > ratio);
|
|
m[3] = (uint8_t)(rand.w > ratio);
|
|
rand = curand_uniform4(&state);
|
|
m[4] = (uint8_t)(rand.x > ratio);
|
|
m[5] = (uint8_t)(rand.y > ratio);
|
|
m[6] = (uint8_t)(rand.z > ratio);
|
|
m[7] = (uint8_t)(rand.w > ratio);
|
|
uint64_t *m8 = reinterpret_cast<uint64_t *>(m);
|
|
mask8[i] = *m8;
|
|
|
|
int bias_i = i % (hidden_size >> 3);
|
|
float4 val_float4 = vals_float4[i];
|
|
const float4 b4 = __ldg(&bias4[bias_i]);
|
|
float4 out_float4;
|
|
|
|
__half2 *val_half2 = reinterpret_cast<__half2 *>(&val_float4);
|
|
__half2 *out_half2 = reinterpret_cast<__half2 *>(&out_float4);
|
|
const __half2 *b_half2 = reinterpret_cast<const __half2 *>(&b4);
|
|
|
|
__half2 scale_mask_1 = __floats2half2_rn(scale * m[0], scale * m[1]);
|
|
__half2 scale_mask_2 = __floats2half2_rn(scale * m[2], scale * m[3]);
|
|
__half2 scale_mask_3 = __floats2half2_rn(scale * m[4], scale * m[5]);
|
|
__half2 scale_mask_4 = __floats2half2_rn(scale * m[6], scale * m[7]);
|
|
out_half2[0] = __hmul2(
|
|
activation_kernel<act_type, __half2>(__hadd2(val_half2[0], b_half2[0])),
|
|
scale_mask_1);
|
|
out_half2[1] = __hmul2(
|
|
activation_kernel<act_type, __half2>(__hadd2(val_half2[1], b_half2[1])),
|
|
scale_mask_2);
|
|
out_half2[2] = __hmul2(
|
|
activation_kernel<act_type, __half2>(__hadd2(val_half2[2], b_half2[2])),
|
|
scale_mask_3);
|
|
out_half2[3] = __hmul2(
|
|
activation_kernel<act_type, __half2>(__hadd2(val_half2[3], b_half2[3])),
|
|
scale_mask_4);
|
|
outs_float4[i] = out_float4;
|
|
}
|
|
|
|
template <>
|
|
void launch_ls_dropout_act_bias<ActivationType::kGelu, float>(
|
|
float *out, const float *vals, uint8_t *mask, const float *bias,
|
|
int total_count, int dim, float ratio, cudaStream_t stream) {
|
|
int grid_dim = total_count >> 10;
|
|
ls_dropout_act_bias_kernel<ActivationType::kGelu>
|
|
<<<grid_dim + 1, 256, 0, stream>>>(
|
|
total_count, ratio, out, vals, mask, bias,
|
|
std::chrono::duration_cast<std::chrono::microseconds>(
|
|
std::chrono::system_clock::now().time_since_epoch())
|
|
.count(),
|
|
dim);
|
|
}
|
|
|
|
template <>
|
|
void launch_ls_dropout_act_bias<ActivationType::kGelu, __half>(
|
|
__half *out, const __half *vals, uint8_t *mask, const __half *bias,
|
|
int total_count, int dim, float ratio, cudaStream_t stream) {
|
|
int grid_dim = total_count >> 11;
|
|
ls_dropout_act_bias_kernel<ActivationType::kGelu>
|
|
<<<grid_dim + 1, 256, 0, stream>>>(
|
|
total_count, ratio, out, vals, mask, bias,
|
|
std::chrono::duration_cast<std::chrono::microseconds>(
|
|
std::chrono::system_clock::now().time_since_epoch())
|
|
.count(),
|
|
dim);
|
|
}
|
|
|
|
template <>
|
|
void launch_ls_dropout_act_bias<ActivationType::kRelu, float>(
|
|
float *out, const float *vals, uint8_t *mask, const float *bias,
|
|
int total_count, int dim, float ratio, cudaStream_t stream) {
|
|
int grid_dim = total_count >> 10;
|
|
ls_dropout_act_bias_kernel<ActivationType::kRelu>
|
|
<<<grid_dim + 1, 256, 0, stream>>>(
|
|
total_count, ratio, out, vals, mask, bias,
|
|
std::chrono::duration_cast<std::chrono::microseconds>(
|
|
std::chrono::system_clock::now().time_since_epoch())
|
|
.count(),
|
|
dim);
|
|
}
|
|
|
|
template <>
|
|
void launch_ls_dropout_act_bias<ActivationType::kRelu, __half>(
|
|
__half *out, const __half *vals, uint8_t *mask, const __half *bias,
|
|
int total_count, int dim, float ratio, cudaStream_t stream) {
|
|
int grid_dim = total_count >> 11;
|
|
ls_dropout_act_bias_kernel<ActivationType::kRelu>
|
|
<<<grid_dim + 1, 256, 0, stream>>>(
|
|
total_count, ratio, out, vals, mask, bias,
|
|
std::chrono::duration_cast<std::chrono::microseconds>(
|
|
std::chrono::system_clock::now().time_since_epoch())
|
|
.count(),
|
|
dim);
|
|
}
|
|
|
|
/**
|
|
* @brief fused bias, activation, and dropout backward
|
|
*
|
|
* @thread
|
|
* gridDim.x = total_count / 1024
|
|
* blockDim.x = 1024
|
|
*
|
|
* @tparam act_type kRelu
|
|
* @param row_size batch_size * seq_len
|
|
* @param ratio dropout ratio
|
|
* @param in_grad [batch_size, seq_len, hidden_size], input grad
|
|
* @param bias_grad [hidden_size], bias grad
|
|
* @param out_grad [batch_size, seq_len, hidden_size], output grad
|
|
* @param mask [batch_size, seq_len, hidden_size], dropout mask
|
|
* @param hidden_size
|
|
* @return void
|
|
*/
|
|
template <ActivationType act_type, typename T>
|
|
__global__ void ls_dropout_act_bias_bwd_kernel(
|
|
const int row_size, const float ratio, T *in_grad,
|
|
T *__restrict__ bias_grad, const T *__restrict__ input,
|
|
const T *__restrict__ bias, const T *out_grad,
|
|
const uint8_t *__restrict__ mask, const int hidden_size) {
|
|
const float scale = 1.f / (1.f - ratio);
|
|
__shared__ float tile[WARP_SIZE][WARP_SIZE + 1];
|
|
|
|
cg::thread_block b = cg::this_thread_block();
|
|
cg::thread_block_tile<WARP_SIZE> g = cg::tiled_partition<WARP_SIZE>(b);
|
|
|
|
int col_idx = flat_2dim(blockIdx.x, threadIdx.x, WARP_SIZE);
|
|
|
|
int stride = hidden_size * WARP_SIZE;
|
|
float local_sum = 0;
|
|
|
|
int idx = flat_2dim(threadIdx.y, col_idx, hidden_size);
|
|
if (col_idx < hidden_size) {
|
|
for (int r = threadIdx.y; r < row_size; r += WARP_SIZE) {
|
|
float val = out_grad[idx];
|
|
float in = input[idx];
|
|
float b = bias[idx % hidden_size];
|
|
val = activation_bwd_kernel<act_type, float>(
|
|
val * scale * static_cast<float>(mask[idx]), in + b);
|
|
local_sum += val;
|
|
in_grad[idx] = val;
|
|
idx += stride;
|
|
}
|
|
}
|
|
|
|
tile[threadIdx.x][threadIdx.y] = local_sum;
|
|
__syncthreads();
|
|
float sum = tile[threadIdx.y][threadIdx.x];
|
|
__syncthreads();
|
|
|
|
for (int i = 1; i < WARP_SIZE; i <<= 1) sum += g.shfl_down(sum, i);
|
|
|
|
if (threadIdx.x == 0) tile[0][threadIdx.y] = sum;
|
|
__syncthreads();
|
|
|
|
if (threadIdx.y == 0) {
|
|
int pos = flat_2dim(blockIdx.x, threadIdx.x, WARP_SIZE);
|
|
bias_grad[pos] = tile[0][threadIdx.x];
|
|
}
|
|
}
|
|
|
|
// @brief fused bias, activation, and dropout backward
|
|
// It is deprecated for precision reason. Keep it for future optimization.
|
|
//
|
|
// template <ActivationType act_type>
|
|
// __global__ void ls_dropout_act_bias_bwd_kernel(
|
|
// const int row_size, const float ratio, __half * in_grad,
|
|
// __half *__restrict__ bias_grad, const __half *__restrict__ input, const
|
|
// __half *__restrict__ bias, const __half * out_grad, const uint8_t
|
|
// *__restrict__ mask, const int hidden_size) {
|
|
// const __half2 scale = __float2half2_rn(1.f / (1.f - ratio));
|
|
// __shared__ __half2 tile[WARP_SIZE][WARP_SIZE + 1];
|
|
|
|
// cg::thread_block b = cg::this_thread_block();
|
|
// cg::thread_block_tile<WARP_SIZE> g = cg::tiled_partition<WARP_SIZE>(b);
|
|
|
|
// __half2 *in_grad2 = reinterpret_cast<__half2 *>(in_grad);
|
|
// __half2 *bias_grad2 = reinterpret_cast<__half2 *>(bias_grad);
|
|
// const __half2 *out_grad2 = reinterpret_cast<const __half2 *>(out_grad);
|
|
// const __half2 *input2 = reinterpret_cast<const __half2 *>(input);
|
|
// const __half2 *bias2 = reinterpret_cast<const __half2 *>(bias);
|
|
|
|
// int col_idx = flat_2dim(blockIdx.x, threadIdx.x, WARP_SIZE);
|
|
|
|
// int stride = hidden_size * WARP_SIZE;
|
|
// __half2 local_sum = __float2half2_rn(0.f);
|
|
|
|
// int idx = flat_2dim(threadIdx.y, col_idx, hidden_size);
|
|
// if (col_idx < hidden_size) {
|
|
// for (int r = threadIdx.y; r < row_size; r += WARP_SIZE) {
|
|
// __half2 val = out_grad2[idx];
|
|
// __half2 in2 = input2[idx];
|
|
// __half2 b2 = bias2[idx % hidden_size ];
|
|
// __half2 m2 = __floats2half2_rn(mask[2 * idx], mask[2 * idx + 1]);
|
|
// val = activation_bwd_kernel<ActivationType::kRelu, __half2>(val * scale
|
|
// *
|
|
// m2,
|
|
// in2+b2);
|
|
// local_sum += val;
|
|
// in_grad2[idx] = val;
|
|
// idx += stride;
|
|
// }
|
|
// }
|
|
|
|
// tile[threadIdx.x][threadIdx.y] = local_sum;
|
|
// __syncthreads();
|
|
// __half2 sum = tile[threadIdx.y][threadIdx.x];
|
|
// __syncthreads();
|
|
|
|
// for (int i = 1; i < WARP_SIZE; i <<= 1) sum += g.shfl_down(sum, i);
|
|
|
|
// if (threadIdx.x == 0) tile[0][threadIdx.y] = sum;
|
|
// __syncthreads();
|
|
|
|
// if (threadIdx.y == 0) {
|
|
// int pos = flat_2dim(blockIdx.x, threadIdx.x, WARP_SIZE);
|
|
// bias_grad2[pos] = tile[0][threadIdx.x];
|
|
// }
|
|
// }
|
|
|
|
template <ActivationType act_type, typename T>
|
|
void launch_ls_dropout_act_bias_bwd(T *in_grad, T *bias_grad, const T *input,
|
|
const T *bias, const T *out_grad,
|
|
const uint8_t *mask, int row_size, int dim,
|
|
float ratio, cudaStream_t stream) {
|
|
dim3 grid_dim((dim - 1) / WARP_SIZE + 1);
|
|
dim3 block_dim(WARP_SIZE, WARP_SIZE);
|
|
ls_dropout_act_bias_bwd_kernel<act_type><<<grid_dim, block_dim, 0, stream>>>(
|
|
row_size, ratio, in_grad, bias_grad, input, bias, out_grad, mask, dim);
|
|
}
|
|
|
|
// template <>
|
|
// void launch_ls_dropout_act_bias_bwd<ActivationType::kRelu, __half>(
|
|
// __half *in_grad, __half *bias_grad,const __half *input, const __half
|
|
// *bias, const __half *out_grad, const uint8_t *mask, int row_size, int
|
|
// dim, float ratio, cudaStream_t stream) {
|
|
// dim >>= 1;
|
|
// dim3 grid_dim((dim - 1) / WARP_SIZE + 1);
|
|
// dim3 block_dim(WARP_SIZE, WARP_SIZE);
|
|
// ls_dropout_act_bias_bwd_kernel<ActivationType::kRelu>
|
|
// <<<grid_dim, block_dim, 0, stream>>>(row_size, ratio, in_grad,
|
|
// bias_grad,
|
|
// input, bias,out_grad, mask, dim);
|
|
// }
|
|
|
|
template void launch_ls_dropout_act_bias_bwd<ActivationType::kRelu, float>(
|
|
float *in_grad, float *bias_grad, const float *input, const float *bias,
|
|
const float *out_grad, const uint8_t *mask, int row_size, int dim,
|
|
float ratio, cudaStream_t stream);
|
|
|
|
template void launch_ls_dropout_act_bias_bwd<ActivationType::kRelu, __half>(
|
|
__half *in_grad, __half *bias_grad, const __half *input, const __half *bias,
|
|
const __half *out_grad, const uint8_t *mask, int row_size, int dim,
|
|
float ratio, cudaStream_t stream);
|
|
|
|
template void launch_ls_dropout_act_bias_bwd<ActivationType::kGelu, float>(
|
|
float *in_grad, float *bias_grad, const float *input, const float *bias,
|
|
const float *out_grad, const uint8_t *mask, int row_size, int dim,
|
|
float ratio, cudaStream_t stream);
|
|
|
|
template void launch_ls_dropout_act_bias_bwd<ActivationType::kGelu, __half>(
|
|
__half *in_grad, __half *bias_grad, const __half *input, const __half *bias,
|
|
const __half *out_grad, const uint8_t *mask, int row_size, int dim,
|
|
float ratio, cudaStream_t stream);
|