[Large Tensor] Fixed Embedding op

[connorgoggins]

Description

The Embedding op was previously breaking on large tensor (dimension >= 2^32) data. With the following input:

nd.Embedding(data=nd.random_normal(shape=(2**32,1)), weight=nd.random_normal(shape=(2**32,1)), input_dim=2**32, output_dim=1)

the following error was thrown:

mxnet.base.MXNetError: MXNetError: Invalid Parameter format for input_dim expect int but value='4294967296', in operator Embedding(name="", output_dim="1", input_dim="4294967296")

To fix this issue, I modified indexing_op.h to switch from storing input_dim as an int to storing it as an index_t. After implementing my fix and rebuilding, the previous input command displayed the correct output:

[[[-0.5190417]]

 [[-1.4388928]]

 [[ 1.1367434]]

 ...

 [[-1.4388928]]

 [[-1.4388928]]

 [[-0.5190417]]]
<NDArray 4294967296x1x1 @cpu(0)>

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented
• To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• M src/operator/tensor/indexing_op.h

Comments

Tested on r5dn.24xl-ubuntu 16.04 and p2.16xl-ubuntu 16.04 with

1. Individual op run
2. Full OpPerf run

Results

The key difference between CPU and GPU tests was the instance type (r5dn.24xl for CPU, p2.16xl for GPU). All relevant build flags remain the same, and both were tested using CPU context.

Single operator test - Embedding op (GPU)
Single operator test - Embedding op (CPU)

Full OpPerf test (GPU)
Full OpPerf test (CPU)

@apeforest @access2rohit @ChaiBapchya

[connorgoggins]

@mxnet-label-bot add [pr-awaiting-review]

[access2rohit]

@connorgoggins can you paste output of a case where sparseEmbedding params are used ? Are both used in the same example use case presented here ?

[connorgoggins]

@access2rohit great question. Here is an example of when sparseEmbedding params are used without my fix, and the resulting error (same error as with Embedding):

>>> mx.contrib.nd.SparseEmbedding(data=mx.nd.random_normal(shape=(2**32,1)), weight=mx.nd.random_normal(shape=(2**32, 1)), input_dim=2**32, output_dim=1)

mxnet.base.MXNetError: MXNetError: Invalid Parameter format for input_dim expect int but value='4294967296', in operator _contrib_SparseEmbedding(name="", output_dim="1", input_dim="4294967296")

With my fix, the call above passes without any issues - output below:

[[[-0.5190417]]

 [[-1.4388928]]

 [[ 1.1367434]]

 ...

 [[-1.4388928]]

 [[-1.4388928]]

 [[-0.5190417]]]
<NDArray 4294967296x1x1 @cpu(0)>

[apeforest]

Would there be a case where output_dim also exceeds 2^32?

[connorgoggins]

@apeforest excellent point, there would be. For example, the following call throws an error:

>>> mx.nd.Embedding(data=mx.nd.random_normal(shape=(1,)), weight=mx.nd.random_normal(shape=(1,2**32)), input_dim=1, output_dim=2**32)

mxnet.base.MXNetError: MXNetError: Invalid Parameter format for output_dim expect int but value='4294967296', in operator Embedding(name="", output_dim="4294967296", input_dim="1")

With my latest update (changing the dtype of output_dim to index_t), here is the result:

>>> mx.nd.Embedding(data=mx.nd.random_normal(shape=(1,)), weight=mx.nd.random_normal(shape=(1,2**32)), input_dim=1, output_dim=2**32)

[[ 1.6323917  -0.33354783 -1.7378405  ... -0.36648417  0.6363522
   2.367109  ]]
<NDArray 1x4294967296 @cpu(0)>

[apeforest]

LGTM. Have you verified that your nightly test can run successfully in your local machine?

[connorgoggins]

@apeforest thanks! The individual nightly tests I wrote for each op run successfully on my local machine. Running the full nightly test suite for each op now.

[connorgoggins]

@apeforest update: full suite of nightly tests passed on r5dn.24xl instances for every one of my PRs with additional tests added.

[apeforest]

LGTM. Thanks for the good work

[apeforest]

@connorgoggins I retriggered CI tests