BERT pretraining

[ZiyueHuang]

Description

@sxjscience

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

cc @dmlc/gluon-nlp-team

[github-actions]

The documentation website for preview: http://gluon-nlp-dev.s3-accelerate.amazonaws.com/PR1376/bert/index.html

[sxjscience]

Have you met some bugs of argmax?

[ZiyueHuang]

mx.npx.topk is faster than mx.np.argmax about 50 times. The issue (apache/mxnet#11061) still exists in MXNet 2.0.

import time
import mxnet as mx
import numpy as np
import os
tmp = mx.np.random.normal(-1, 1, (64,300000), ctx=mx.gpu())
tic = time.time()
for i in range(20):
   if (i == 5):
        begin = time.time();
   elif (i == 15):
        end = time.time();
   tic = time.time()
   #out = mx.np.argmax(tmp, axis=1)
   out = mx.npx.topk(tmp, k=1, ret_typ='indices', axis=1, dtype=np.int32)
   out.wait_to_read()
   toc = time.time() - tic
   print ("used time %f:"%toc)
avg = (end - begin) / 10
print ("avg time %f"%avg)

[sxjscience]

Thanks Ziyue, I can reproduce....

@szha Our argmax is too slow...

[sxjscience]

I'm actually not aware of this. Argmax is quite easy to accelerate and we can even call packages like cub.

[codecov]

Codecov Report

Merging #1376 into master will decrease coverage by 0.04%.
The diff coverage is n/a.

@@            Coverage Diff             @@
##           master    #1376      +/-   ##
==========================================
- Coverage   71.09%   71.04%   -0.05%     
==========================================
  Files         107      107              
  Lines       12607    12607              
==========================================
- Hits         8963     8957       -6     
- Misses       3644     3650       +6     

Impacted Files	Coverage Δ
src/gluonnlp/data/filtering.py	78.26% <0.00%> (-4.35%)	⬇️
src/gluonnlp/data/tokenizers/yttm.py	81.89% <0.00%> (-0.87%)	⬇️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 4fb41d7...9f83ab7. Read the comment docs.

[github-actions]

The documentation website for preview: http://gluon-nlp-dev.s3-accelerate.amazonaws.com/PR1376/bert/index.html

[github-actions]

The documentation website for preview: http://gluon-nlp-dev.s3-accelerate.amazonaws.com/PR1376/bert/index.html

[sxjscience]

LGTM overall. But we may later need to consider to merge the Electra + Bert (nsp + mlm) + Albert (sop + mlm) implementations.

[sxjscience]

Would you try to merge the upstream/master? It seems that there's something wrong in the GPU test.

[ZiyueHuang]

Would you try to merge the upstream/master? It seems that there's something wrong in the GPU test.

done

[github-actions]

The documentation website for preview: http://gluon-nlp-dev.s3-accelerate.amazonaws.com/PR1376/bert/index.html