Support table batched TT embedding lookup

facebookresearch / FBTT-Embedding

This is a Tensor Train based compression library to compress sparse embedding tables used in large-scale machine learning models such as recommendation and natural language processing. We showed this library can reduce the total model size by up to 100x in Facebook’s open sourced DLRM model while achieving same model quality. Our implementation is faster than the state-of-the-art implementations. Existing the state-of-the-art library also decompresses the whole embedding tables on the fly therefore they do not provide memory reduction during runtime of the training. Our library decompresses only the requested rows therefore can provide 10,000 times memory footprint reduction per embedding table. The library also includes a software cache to store a portion of the entries in the table in decompressed format for faster lookup and process.

MIT License

192 stars 27 forks source link

Support table batched TT embedding lookup #13

Closed bilgeacun closed 3 years ago

bilgeacun commented 3 years ago

Summary: Support table batched TT embedding lookup assuming the multiple tables have exactly the same hash size and embedding dim.

Make each tt_cores 3 dimensional: [num_tables, tt_p_shape[i], ...]
Handle the additional num_tables dimension correctly in the computation
Have a TableBatchedTTEmbeddingBag that supports passing num_tables
Make TTEmbeddingBag as special case of TableBatchedTTEmbeddingBag when num_tables = 1
Add unit test to compare TableBatchedTTEmbeddingBag with running multiple TTEmbeddingBags

Differential Revision: D27091933

facebook-github-bot commented 3 years ago

This pull request was exported from Phabricator. Differential Revision: D27091933

facebook-github-bot commented 3 years ago

This pull request was exported from Phabricator. Differential Revision: D27091933

facebook-github-bot commented 3 years ago

This pull request has been merged in facebookresearch/FBTT-Embedding@1662e0e677b95ab20d3b29dc905f99fae959ce75.