Skip to content
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

Supporting different hidden dimensions for transformer kernels #934

Merged
merged 4 commits into from
Apr 8, 2021
Merged

Supporting different hidden dimensions for transformer kernels #934

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a080819
Supporting different hidden dimensions for transformer kernels-v2
Mar 29, 2021
ab488b5
Merge branch 'master' into transformer-kernel/all-hidden-dimensions
RezaYazdaniAminabadi Apr 7, 2021
f2f7457
Merge branch 'master' into transformer-kernel/all-hidden-dimensions
jeffra Apr 7, 2021
2550624
Merge branch 'master' into transformer-kernel/all-hidden-dimensions
jeffra Apr 8, 2021
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
1 change: 1 addition & 0 deletions csrc/transformer/ds_transformer_cuda.cpp
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ size_t get_workspace_size(int maxBatchSize,
{
size_t workSpacesize = 4 * (size_t(maxBatchSize) * seq_len * hidden_size);
if (training) {
workSpacesize += 2 * (size_t(maxBatchSize) * seq_len * hidden_size);
workSpacesize += ((std::max)((size_t(maxBatchSize) * seq_len * intermediate_size),
2 * (size_t(maxBatchSize) * heads * seq_len * seq_len)));
if (gelu_checkpoint)
Expand Down
43 changes: 19 additions & 24 deletions csrc/transformer/gelu_kernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,11 @@ the intermediate hyperbolic tangent, since there's no intrinsic
that computes it directly.
*/

__global__ void gelu_kernel(const float* input, float* vals, int intermediate_size)
__global__ void gelu_kernel(const float* input, float* vals, int row_stride, int iterations)
{
int row = blockIdx.x;
int id = threadIdx.x;
int loop_stride = blockDim.x;
int iterations = intermediate_size / blockDim.x / 4;
int row_stride = intermediate_size / 4;

const float4* input_cast = reinterpret_cast<const float4*>(input);
float4* vals_cast = reinterpret_cast<float4*>(vals);
Expand All @@ -60,14 +58,12 @@ __global__ void gelu_kernel(const float* input, float* vals, int intermediate_si
}
}

__global__ void gelu_kernel(const __half* input, __half* vals, int intermediate_size)
__global__ void gelu_kernel(const __half* input, __half* vals, int row_stride, int iterations)
{
#if __CUDA_ARCH__ >= 700
int row = blockIdx.x;
int id = threadIdx.x;
int loop_stride = blockDim.x;
int iterations = intermediate_size / blockDim.x / 4;
int row_stride = intermediate_size / 4;

const float2* input_cast = reinterpret_cast<const float2*>(input);
float2* vals_cast = reinterpret_cast<float2*>(vals);
Expand Down Expand Up @@ -98,13 +94,12 @@ __global__ void gelu_kernel(const __half* input, __half* vals, int intermediate_
__global__ void fused_bias_gelu(const float* input,
const float* bias,
float* vals,
int intermediate_size)
int row_stride,
int iterations)
{
int row = blockIdx.x;
int id = threadIdx.x;
int loop_stride = blockDim.x;
int iterations = intermediate_size / blockDim.x / 4;
int row_stride = intermediate_size / 4;

const float4* input_cast = reinterpret_cast<const float4*>(input);
float4* vals_cast = reinterpret_cast<float4*>(vals);
Expand Down Expand Up @@ -133,14 +128,13 @@ __global__ void fused_bias_gelu(const float* input,
__global__ void fused_bias_gelu(const __half* input,
const __half* bias,
__half* vals,
int intermediate_size)
int row_stride,
int iterations)
{
#if __CUDA_ARCH__ >= 700
int row = blockIdx.x;
int id = threadIdx.x;
int loop_stride = blockDim.x;
int iterations = intermediate_size / blockDim.x / 4;
int row_stride = intermediate_size / 4;

const float2* input_cast = reinterpret_cast<const float2*>(input);
float2* vals_cast = reinterpret_cast<float2*>(vals);
Expand Down Expand Up @@ -182,13 +176,12 @@ __global__ void fused_bias_gelu(const __half* input,
__global__ void d_gelu_func(float* d_output,
const float* gelu_input,
const float* bias,
int intermediate_size)
int row_stride,
int iterations)
{
int row = blockIdx.x;
int id = threadIdx.x;
int loop_stride = blockDim.x;
int iterations = intermediate_size / blockDim.x / 4;
int row_stride = intermediate_size / 4;

float4* d_output_cast = reinterpret_cast<float4*>(d_output);
const float4* gelu_input_cast = reinterpret_cast<const float4*>(gelu_input);
Expand Down Expand Up @@ -218,14 +211,13 @@ __global__ void d_gelu_func(float* d_output,
__global__ void d_gelu_func(__half* d_output,
const __half* gelu_input,
const __half* bias,
int intermediate_size)
int row_stride,
int iterations)
{
#if __CUDA_ARCH__ >= 700
int row = blockIdx.x;
int id = threadIdx.x;
int loop_stride = blockDim.x;
int iterations = intermediate_size / blockDim.x / 4;
int row_stride = intermediate_size / 4;

float2* d_output_cast = reinterpret_cast<float2*>(d_output);
const float2* gelu_input_cast = reinterpret_cast<const float2*>(gelu_input);
Expand Down Expand Up @@ -282,11 +274,12 @@ void launch_bias_gelu(const T* input,
cudaStream_t stream)
{
int iterations = (intermediate_size + 1023) / 1024;
int threads = intermediate_size / iterations / 4;
int threads = (intermediate_size - 1) / (iterations * 4) + 1;
dim3 block_dims(threads);
dim3 grid_dims(batch_size);

fused_bias_gelu<<<grid_dims, block_dims, 0, stream>>>(input, bias, output, intermediate_size);
fused_bias_gelu<<<grid_dims, block_dims, 0, stream>>>(
input, bias, output, intermediate_size / 4, iterations);
}

template <typename T>
Expand All @@ -297,11 +290,12 @@ void launch_gelu(const T* input,
cudaStream_t stream)
{
int iterations = (intermediate_size + 1023) / 1024;
int threads = intermediate_size / iterations / 4;
int threads = (intermediate_size - 1) / (iterations * 4) + 1;
dim3 block_dims(threads);
dim3 grid_dims(batch_size);

gelu_kernel<<<grid_dims, block_dims, 0, stream>>>(input, output, intermediate_size);
gelu_kernel<<<grid_dims, block_dims, 0, stream>>>(
input, output, intermediate_size / 4, iterations);
}

template void launch_bias_gelu<float>(const float*, const float*, float*, int, int, cudaStream_t);
Expand All @@ -324,11 +318,12 @@ void launch_d_gelu(T* d_output,
cudaStream_t stream)
{
int iterations = (intermediate_size + 1023) / 1024;
int threads = intermediate_size / iterations / 4;
int threads = (intermediate_size - 1) / (iterations * 4) + 1;
dim3 block_dims(threads);
dim3 grid_dims(batch_size);

d_gelu_func<<<grid_dims, block_dims, 0, stream>>>(d_output, input, bias, intermediate_size);
d_gelu_func<<<grid_dims, block_dims, 0, stream>>>(
d_output, input, bias, intermediate_size / 4, iterations);
}

template void launch_d_gelu<float>(float*, const float*, const float*, int, int, cudaStream_t);
Expand Down
11 changes: 7 additions & 4 deletions tests/unit/test_cuda_backward.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,9 @@
import sys

#if not deepspeed.ops.__installed_ops__['transformer']:
pytest.skip(
"transformer kernels are temporarily disabled because of unexplained failures",
allow_module_level=True)
#pytest.skip(
# "transformer kernels are temporarily disabled because of unexplained failures",
# allow_module_level=True)


def check_equal(first, second, atol=1e-2, verbose=False):
Expand Down Expand Up @@ -258,6 +258,9 @@ 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',
[
(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),
(1024,128,10,2,2,False,False, 0.1),
Expand Down Expand Up @@ -291,7 +294,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)
run_backward(ds_config, seq_len, atol=atol, verbose=False)


#@pytest.mark.parametrize('batch_size, hidden_size, seq_len, heads, num_layers, is_preln, use_fp16, atol',
Expand Down
22 changes: 13 additions & 9 deletions tests/unit/test_cuda_forward.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -199,7 +199,11 @@ 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',
[
(8,256,53,4,3,True,False),
(8,160,128,2,3,True,True),
(8,160,128,2,3,False,True),
(8,1600,128,2,3,True,True),
(8,1600,128,25,3,True,True),
(8,1600,128,25,3,False,True),
(8,256,52,4,3,True,True),
(3,1024,51,16,3,True,False),
(3,1024,54,16,3,True,True),
Expand Down Expand Up @@ -259,10 +263,10 @@ def test_forward(batch_size,

@pytest.mark.parametrize('batch_size, small_bsz, hidden_size, seq_len, heads, num_layers, is_preln, use_fp16',
[
(8,3,1024,512,16,3,True,False),
(8,7,1024,512,16,3,True,True),
(8,3,1024,512,16,3,False,False),
(8,7,1024,512,16,3,False,True),
#(8,3,1024,512,16,3,True,False),
#(8,7,1024,512,16,3,True,True),
#(8,3,1024,512,16,3,False,False),
#(8,7,1024,512,16,3,False,True),
]) # yapf: disable
def test_forward_with_small_bsz(batch_size,
small_bsz,
Expand Down Expand Up @@ -294,10 +298,10 @@ def test_forward_with_small_bsz(batch_size,

@pytest.mark.parametrize('batch_size, hidden_size, seq_len, heads, num_layers, is_preln, use_fp16',
[
(64,1024,128,16,3,True,False),
(64,1024,128,16,3,True,True),
(64,1024,128,16,3,False,False),
(64,1024,128,16,3,False,True),
#(64,1024,128,16,3,True,False),
#(64,1024,128,16,3,True,True),
#(64,1024,128,16,3,False,False),
#(64,1024,128,16,3,False,True),
]) # yapf: disable
def test_forward_stochastic(batch_size,
hidden_size,
Expand Down