зеркало из https://github.com/microsoft/DeepSpeed.git
Transformer kernel/fix layer norm (#1587)
* fixing the softmax masking when using triangular masking * fix a bug in the the layernorm backward kernels * revert back some changes & remove debug code * change the constants to a macro Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
This commit is contained in:
Reza Yazdani
GitHub
Родитель
fc2f378ece
Коммит
8e891aa568
|
|
@ -35,6 +35,8 @@
|
|||
|
||||
#define MAX_REG 256
|
||||
|
||||
#define WARP_SIZE_BITS 5
|
||||
|
||||
template <typename T>
|
||||
void launch_quantize_kernel(T* vals,
|
||||
int total_count,
|
||||
|
|
|
|||
|
|
@ -59,13 +59,15 @@ __global__ void fused_bias_residual_layer_norm(float* vals,
|
|||
|
||||
b.sync();
|
||||
|
||||
if (g.thread_rank() < (iteration_stride >> 5)) sum = shr[g.thread_rank()];
|
||||
if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) sum = shr[g.thread_rank()];
|
||||
|
||||
#if !defined(__STOCHASTIC_MODE__) || __CUDA_ARCH__ < 700
|
||||
b.sync();
|
||||
#endif
|
||||
|
||||
for (int i = 1; i < (iteration_stride >> 5); i *= 2) { sum += g.shfl_down(sum, i); }
|
||||
for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) {
|
||||
sum += g.shfl_down(sum, i);
|
||||
}
|
||||
|
||||
sum = g.shfl(sum, 0);
|
||||
float mean = sum / row_stride;
|
||||
|
|
@ -83,13 +85,15 @@ __global__ void fused_bias_residual_layer_norm(float* vals,
|
|||
|
||||
b.sync();
|
||||
|
||||
if (g.thread_rank() < (iteration_stride >> 5)) variance = shr[g.thread_rank()];
|
||||
if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) variance = shr[g.thread_rank()];
|
||||
|
||||
#ifndef __STOCHASTIC_MODE__
|
||||
b.sync();
|
||||
#endif
|
||||
|
||||
for (int i = 1; i < (iteration_stride >> 5); i *= 2) { variance += g.shfl_down(variance, i); }
|
||||
for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) {
|
||||
variance += g.shfl_down(variance, i);
|
||||
}
|
||||
variance = g.shfl(variance, 0);
|
||||
variance /= row_stride;
|
||||
variance += epsilon;
|
||||
|
|
@ -130,7 +134,7 @@ __global__ void fused_bias_residual_layer_norm(__half* vals,
|
|||
|
||||
int row = blockIdx.x;
|
||||
int id = threadIdx.x;
|
||||
int gid = id >> 5;
|
||||
int gid = id >> WARP_SIZE_BITS;
|
||||
|
||||
float2 vals_f[NORM_REG];
|
||||
__shared__ float shr[MAX_WARP_NUM];
|
||||
|
|
@ -162,13 +166,15 @@ __global__ void fused_bias_residual_layer_norm(__half* vals,
|
|||
|
||||
b.sync();
|
||||
|
||||
if (g.thread_rank() < (iteration_stride >> 5)) sum = shr[g.thread_rank()];
|
||||
if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) sum = shr[g.thread_rank()];
|
||||
|
||||
#ifndef __STOCHASTIC_MODE__
|
||||
b.sync();
|
||||
#endif
|
||||
|
||||
for (int i = 1; i < (iteration_stride >> 5); i *= 2) { sum += g.shfl_down(sum, i); }
|
||||
for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) {
|
||||
sum += g.shfl_down(sum, i);
|
||||
}
|
||||
sum = g.shfl(sum, 0);
|
||||
float mean = sum / (row_stride * 2);
|
||||
|
||||
|
|
@ -186,13 +192,15 @@ __global__ void fused_bias_residual_layer_norm(__half* vals,
|
|||
|
||||
b.sync();
|
||||
|
||||
if (g.thread_rank() < (iteration_stride >> 5)) variance = shr[g.thread_rank()];
|
||||
if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) variance = shr[g.thread_rank()];
|
||||
|
||||
#ifndef __STOCHASTIC_MODE__
|
||||
b.sync();
|
||||
#endif
|
||||
|
||||
for (int i = 1; i < (iteration_stride >> 5); i *= 2) { variance += g.shfl_down(variance, i); }
|
||||
for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) {
|
||||
variance += g.shfl_down(variance, i);
|
||||
}
|
||||
variance = g.shfl(variance, 0);
|
||||
variance /= (row_stride * 2);
|
||||
variance += epsilon;
|
||||
|
|
@ -345,13 +353,15 @@ __global__ void fused_bias_residual_layer_norm(float* vals,
|
|||
|
||||
b.sync();
|
||||
|
||||
if (g.thread_rank() < (iteration_stride >> 5)) sum = shr[g.thread_rank()];
|
||||
if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) sum = shr[g.thread_rank()];
|
||||
|
||||
#if !defined(__STOCHASTIC_MODE__) || __CUDA_ARCH__ < 700
|
||||
b.sync();
|
||||
#endif
|
||||
|
||||
for (int i = 1; i < (iteration_stride >> 5); i *= 2) { sum += g.shfl_down(sum, i); }
|
||||
for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) {
|
||||
sum += g.shfl_down(sum, i);
|
||||
}
|
||||
|
||||
sum = g.shfl(sum, 0);
|
||||
float mean = sum / row_stride;
|
||||
|
|
@ -367,13 +377,15 @@ __global__ void fused_bias_residual_layer_norm(float* vals,
|
|||
|
||||
b.sync();
|
||||
|
||||
if (g.thread_rank() < (iteration_stride >> 5)) variance = shr[g.thread_rank()];
|
||||
if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) variance = shr[g.thread_rank()];
|
||||
|
||||
#ifndef __STOCHASTIC_MODE__
|
||||
b.sync();
|
||||
#endif
|
||||
|
||||
for (int i = 1; i < (iteration_stride >> 5); i *= 2) { variance += g.shfl_down(variance, i); }
|
||||
for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) {
|
||||
variance += g.shfl_down(variance, i);
|
||||
}
|
||||
variance = g.shfl(variance, 0);
|
||||
variance /= row_stride;
|
||||
variance += epsilon;
|
||||
|
|
@ -414,7 +426,7 @@ __global__ void fused_bias_residual_layer_norm(__half* vals,
|
|||
|
||||
int row = blockIdx.x;
|
||||
int id = threadIdx.x;
|
||||
int gid = id >> 5;
|
||||
int gid = id >> WARP_SIZE_BITS;
|
||||
|
||||
float2 vals_f[NORM_REG];
|
||||
__shared__ float shr[MAX_WARP_NUM];
|
||||
|
|
@ -446,13 +458,15 @@ __global__ void fused_bias_residual_layer_norm(__half* vals,
|
|||
|
||||
b.sync();
|
||||
|
||||
if (g.thread_rank() < (iteration_stride >> 5)) sum = shr[g.thread_rank()];
|
||||
if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) sum = shr[g.thread_rank()];
|
||||
|
||||
#ifndef __STOCHASTIC_MODE__
|
||||
b.sync();
|
||||
#endif
|
||||
|
||||
for (int i = 1; i < (iteration_stride >> 5); i *= 2) { sum += g.shfl_down(sum, i); }
|
||||
for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) {
|
||||
sum += g.shfl_down(sum, i);
|
||||
}
|
||||
sum = g.shfl(sum, 0);
|
||||
float mean = sum / (row_stride * 2);
|
||||
|
||||
|
|
@ -470,13 +484,15 @@ __global__ void fused_bias_residual_layer_norm(__half* vals,
|
|||
|
||||
b.sync();
|
||||
|
||||
if (g.thread_rank() < (iteration_stride >> 5)) variance = shr[g.thread_rank()];
|
||||
if (g.thread_rank() < (iteration_stride >> WARP_SIZE_BITS)) variance = shr[g.thread_rank()];
|
||||
|
||||
#ifndef __STOCHASTIC_MODE__
|
||||
b.sync();
|
||||
#endif
|
||||
|
||||
for (int i = 1; i < (iteration_stride >> 5); i *= 2) { variance += g.shfl_down(variance, i); }
|
||||
for (int i = 1; i < (iteration_stride >> WARP_SIZE_BITS); i *= 2) {
|
||||
variance += g.shfl_down(variance, i);
|
||||
}
|
||||
variance = g.shfl(variance, 0);
|
||||
variance /= (row_stride * 2);
|
||||
variance += epsilon;
|
||||
|
|
@ -755,7 +771,7 @@ __global__ void LayerNormBackward2(const float* out_grad,
|
|||
int row = blockIdx.x;
|
||||
int id = threadIdx.x;
|
||||
int wid = id / WARP_SIZE;
|
||||
int warp_num = (THREADS < row_stride ? THREADS : row_stride) / WARP_SIZE;
|
||||
int warp_num = iteration_stride >> WARP_SIZE_BITS;
|
||||
__shared__ float partialSum[MAX_WARP_NUM];
|
||||
|
||||
out_grad += (row * row_stride);
|
||||
|
|
@ -855,7 +871,7 @@ __global__ void LayerNormBackward2(const __half* out_grad,
|
|||
int row = blockIdx.x;
|
||||
int id = threadIdx.x;
|
||||
int wid = id / WARP_SIZE;
|
||||
int warp_num = (iteration_stride < row_stride ? iteration_stride : row_stride) / WARP_SIZE;
|
||||
int warp_num = iteration_stride >> WARP_SIZE_BITS;
|
||||
__shared__ float partialSum[MAX_WARP_NUM];
|
||||
|
||||
__half2 vals_arr[NORM_REG];
|
||||
|
|
@ -1027,8 +1043,8 @@ void launch_layerNorm_backward<__half>(const __half* out_grad,
|
|||
dim3 grid_dim(hidden_dim / TILE_DIM);
|
||||
dim3 block_dim(TILE_DIM, TILE_DIM);
|
||||
|
||||
LayerNormBackward1<__half><<<grid_dim, block_dim, 0, stream[0]>>>(
|
||||
out_grad, vals_hat, gamma, betta, gamma_grad, betta_grad, batch, hidden_dim, invertible);
|
||||
// LayerNormBackward1<__half><<<grid_dim, block_dim, 0, stream[0]>>>(
|
||||
// out_grad, vals_hat, gamma, betta, gamma_grad, betta_grad, batch, hidden_dim, invertible);
|
||||
|
||||
dim3 grid_dim2(batch);
|
||||
|
||||
|
|
@ -1069,8 +1085,8 @@ __global__ void LayerNormBackward2(const float* out_grad,
|
|||
|
||||
int row = blockIdx.x;
|
||||
int id = threadIdx.x;
|
||||
int wid = id / WARP_SIZE;
|
||||
int warp_num = (THREADS < row_stride ? THREADS : row_stride) / WARP_SIZE;
|
||||
int wid = id >> WARP_SIZE_BITS;
|
||||
int warp_num = iteration_stride >> WARP_SIZE_BITS;
|
||||
__shared__ float partialSum[MAX_WARP_NUM];
|
||||
|
||||
out_grad += (row * row_stride);
|
||||
|
|
@ -1164,13 +1180,14 @@ __global__ void LayerNormBackward2(const __half* out_grad,
|
|||
|
||||
int row = blockIdx.x;
|
||||
int id = threadIdx.x;
|
||||
int wid = id / WARP_SIZE;
|
||||
int warp_num = (iteration_stride < row_stride ? iteration_stride : row_stride) / WARP_SIZE;
|
||||
int wid = id >> WARP_SIZE_BITS;
|
||||
int warp_num = iteration_stride >> WARP_SIZE_BITS;
|
||||
|
||||
__shared__ float partialSum[MAX_WARP_NUM];
|
||||
|
||||
__half2 vals_arr[NORM_REG];
|
||||
float2 vals_arr_f[NORM_REG];
|
||||
__half2 xu[NORM_REG];
|
||||
|
||||
__half2* inp_grad_h = reinterpret_cast<__half2*>(inp_grad);
|
||||
const __half2* out_grad_h = reinterpret_cast<const __half2*>(out_grad);
|
||||
|
|
@ -1182,27 +1199,28 @@ __global__ void LayerNormBackward2(const __half* out_grad,
|
|||
|
||||
const __half2* gamma_h = reinterpret_cast<const __half2*>(gamma);
|
||||
int high_index = iterations * iteration_stride + id;
|
||||
|
||||
__half mean_h = means[row];
|
||||
__half2 mean_reg = __halves2half2(mean_h, mean_h);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < iterations; i++) {
|
||||
__half2 gamma_reg = gamma_h[i * iteration_stride + id];
|
||||
vals_arr[i] = out_grad_h[i * iteration_stride + id];
|
||||
vals_arr[i] *= gamma_reg; // out_grad * gamma
|
||||
xu[i] = (vals_hat_h[i * iteration_stride + id] - mean_reg);
|
||||
}
|
||||
if ((high_index) < row_stride) {
|
||||
__half2 gamma_reg = gamma_h[high_index];
|
||||
vals_arr[iterations] = out_grad_h[high_index];
|
||||
vals_arr[iterations] *= gamma_reg; // out_grad * gamma
|
||||
xu[iterations] = (vals_hat_h[high_index] - mean_reg);
|
||||
iterations++;
|
||||
}
|
||||
__half mean_h = means[row];
|
||||
__half var_h = vars[row];
|
||||
__half2 var_reg = __halves2half2(var_h, var_h);
|
||||
__half2 mean_reg = __halves2half2(mean_h, mean_h);
|
||||
__half2 xu[NORM_REG];
|
||||
|
||||
float sum = 0.f;
|
||||
for (int i = 0; i < iterations; i++) {
|
||||
xu[i] = (vals_hat_h[i * iteration_stride + id] - mean_reg);
|
||||
__half2 result_h = (xu[i] * vals_arr[i]);
|
||||
float2 result_f = __half22float2(result_h);
|
||||
sum += result_f.x;
|
||||
|
|
@ -1488,7 +1506,7 @@ __global__ void LayerNormBackward2_fused_add(const float* out_grad1,
|
|||
int row = blockIdx.x;
|
||||
int id = threadIdx.x;
|
||||
int wid = id / WARP_SIZE;
|
||||
int warp_num = (THREADS < row_stride ? THREADS : row_stride) / WARP_SIZE;
|
||||
int warp_num = iteration_stride >> WARP_SIZE_BITS;
|
||||
__shared__ float partialSum[MAX_WARP_NUM];
|
||||
|
||||
out_grad1 += (row * row_stride);
|
||||
|
|
@ -1592,7 +1610,7 @@ __global__ void LayerNormBackward2_fused_add(const __half* out_grad1,
|
|||
int row = blockIdx.x;
|
||||
int id = threadIdx.x;
|
||||
int wid = id / WARP_SIZE;
|
||||
int warp_num = (iteration_stride < row_stride ? iteration_stride : row_stride) / WARP_SIZE;
|
||||
int warp_num = iteration_stride >> WARP_SIZE_BITS;
|
||||
__shared__ float partialSum[MAX_WARP_NUM];
|
||||
|
||||
__half2 vals_arr[NORM_REG];
|
||||
|
|
@ -1810,7 +1828,7 @@ __global__ void LayerNormBackward2_fused_add(const float* out_grad1,
|
|||
int row = blockIdx.x;
|
||||
int id = threadIdx.x;
|
||||
int wid = id / WARP_SIZE;
|
||||
int warp_num = (THREADS < row_stride ? THREADS : row_stride) / WARP_SIZE;
|
||||
int warp_num = iteration_stride >> WARP_SIZE_BITS;
|
||||
__shared__ float partialSum[MAX_WARP_NUM];
|
||||
|
||||
float vals_arr[NORM_REG];
|
||||
|
|
@ -1913,7 +1931,7 @@ __global__ void LayerNormBackward2_fused_add(const __half* out_grad1,
|
|||
int row = blockIdx.x;
|
||||
int id = threadIdx.x;
|
||||
int wid = id / WARP_SIZE;
|
||||
int warp_num = (iteration_stride < row_stride ? iteration_stride : row_stride) / WARP_SIZE;
|
||||
int warp_num = iteration_stride >> WARP_SIZE_BITS;
|
||||
|
||||
__shared__ float partialSum[MAX_WARP_NUM];
|
||||
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ __global__ void attn_softmax(float* vals,
|
|||
{
|
||||
__shared__ float partialSum[MAX_WARP_NUM];
|
||||
|
||||
int warp_num = blockDim.x >> 5;
|
||||
int warp_num = blockDim.x >> WARP_SIZE_BITS;
|
||||
|
||||
int iteration_stride = blockDim.x;
|
||||
int block_width = blockStride * seq_length;
|
||||
|
|
@ -45,7 +45,7 @@ __global__ void attn_softmax(float* vals,
|
|||
(threadIdx.x / max_threads_in_sequence) * seq_length;
|
||||
int mask_offset = batch * seq_length;
|
||||
|
||||
int wid = threadIdx.x >> 5;
|
||||
int wid = threadIdx.x >> WARP_SIZE_BITS;
|
||||
int lane = threadIdx.x & 0x1f;
|
||||
|
||||
float4* val_cast = reinterpret_cast<float4*>(vals);
|
||||
|
|
@ -159,7 +159,7 @@ __global__ void attn_softmax(__half* vals,
|
|||
#if __CUDA_ARCH__ >= 700
|
||||
__shared__ float partialSum[MAX_WARP_NUM];
|
||||
|
||||
int warp_num = blockDim.x >> 5;
|
||||
int warp_num = blockDim.x >> WARP_SIZE_BITS;
|
||||
|
||||
int iteration_stride = blockDim.x;
|
||||
int block_width = blockStride * seq_length;
|
||||
|
|
@ -176,7 +176,7 @@ __global__ void attn_softmax(__half* vals,
|
|||
(threadIdx.x / max_threads_in_sequence) * seq_length;
|
||||
int mask_offset = batch * seq_length;
|
||||
|
||||
int wid = threadIdx.x >> 5;
|
||||
int wid = threadIdx.x >> WARP_SIZE_BITS;
|
||||
int lane = threadIdx.x & 0x1f;
|
||||
|
||||
float2* val_cast = reinterpret_cast<float2*>(vals);
|
||||
|
|
@ -439,7 +439,7 @@ __global__ void softmax_backward_kernel(T* out_grad, const T* soft_inp, int seq_
|
|||
{
|
||||
__shared__ float partialSum[MAX_WARP_NUM];
|
||||
|
||||
int warp_num = blockDim.x >> 5; // warp-count = num_threads / WARP_SIZE (32)
|
||||
int warp_num = blockDim.x >> WARP_SIZE_BITS; // warp-count = num_threads / WARP_SIZE (32)
|
||||
|
||||
int iteration_stride = blockDim.x;
|
||||
int block_width = blockStride * seq_length;
|
||||
|
|
@ -454,7 +454,7 @@ __global__ void softmax_backward_kernel(T* out_grad, const T* soft_inp, int seq_
|
|||
int row = blockIdx.x;
|
||||
int id = threadIdx.x;
|
||||
|
||||
int wid = id >> 5;
|
||||
int wid = id >> WARP_SIZE_BITS;
|
||||
int lane = id & 0x1f;
|
||||
|
||||
T val_reg[MAX_THREAD_ITERATIONS];
|
||||
|
|
|
|||
|
|
@ -76,7 +76,7 @@ class InferenceEngine(Module):
|
|||
elif self.mp_world_size > 1:
|
||||
self._create_model_parallel_group()
|
||||
# apply injection policy
|
||||
if self.injection_dict:
|
||||
if self.injection_dict is not None:
|
||||
for client_module, injection_policy in self.injection_dict.items():
|
||||
self._apply_injection_policy(client_module,
|
||||
injection_policy,
|
||||
|
|
|
|||
|
|
@ -268,12 +268,13 @@ def run_backward(ds_config, seq_len, atol=1e-2, verbose=False):
|
|||
# 3-128-54-2-24-False-True-0.2
|
||||
@pytest.mark.parametrize('batch_size, hidden_size, seq_len, heads, num_layers, is_preln, use_fp16, atol',
|
||||
[
|
||||
(64,160,128,2,24,False,True, 0.2),
|
||||
(64,1600,128,2,4,False,True, 0.2),
|
||||
(8,1600,128,25,3,True,True, 0.05),
|
||||
(8,160,128,2,3,True,True, 0.1),
|
||||
(8,1600,128,2,3,True,True, 0.05),
|
||||
(3,1024,119,16,24,True,False, 0.05),
|
||||
(3,1024,115,16,24,True,True, 0.05),
|
||||
#(3,1024,119,16,24,True,False, 0.05),
|
||||
#(3,1024,115,16,24,True,True, 0.05),
|
||||
#(1024,128,10,2,2,False,False, 0.1),
|
||||
#(3,1024,52,16,24,False,True, 0.2),
|
||||
#(3,128,51,2,24,False,False, 0.1),
|
||||
|
|
@ -305,7 +306,7 @@ def test_backward(batch_size,
|
|||
ds_config.initializer_range = 0.02
|
||||
ds_config.fp16 = use_fp16
|
||||
|
||||
run_backward(ds_config, seq_len, atol=atol, verbose=False)
|
||||
run_backward(ds_config, seq_len, atol=atol, verbose=True)
|
||||
|
||||
|
||||
#@pytest.mark.parametrize('batch_size, hidden_size, seq_len, heads, num_layers, is_preln, use_fp16, atol',
|
||||
|
|
|
|||
|
|
@ -199,6 +199,7 @@ def run_forward(ds_config, seq_len, atol=1e-2, verbose=False, test_bsz=None):
|
|||
# FP16 test cases can only run on the devices support FP16.
|
||||
@pytest.mark.parametrize('batch_size, hidden_size, seq_len, heads, num_layers, is_preln, use_fp16',
|
||||
[
|
||||
(64,160,128,2,24,False,True),
|
||||
#(8,2048,2048,32,1,True,True),
|
||||
(8,160,128,2,3,True,True),
|
||||
(8,160,128,2,3,False,True),
|
||||
|
|
|
|||
Загрузка…
Ссылка в новой задаче