Merge pull request #182 from NVIDIA/gptj_fix

fix: fix bug of GPT-J due to the setting of rotary_embedding_dim
NVIDIA · Jan 12, 2022 · f8f0ee8 · f8f0ee8
2 parents 8a1c109 + ea7b880
commit f8f0ee8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/fastertransformer/layers/attention_layers/DecoderSelfAttentionLayer.cc b/src/fastertransformer/layers/attention_layers/DecoderSelfAttentionLayer.cc
@@ -76,6 +76,7 @@ void fusedQKV_masked_attention_dispatch(const T* qkv_buf,
     params.length_per_sample = sequence_lengths;
     params.timestep = step - 1;
     params.num_heads = head_num;
+    params.rotary_embedding_dim = rotary_embedding_dim;
     params.hidden_size_per_head = size_per_head;
     params.inv_sqrt_dh = 1.F / (sqrtf((float)params.hidden_size_per_head) * q_scaling);