Speedup batch matmul in pytorch
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
Hi there,

I'm trying to combine Sparse Transformer into Vision Transformer to speed up the time-consuming training also inference.
However, the speed is really slow with my code which is shown below. Can someone help me point out where I was wrong?

Thanks.
`import torch
import math

query_layer = torch.randn(32, 12, 197, 64)
key_layer = torch.randn(32, 12, 197, 64)
key_layer_transpose = key_layer.transpose(-1, -2)

dim0 = query_layer.shape[0]
dim1 = query_layer.shape[1]
dim2 = query_layer.shape[2]
dim3 = query_layer.shape[3]

print(dim0, dim1, dim2, dim3)

#orignal transformer attention score calculation
attention_scores = torch.matmul(query_layer, key_layer_transpose)
print(attention_scores)

#my modified based on Sparse Transformer for speeding up the training (but actually lower than at least 20 times in comparison with the original transformer
N = math.sqrt(dim3)
attention_scores = torch.zeros(dim0, dim1, dim2, dim2, device='cuda:0')
for i_dim0 in range(dim0):
  for i_dim1 in range(dim1):
    for i in range(dim2):
      for j in range(dim2):
        if (i == j) or ((i - j) % N == 0 and i - j > 0):
          attention_scores[i_dim0, i_dim1, i, j] = torch.matmul(query_layer[i_dim0, i_dim1, i, :], key_layer_transpose[i_dim0, i_dim1, :, j])
attention_scores.shape
print(attention_scores)
`


Hello, thanks for opening an issue! We try to keep the github issues for bugs/feature requests.
Could you ask your question on the [forum](https://discuss.huggingface.co) instead?

Thanks!<|||||>I've found the solution by using deepspeed.

However, I encountered the error and I opened an issue here. Please help me
with this issue. Thanks.

https://github.com/microsoft/DeepSpeed/issues/1153

On Thu, Jun 10, 2021 at 2:04 PM Lysandre Debut ***@***.***>
wrote:

> Hello, thanks for opening an issue! We try to keep the github issues for
> bugs/feature requests.
> Could you ask your question on the forum <https://discuss.huggingface.co>
> instead?
>
> Thanks!
>
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/issues/12093#issuecomment-858369696>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AJY7KIWLCP5AOY7CA3JI6ZDTSBPY3ANCNFSM46ND2SZQ>
> .
>
