diff --git a/csrc/transformer/ds_transformer_cuda.cpp b/csrc/transformer/ds_transformer_cuda.cpp
old mode 100644
new mode 100755
index 8e605d7558d5..6844fa563705
--- a/csrc/transformer/ds_transformer_cuda.cpp
+++ b/csrc/transformer/ds_transformer_cuda.cpp
@@ -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)
diff --git a/csrc/transformer/gelu_kernels.cu b/csrc/transformer/gelu_kernels.cu
index 209b64a90902..12048006266e 100644
--- a/csrc/transformer/gelu_kernels.cu
+++ b/csrc/transformer/gelu_kernels.cu
@@ -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);
@@ -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);
@@ -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);
@@ -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);
@@ -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);
@@ -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);
@@ -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>
@@ -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);
@@ -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);
diff --git a/tests/unit/test_cuda_backward.py b/tests/unit/test_cuda_backward.py
index 2c7e07aa8b31..e05cb1190dde 100755
--- a/tests/unit/test_cuda_backward.py
+++ b/tests/unit/test_cuda_backward.py
@@ -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):
@@ -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),
@@ -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',
diff --git a/tests/unit/test_cuda_forward.py b/tests/unit/test_cuda_forward.py
index 5add5e152a91..73e847aa3ac4 100755
--- a/tests/unit/test_cuda_forward.py
+++ b/tests/unit/test_cuda_forward.py
@@ -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),
@@ -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,
@@ -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,