Gluon-nlp: Different tensor contexts in nlp.optimizer.lamb

[zburning]

Description

In the update() function in lamb.py:

# preprocess grad
   grad *= self.rescale_grad
   if self.clip_gradient is not None:
       grad = clip(grad, -self.clip_gradient, self.clip_gradient)

The rescale_grad should be as_in_context(grad.context)? Otherwise it will raise illegal memory access in multi-gpu training

Error Message

(Paste the complete error message, including stack trace.)

To Reproduce

(If you developed your own code, please provide a short script that reproduces the error. For existing examples, please provide link.)

Steps to reproduce

(Paste the commands you ran that produced the error.)

1.
2.

What have you tried to solve it?

1.
2.

Environment

We recommend using our script for collecting the diagnositc information. Run the following command and paste the outputs below:

curl --retry 10 -s https://raw.githubusercontent.com/dmlc/gluon-nlp/master/tools/diagnose.py | python

# paste outputs here

[eric-haibin-lin]

Since mxnet has lamb optimizer available in 1.6, we shall just reuse that one and remove the one in gluonnlp.

[leezu]

Related: https://github.com/dmlc/gluon-nlp/issues/1019

[eric-haibin-lin]

[API] Drop LAMB optimizer from GluonNLP in favor of MXNet version #1116