Exception: The provided qkv memory layout is not supported! when using RoPE

[halilakin]

Both MultiheadAttention and TransformerLayer work fine without the rotary_pos_emb in the forward pass. However, when passing rotary_pos_emb, I get the error above.

It seems using RoPE (i.e., passing rotary embeddings in the forward of MHA) is causing query_layer and key_layer's memory layouts to change. Is this expected and is there a workaround?

[ptrendx]

@cyanguwa Could you investigate this? It sounds like a bug.

[yongyanrao]

Noticed the same issue. It seems due to rotary is not an inplace operation.

[cyanguwa]

@halilakin @yongyanrao , please give PR 453 a try and let me know if there's still a problem. Thanks!

[yongyanrao]

@cyanguwa Thank you for the quick response. We tested the changes, along with the newly implemented RotaryPositionEmbedding. They worked perfectly from our side so far. Also doing some further tests, will keep you posted on the result. Thanks again.

[ptrendx]

Seems that the issue has been resolved - closing. Please reopen if your tests show that there are still some problems here.