Introduce distributed embeddings

[edknv]

Part of NVIDIA-Merlin/Merlin#733.

Goals ⚽

There is a package called distributed-embeddings, a library for building large embedding based (e.g. recommender) models in Tensorflow. It's an alternative approach to SOK.

This PR introduces DistributedEmbedding for multi-GPU embedding table support.

Implementation Details 🚧

• distributed-embeddings by default will round-robin the entire embedding tables across the GPUs, e.g., the first embedding table on GPU 1, the second one on GPU 2, etc.
• In theory the tables can be sharded by using column_slice but this has not been tested thoroughly from Models side.
• Most of the logic is for inferring the embedding size from the schema using the cardinality in int_domain (similarly to the existing EmbeddingTable), determining shapes, and translating a dictionary input into an ordered list input (because distributed-embeddings doesn't support dictionaries yet).
• From the user perspective, they can replace mm.Embeddings with mm.DistributedEmbeddings in their models when they wish to use multi-GPU embedding tables. (See the unit test for DLRM.)
• Depends on upstream fix: Use tf.shape for graph mode support distributed-embeddings#6
• Added a Github actions for running unit tests that depend on horovod. distributed-embeddings is for now installed via a script that clones the github repo and installs from source, because there is no pypi package.

Testing Details 🔍

Unit tests: tests/unit/tf/horovod/test_embedding.py
Performance tests: TBD

[github-actions]

Documentation preview

https://nvidia-merlin.github.io/models/review/pr-974

[bschifferer]

The distributed embedding examples uses a custom train step functions:
https://github.com/NVIDIA-Merlin/distributed-embeddings/blob/main/examples/dlrm/main.py#L201-L215

In my understanding, distributed embedding does NOT work with keras model.fit function:
https://github.com/NVIDIA-Merlin/models/pull/974/files#diff-1e42e5c4771f01c26b3c78c545eb341590a4406b2c5af8da0491ab4b7ea51464R80

I think we need the distributed embedding team to review the PR

[FDecaYed]

Looks good to me.
On model fit support:
Current code using model.fit should work since nothing in distributed-embedding conflicts with keras model fit/compile.

The reason we have custom train_step() example is for when user wants hybrid data/model parallel. The way horovod data parallel support model.fit() api is wrapping optimizer, which will break now if distributed model parallel is also used. By my understanding, merlin-model also have integration with horovod data parallel through distributedoptimizer? if so user could run into problem when they use both integration

We will support model.fit+hvd.distributedoptimizer in next release, so code here should just work. One caveat is the fix(and later version of DE) will be depend on a later horovod version.

Alternatively, I think merlin-model would need to implement custom train_step in block/model using DE? That'll be a much bigger change though.

[rnyak]

The PR needs to be update based on the dataloader changes. There is a new version of DE. we need to add an integration test as well to be sure that the functionality is working.

[rnyak]

@FDecaYed hello. do you have any updates for this PR? thanks.

[FDecaYed]

@rnyak Sorry, this fell off my list.
As of now, DE already added support for model fit. So some of the problems should be gone. I'm willing to jump in and help if needed.

On the other hand, I'm not familiar with merlin models and the dataloader change you mentioned. @edknv do you know what it is and could you help bring the code up to date?