[API] use softmax with length, and interleaved matmul for BERT

[eric-haibin-lin]

Description

Checklist

Essentials

• PR's title starts with a category (e.g. [BUGFIX], [MODEL], [TUTORIAL], [FEATURE], [DOC], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

Changes

• Feature1, tests, (and when applicable, API doc)
• Feature2, tests, (and when applicable, API doc)

Comments

• If this change is a backward incompatible change, why must this change be made.
• Interesting edge cases to note here

[codecov]

Codecov Report

Merging #1136 into master will decrease coverage by 0.41%.
The diff coverage is 98.8%.

@@            Coverage Diff             @@
##           master    #1136      +/-   ##
==========================================
- Coverage   87.76%   87.34%   -0.42%     
==========================================
  Files          67       67              
  Lines        6310     6386      +76     
==========================================
+ Hits         5538     5578      +40     
- Misses        772      808      +36

Impacted Files	Coverage Δ
src/gluonnlp/model/bert.py	94.28% <98.8%> (+1.63%)	⬆️
src/gluonnlp/data/batchify/embedding.py	45.16% <0%> (-52.42%)	⬇️
src/gluonnlp/utils/files.py	42.62% <0%> (-3.28%)	⬇️
src/gluonnlp/vocab/subwords.py	85.1% <0%> (-2.13%)	⬇️
src/gluonnlp/data/question_answering.py	100% <0%> (ø)	⬆️
src/gluonnlp/model/attention_cell.py	91.06% <0%> (+0.55%)	⬆️
src/gluonnlp/model/transformer.py	91.66% <0%> (+4.48%)	⬆️
src/gluonnlp/model/utils.py	77.69% <0%> (+6.92%)	⬆️
src/gluonnlp/model/seq2seq_encoder_decoder.py	80% <0%> (+30%)	⬆️

[fhieber]

Hi @eric-haibin-lin,
could you help me understand the difference between this PR and the previously reverted commit? That is, what has changed to address #1127 ?
We are about to switch to the fast multi-head attention ops in Sockeye as well. I also observed perplexity differences there, so I'd like to understand the root cause of the reverted commit for gluon-nlp. Thanks!

[mli]

Job PR-1136/1 is complete.
Docs are uploaded to http://gluon-nlp-staging.s3-accelerate.dualstack.amazonaws.com/PR-1136/1/index.html

[eric-haibin-lin]

Hi @fhieber the main difference is in this code block:

        query_weight = query_weight.reshape(shape=(self._num_heads, -1, 0), reverse=True)
        key_weight = key_weight.reshape(shape=(self._num_heads, -1, 0), reverse=True)
        value_weight = value_weight.reshape(shape=(self._num_heads, -1, 0), reverse=True)
        in_weight = F.concat(query_weight, key_weight, value_weight, dim=-2)
        in_weight = in_weight.reshape(shape=(-1, 0), reverse=True)

I think the new ops assume the projection is done with interleaving weights for k/q/v. The concatenated weight should have shape (num_heads, C_out/num_heads * 3, C_in).

[mli]

Job PR-1136/2 is complete.
Docs are uploaded to http://gluon-nlp-staging.s3-accelerate.dualstack.amazonaws.com/PR-1136/2/index.html

[mli]

Job PR-1136/3 is complete.
Docs are uploaded to http://gluon-nlp-staging.s3-accelerate.dualstack.amazonaws.com/PR-1136/3/index.html

[mli]

Job PR-1136/4 is complete.
Docs are uploaded to http://gluon-nlp-staging.s3-accelerate.dualstack.amazonaws.com/PR-1136/4/index.html

[eric-haibin-lin]

replaced #1091

[TaoLv]

Is it possible to avoid concat for every iteration? Or at least for inference, we only need concat once, right?

[eric-haibin-lin]

For inference, yes that's true. It's similar to RNN. If we figure out a way to avoid the weight concat in RNN, we can apply that here, too. @TaoLv do you have any idea/suggestion?