Add layout + compute_layout support: TransformerNMT, BERT, ALBERT, ELECTRA, MobileBERT, RoBERTA, XLMR

[sxjscience]

@MoisesHer @dmlc/gluon-nlp-team

This PR adds two additional flags to backbone models to enhance the computational speed and usability.

• layout
  The layout of the inputs + outputs, where
  • NT: (batch_size, sequence_length, ...)
  • TN: (sequence_length, batch_size, ...)
• compute_layout
  The layout of the inner computation. By default, we have the "auto" option, in which GluonNLP will determine the best layout based on heuristics.

The technical insights about why layouts may matter lies as follows (also documented here

gluon-nlp/src/gluonnlp/attention_cell.py

Lines 540 to 546 in cd48efd

	# 2. Calculate the attention weights
	# Score: (L_query, B, N, C_Q) X (L_mem, B, N, C_Q) --> (B, N, L_query, L_mem)
	# This layout structure can be implemented very efficiently because B, N are consecutive
	# to each other. To have a clear picture of what's happening, we may consider the
	# (i, j)th element of the output
	# out[i, j, :, :] = query[:, i, j, :] X key[:, i, j, :].T, which is just one GEMM call
	# We can thus implement the whole kernel via a single call of batched GEMM with stride.

):

When the layouts of memory and query are "TNC", they have the shape:

• query: (L_query, B, N, C_Q)
• memory: (L_mem, B, N, C_Q)

One step in the AttentionCell is to obtain the multi-head attention scores, which can be written as follows:

(L_query, B, N, C_Q) X (L_mem, B, N, C_Q) -> (B, N, L_query, L_mem)

This layout structure can be implemented very efficiently because B, N are consecutive to each other.

To have a clear picture of what's happening, we may consider the (i, j)th element of the output, i.e.,

out[i, j, :, :] = query[:, i, j, :] X key[:, i, j, :].T

This is just one GEMM call. We can thus implement the whole kernel via a single call of batched GEMM by correctly specifying the strides.

Also, in fairseq, the inner computation of the TransformerEncoder is using the TN layout:

https://github.com/pytorch/fairseq/blob/108bb2560b1ec01524ba723bc7c69186875afa0a/fairseq/models/transformer.py#L399-L407

After this PR, these models will have the layout flag:

• TransformerEncoder, TransformerDecoder, TransformerNMTModel
• BertModel
• AlbertModel
• ElectraModel
• MobileBertModel
• RobertaModel
• XLMRModel
• Test cases

[codecov]

Codecov Report

Merging #1258 into numpy will increase coverage by 1.35%.
The diff coverage is 89.17%.

@@            Coverage Diff             @@
##            numpy    #1258      +/-   ##
==========================================
+ Coverage   82.56%   83.92%   +1.35%     
==========================================
  Files          38       41       +3     
  Lines        5534     6157     +623     
==========================================
+ Hits         4569     5167     +598     
- Misses        965      990      +25     

Impacted Files	Coverage Δ
setup.py	0.00% <ø> (ø)
src/gluonnlp/models/transformer_xl.py	82.52% <66.66%> (-0.20%)	⬇️
src/gluonnlp/models/xlmr.py	88.23% <70.00%> (+1.35%)	⬆️
src/gluonnlp/models/mobilebert.py	87.72% <81.67%> (+6.37%)	⬆️
src/gluonnlp/models/electra.py	80.78% <84.68%> (+1.83%)	⬆️
src/gluonnlp/attention_cell.py	79.52% <86.66%> (-0.39%)	⬇️
src/gluonnlp/models/roberta.py	93.67% <87.50%> (+4.89%)	⬆️
src/gluonnlp/models/bert.py	94.86% <92.96%> (+9.93%)	⬆️
src/gluonnlp/utils/testing.py	94.11% <93.47%> (-3.03%)	⬇️
src/gluonnlp/models/albert.py	95.47% <93.80%> (-1.22%)	⬇️
... and 17 more

[sxjscience]

Now I'm waiting for #1261. We added two flags:

• layout
  It means the layout of the input + output
• compute_layout
  It means the layout of the inner computation. By default, we will always use compute_layout=auto, which will automatically determine the best compute layout based on our heuristic.

[sxjscience]

@zheyuye I'm moving the cfgs to the file to make it consistent to Transformer + RoBERTA.

[zheyuye]

ok, I am going to revise other models

[sxjscience]

@dmlc/gluon-nlp-committers @MoisesHer @zheyuye Should be ready for review.