Fixing several bugs in the inference-api and the kernels (#1951)

Co-authored-by: Jeff Rasley <[email protected]>

.github/workflows/amd.yml

@@ -37,12 +37,23 @@ jobs:
           python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
           sudo apt-get update
           sudo apt-get install -y libaio-dev
+
+      - name: Install transformers
+        run: |
+          git clone https://github.com/huggingface/transformers
+          cd transformers
+          # if needed switch to the last known good SHA until transformers@master is fixed
+          # git checkout 1cc453d33
+          git rev-parse --short HEAD
+          pip install .
+
       # Runs a set of commands using the runners shell
       - name: Install deepspeed
         run: |
           sudo /opt/conda/bin/pip install .[dev,1bit,autotuning]
           #python -c "from deepspeed.env_report import cli_main; cli_main()"
           ds_report
+
       # Runs a set of commands using the runners shell
       - name: Unit tests
         run: |

.github/workflows/nv-torch12-p40.yml

@@ -32,6 +32,15 @@ jobs:
           python -c "import torch; print('torch:', torch.__version__, torch)"
           python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
 
+      - name: Install transformers
+        run: |
+          git clone https://github.com/huggingface/transformers
+          cd transformers
+          # if needed switch to the last known good SHA until transformers@master is fixed
+          # git checkout 1cc453d33
+          git rev-parse --short HEAD
+          pip install .
+
       - name: Install deepspeed
         run: |
           pip install .[dev,autotuning]

.github/workflows/nv-torch18-v100.yml

@@ -32,10 +32,21 @@ jobs:
           pip install torch==1.8.2+cu111 torchvision==0.9.2+cu111 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html
           python -c "import torch; print('torch:', torch.__version__, torch)"
           python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
+
+      - name: Install transformers
+        run: |
+          git clone https://github.com/huggingface/transformers
+          cd transformers
+          # if needed switch to the last known good SHA until transformers@master is fixed
+          # git checkout 1cc453d33
+          git rev-parse --short HEAD
+          pip install .
+
       - name: Install deepspeed
         run: |
           pip install .[dev,1bit,autotuning,sparse_attn]
           ds_report
+
       - name: Unit tests
         run: |
           unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch

csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu

@@ -4,6 +4,7 @@
 #include <cuda_profiler_api.h>
 #endif
 
+namespace cg = cooperative_groups;
 namespace cg = cooperative_groups;
 
 __global__ void apply_rotary_pos_emb(float* mixed_query,
@@ -153,7 +154,9 @@ __global__ void apply_rotary_pos_emb1(__half* mixed_query,
     int lane = id & 0x1f;
 
     unsigned head_id = blockIdx.x * MAX_WARP_NUM + gid;
+    unsigned seq_index = head_id % seq_len;
     unsigned offset = head_id * head_size;
+    unsigned k_offset = (seq_index + (head_id / seq_len) * MAX_OUT_TOKES) * head_size;
 
     constexpr unsigned mask[32] = {
         0x1 | 0x1000,     0x2 | 0x2000,     0x4 | 0x4000,     0x8 | 0x8000,     0x10 | 0x10000,
@@ -171,7 +174,7 @@ __global__ void apply_rotary_pos_emb1(__half* mixed_query,
             float inv_freq = (float)((lane % half_dim) * 2) / (float)rotary_dim;
             inv_freq = 1.0 / powf(10000.0, inv_freq) * (float)seq_id;
             float q = (float)mixed_query[offset + lane];
-            float k = (float)key_layer[offset + lane];
+            float k = (float)key_layer[k_offset + lane];
             float rotary_sign = (lane > (half_dim - 1) ? -1.0 : 1.0);
             float q_rot = (q * rotary_sign);
             float k_rot = (k * rotary_sign);
@@ -183,7 +186,7 @@ __global__ void apply_rotary_pos_emb1(__half* mixed_query,
             k = k * cosf(inv_freq) + k_rot_tmp * sinf(inv_freq);
 
             mixed_query[offset + lane] = (__half)q;
-            key_layer[offset + lane] = (__half)k;
+            key_layer[k_offset + lane] = (__half)k;
 
             lane += WARP_SIZE;
         }
@@ -237,6 +240,7 @@ template void launch_apply_rotary_pos_emb<__half>(__half*,
                                                   bool,
                                                   bool,
                                                   cudaStream_t);
+
 /*
 __global__ void apply_rotary_pos_emb(float* mixed_query,
 float* key_layer,

csrc/transformer/inference/csrc/gelu.cu

@@ -317,12 +317,18 @@ __global__ void gptj_residual_add(float* input,
         float4 out = output_cast[offset];
         float4 res_vec = attn_cast[offset];
         float4 bias_data = bias_cast[offset % intermediate_size];
-        float4 attn_bias = attnbias_cast[offset % intermediate_size];
 
-        data.x = data.x * mp_size + (out.x + res_vec.x + bias_data.x + attn_bias.x);
-        data.y = data.y * mp_size + (out.y + res_vec.y + bias_data.y + attn_bias.y);
-        data.z = data.z * mp_size + (out.z + res_vec.z + bias_data.z + attn_bias.z);
-        data.w = data.w * mp_size + (out.w + res_vec.w + bias_data.w + attn_bias.w);
+        if (attnbias) {
+            float4 attn_bias = attnbias_cast[offset % intermediate_size];
+            data.x += attn_bias.x;
+            data.y += attn_bias.y;
+            data.z += attn_bias.z;
+            data.w += attn_bias.w;
+        }
+        data.x = data.x * mp_size + (out.x + res_vec.x + bias_data.x);
+        data.y = data.y * mp_size + (out.y + res_vec.y + bias_data.y);
+        data.z = data.z * mp_size + (out.z + res_vec.z + bias_data.z);
+        data.w = data.w * mp_size + (out.w + res_vec.w + bias_data.w);
 
         output_cast[offset] = data;
     }
@@ -354,13 +360,11 @@ __global__ void gptj_residual_add(__half* input,
         float2 res_vec = attn_cast[offset];
 
         float2 bias_vec = bias_cast[offset % intermediate_size];
-        float2 attn_bias_vec = attnbias_cast[offset % intermediate_size];
 
         __half2* vals_half = reinterpret_cast<__half2*>(&vals_vec);
         __half2* out_half = reinterpret_cast<__half2*>(&out_vec);
         __half2* res_half = reinterpret_cast<__half2*>(&res_vec);
         __half2* bias_half = reinterpret_cast<__half2*>(&bias_vec);
-        __half2* attnbias_half = reinterpret_cast<__half2*>(&attn_bias_vec);
 
         float2 low_data = __half22float2(vals_half[0]);
         float2 high_data = __half22float2(vals_half[1]);
@@ -373,18 +377,21 @@ __global__ void gptj_residual_add(__half* input,
 
         float2 low_bias = __half22float2(bias_half[0]);
         float2 high_bias = __half22float2(bias_half[1]);
-
-        float2 attn_low_bias = __half22float2(attnbias_half[0]);
-        float2 attn_high_bias = __half22float2(attnbias_half[1]);
-
-        low_data.x =
-            low_data.x * mp_size + (low_out.x + low_res.x + (low_bias.x + attn_low_bias.x));
-        low_data.y =
-            low_data.y * mp_size + (low_out.y + low_res.y + (low_bias.y + attn_low_bias.y));
-        high_data.x =
-            high_data.x * mp_size + (high_out.x + high_res.x + (high_bias.x + attn_high_bias.x));
-        high_data.y =
-            high_data.y * mp_size + (high_out.y + high_res.y + (high_bias.y + attn_high_bias.y));
+        if (attn_bias) {
+            float2 attn_bias_vec = attnbias_cast[offset % intermediate_size];
+            __half2* attnbias_half = reinterpret_cast<__half2*>(&attn_bias_vec);
+            float2 attn_low_bias = __half22float2(attnbias_half[0]);
+            float2 attn_high_bias = __half22float2(attnbias_half[1]);
+            low_data.x += attn_low_bias.x;
+            low_data.y += attn_low_bias.y;
+            high_data.x += attn_high_bias.x;
+            high_data.y += attn_high_bias.y;
+        }
+
+        low_data.x = low_data.x * mp_size + (low_out.x + low_res.x + (low_bias.x));
+        low_data.y = low_data.y * mp_size + (low_out.y + low_res.y + (low_bias.y));
+        high_data.x = high_data.x * mp_size + (high_out.x + high_res.x + (high_bias.x));
+        high_data.y = high_data.y * mp_size + (high_out.y + high_res.y + (high_bias.y));
 
         vals_half[0] = __float22half2_rn(low_data);
         vals_half[1] = __float22half2_rn(high_data);