Use scaled_dot_product_attention in WavLM attention

[nateanl]

Fix #3219.

torch.nn.MultiheadAttention will throw an error if torch.no_grad() and mask are both given. The pull request fixes it by replacing the forward method with torch.nn.functional.scaled_dot_product_attention.

[nateanl]

Here are the benchmark results with new changes. The benchmark script can be found in https://gist.github.com/nateanl/97b2f9adb39c05a4e854fbd924de01f6.

MultiheadAttention
[--------------- WavLM benchmark ---------------]
                              |    CPU    |  CUDA
1 threads: --------------------------------------
      wavlm_base 5 seconds    |   1561.1  |  14.2
      wavlm_base 10 seconds   |   3325.5  |  22.1
      wavlm_base 15 seconds   |   5423.0  |  34.4
      wavlm_base 20 seconds   |   7722.4  |  50.7
      wavlm_large 5 seconds   |   4451.5  |  25.2
      wavlm_large 10 seconds  |   9465.5  |  48.3
      wavlm_large 15 seconds  |  15140.8  |  72.6
      wavlm_large 20 seconds  |  22030.9  |  97.9
Times are in milliseconds (ms).


scaled_dot_product_attention
[--------------- WavLM benchmark ---------------]
                              |    CPU    |  CUDA
1 threads: --------------------------------------
      wavlm_base 5 seconds    |   1562.8  |  13.2
      wavlm_base 10 seconds   |   3164.2  |  21.1
      wavlm_base 15 seconds   |   5175.0  |  33.1
      wavlm_base 20 seconds   |   7357.2  |  48.9
      wavlm_large 5 seconds   |   4415.9  |  23.3
      wavlm_large 10 seconds  |   9279.8  |  47.8
      wavlm_large 15 seconds  |  14930.8  |  71.1
      wavlm_large 20 seconds  |  21524.6  |  97.3
Times are in milliseconds (ms).

[facebook-github-bot]

@nateanl has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[nateanl]

Here is the script which shows the new scaled_dot_product_attention achieves identical outputs as original implementation.

[facebook-github-bot]

@nateanl merged this pull request in adb0338.

[github-actions]

Hey @nateanl.
You merged this PR, but labels were not properly added. Please add a primary and secondary label (See https://github.com/pytorch/audio/blob/main/.github/process_commit.py)