Add RETRO indexed dataset and `set_inference_key_value_memory` inference #4220

yidong72 · 2022-05-20T23:07:54Z

What does this PR do ?

Added the RETRO index dataset for both data and retrieval data
Added the KNN Map index dataset
Added unit tests for all the index dataset
Added RETRO dataset that uses index dataset
Added the scripts to preprocess, build Faiss index and generate KNN map index for training
Added the set_inference_key_value_memory capability to RETRO modules so it can run inference efficiently
Added unit tests for efficient inference of all RETRO modules.

Signed-off-by: Yi Dong <yidong@nvidia.com>

lgtm-com · 2022-05-30T16:58:28Z

This pull request introduces 1 alert when merging cd0a84a into 7a9a8f0 - view on LGTM.com

new alerts:

1 for Unused import

Signed-off-by: Yi Dong <yidong@nvidia.com>

lgtm-com · 2022-05-30T17:29:11Z

This pull request introduces 1 alert when merging e4809f7 into 7a9a8f0 - view on LGTM.com

new alerts:

1 for Unused import

Signed-off-by: Yi Dong <yidong@nvidia.com>

lgtm-com · 2022-05-31T19:47:00Z

This pull request introduces 1 alert when merging ab14327 into e838862 - view on LGTM.com

new alerts:

1 for Unused import

Signed-off-by: Yi Dong <yidong@nvidia.com>

okuchaiev · 2022-06-01T17:36:36Z

/blossom-ci

michalivne

Very few minor comments. LGTM.

nemo/collections/nlp/models/language_modeling/megatron_retrieval_model.py

examples/nlp/language_modeling/conf/megatron_retro_config.yaml

michalivne · 2022-06-01T21:44:10Z

nemo/collections/nlp/data/language_modeling/megatron/retro_dataset.py

+        retrieval_index: MMapRetrievalIndexedDataset,
+    ):
+        if not HAVE_APEX:
+            raise ImportError(


@yidong72 can you add a description of the arguments? (e.g., what is documents?)

added comment

michalivne · 2022-06-01T21:57:28Z

scripts/nlp_language_modeling/example_knn_map_quality.py

+            assert half == data_ds._index.chunk_size
+            neighbor_match = tokenizer.ids_to_text(token_ids[:half])
+            neighbor_extend = tokenizer.ids_to_text(token_ids[half:])
+            print(f' ->K{i}: {neighbor_match} --- {neighbor_extend}')


Why break retrieval into 2? Only half of each chunk is used for embedding?

because the retrieved data has two parts. The first part is used to match the query chunk, the second half is the continuation chunk.

Signed-off-by: Yi Dong <yidong@nvidia.com>

MaximumEntropy

Thanks for making the changes! transformer.py is getting really big :) I think we should try and minimize modifications to it in the future.

yidong72 · 2022-06-02T13:04:12Z

Thanks for making the changes! transformer.py is getting really big :) I think we should try and minimize modifications to it in the future.

Agree. Also, we should refactor the TransformerLayer class to make it abstract. For different models, we can add different implementations.

…nce (NVIDIA#4220) * added retrieval index dataset Signed-off-by: Yi Dong <yidong@nvidia.com> * added retrieval db handling Signed-off-by: Yi Dong <yidong@nvidia.com> * create and load data Signed-off-by: Yi Dong <yidong@nvidia.com> * working chunk retrieval Signed-off-by: Yi Dong <yidong@nvidia.com> * retrieval fetch works Signed-off-by: Yi Dong <yidong@nvidia.com> * unit test passes Signed-off-by: Yi Dong <yidong@nvidia.com> * add option to run retrieval preprocess Signed-off-by: Yi Dong <yidong@nvidia.com> * slice into chunks Signed-off-by: Yi Dong <yidong@nvidia.com> * add script to build index Signed-off-by: Yi Dong <yidong@nvidia.com> * building faiss index works Signed-off-by: Yi Dong <yidong@nvidia.com> * speed up the index building Signed-off-by: Yi Dong <yidong@nvidia.com> * added knn map index file Signed-off-by: Yi Dong <yidong@nvidia.com> * workign build knn map Signed-off-by: Yi Dong <yidong@nvidia.com> * added docstring Signed-off-by: Yi Dong <yidong@nvidia.com> * add retro dataset Signed-off-by: Yi Dong <yidong@nvidia.com> * dataset test passes Signed-off-by: Yi Dong <yidong@nvidia.com> * added unittest asserts Signed-off-by: Yi Dong <yidong@nvidia.com> * added dataset Signed-off-by: Yi Dong <yidong@nvidia.com> * training of retro is working Signed-off-by: Yi Dong <yidong@nvidia.com> * remove unused imports Signed-off-by: Yi Dong <yidong@nvidia.com> * fix jenkins Signed-off-by: Yi Dong <yidong@nvidia.com> * added knn example data Signed-off-by: Yi Dong <yidong@nvidia.com> * better print format Signed-off-by: Yi Dong <yidong@nvidia.com> * configure the number of neighbors Signed-off-by: Yi Dong <yidong@nvidia.com> * removed non-used cfg Signed-off-by: Yi Dong <yidong@nvidia.com> * turn on normliazaiton option Signed-off-by: Yi Dong <yidong@nvidia.com> * add layer number offset Signed-off-by: Yi Dong <yidong@nvidia.com> * need to add one Signed-off-by: Yi Dong <yidong@nvidia.com> * use at leaset one layer Signed-off-by: Yi Dong <yidong@nvidia.com> * added inference unit test Signed-off-by: Yi Dong <yidong@nvidia.com> * encoder inference test pass Signed-off-by: Yi Dong <yidong@nvidia.com> * encoder inference is confirmed to work Signed-off-by: Yi Dong <yidong@nvidia.com> * handles another edge case Signed-off-by: Yi Dong <yidong@nvidia.com> * alige relative position to the context Signed-off-by: Yi Dong <yidong@nvidia.com> * chunked cross attention passes test Signed-off-by: Yi Dong <yidong@nvidia.com> * fixed chunk cross attention masked attention Signed-off-by: Yi Dong <yidong@nvidia.com> * fix the padding Signed-off-by: Yi Dong <yidong@nvidia.com> * chunked cross attention layer inference passes unit test Signed-off-by: Yi Dong <yidong@nvidia.com> * working on the decoder Signed-off-by: Yi Dong <yidong@nvidia.com> * one mile stone for decoder Signed-off-by: Yi Dong <yidong@nvidia.com> * decoder is working Signed-off-by: Yi Dong <yidong@nvidia.com> * make encoder infer behave nicely Signed-off-by: Yi Dong <yidong@nvidia.com> * added encoder decoder inference unittest Signed-off-by: Yi Dong <yidong@nvidia.com> * remove bad imports Signed-off-by: Yi Dong <yidong@nvidia.com> * make training work Signed-off-by: Yi Dong <yidong@nvidia.com> * fix failed unit test Signed-off-by: Yi Dong <yidong@nvidia.com> * remove unused variables Signed-off-by: Yi Dong <yidong@nvidia.com> * added run on GPU for unittest Signed-off-by: Yi Dong <yidong@nvidia.com> * add pad id to the preprocessing script Signed-off-by: Yi Dong <yidong@nvidia.com> * added doc string for indexed dataset Signed-off-by: Yi Dong <yidong@nvidia.com> * efficient deduplicate doc Signed-off-by: Yi Dong <yidong@nvidia.com> * add test case data and retrieval use the same indexed data. handles no neighbor padding Signed-off-by: Yi Dong <yidong@nvidia.com> * address the reviewer comments Signed-off-by: Yi Dong <yidong@nvidia.com> * preserve some fraction data not used for retrieval index Signed-off-by: Yi Dong <yidong@nvidia.com> * added perplexity Signed-off-by: Yi Dong <yidong@nvidia.com> * remove the default batch limits Signed-off-by: Yi Dong <yidong@nvidia.com> * address review comment Signed-off-by: Yi Dong <yidong@nvidia.com> * added pad_id valdiation logics Signed-off-by: Yi Dong <yidong@nvidia.com> * fix no attention issue Signed-off-by: Yi Dong <yidong@nvidia.com> * comment the fix, waiting for the fix from apex Signed-off-by: Yi Dong <yidong@nvidia.com> * get rid of pre_decoder final layernorm Signed-off-by: Yi Dong <yidong@nvidia.com> * add index check Signed-off-by: Yi Dong <yidong@nvidia.com> * same implementation Signed-off-by: Yi Dong <yidong@nvidia.com> * fix merge error Signed-off-by: Yi Dong <yidong@nvidia.com> * fix style Signed-off-by: Yi Dong <yidong@nvidia.com> * use the dec num layers to encoder scaling Signed-off-by: Yi Dong <yidong@nvidia.com> * fix style Signed-off-by: Yi Dong <yidong@nvidia.com> * address reviewer comments Signed-off-by: Yi Dong <yidong@nvidia.com> * added the headscale option Signed-off-by: Yi Dong <yidong@nvidia.com> Co-authored-by: Eric Harper <complex451@gmail.com> Signed-off-by: Georg Kucsko <gkucsko@gmail.com>

…nce (NVIDIA#4220) * added retrieval index dataset Signed-off-by: Yi Dong <yidong@nvidia.com> * added retrieval db handling Signed-off-by: Yi Dong <yidong@nvidia.com> * create and load data Signed-off-by: Yi Dong <yidong@nvidia.com> * working chunk retrieval Signed-off-by: Yi Dong <yidong@nvidia.com> * retrieval fetch works Signed-off-by: Yi Dong <yidong@nvidia.com> * unit test passes Signed-off-by: Yi Dong <yidong@nvidia.com> * add option to run retrieval preprocess Signed-off-by: Yi Dong <yidong@nvidia.com> * slice into chunks Signed-off-by: Yi Dong <yidong@nvidia.com> * add script to build index Signed-off-by: Yi Dong <yidong@nvidia.com> * building faiss index works Signed-off-by: Yi Dong <yidong@nvidia.com> * speed up the index building Signed-off-by: Yi Dong <yidong@nvidia.com> * added knn map index file Signed-off-by: Yi Dong <yidong@nvidia.com> * workign build knn map Signed-off-by: Yi Dong <yidong@nvidia.com> * added docstring Signed-off-by: Yi Dong <yidong@nvidia.com> * add retro dataset Signed-off-by: Yi Dong <yidong@nvidia.com> * dataset test passes Signed-off-by: Yi Dong <yidong@nvidia.com> * added unittest asserts Signed-off-by: Yi Dong <yidong@nvidia.com> * added dataset Signed-off-by: Yi Dong <yidong@nvidia.com> * training of retro is working Signed-off-by: Yi Dong <yidong@nvidia.com> * remove unused imports Signed-off-by: Yi Dong <yidong@nvidia.com> * fix jenkins Signed-off-by: Yi Dong <yidong@nvidia.com> * added knn example data Signed-off-by: Yi Dong <yidong@nvidia.com> * better print format Signed-off-by: Yi Dong <yidong@nvidia.com> * configure the number of neighbors Signed-off-by: Yi Dong <yidong@nvidia.com> * removed non-used cfg Signed-off-by: Yi Dong <yidong@nvidia.com> * turn on normliazaiton option Signed-off-by: Yi Dong <yidong@nvidia.com> * add layer number offset Signed-off-by: Yi Dong <yidong@nvidia.com> * need to add one Signed-off-by: Yi Dong <yidong@nvidia.com> * use at leaset one layer Signed-off-by: Yi Dong <yidong@nvidia.com> * added inference unit test Signed-off-by: Yi Dong <yidong@nvidia.com> * encoder inference test pass Signed-off-by: Yi Dong <yidong@nvidia.com> * encoder inference is confirmed to work Signed-off-by: Yi Dong <yidong@nvidia.com> * handles another edge case Signed-off-by: Yi Dong <yidong@nvidia.com> * alige relative position to the context Signed-off-by: Yi Dong <yidong@nvidia.com> * chunked cross attention passes test Signed-off-by: Yi Dong <yidong@nvidia.com> * fixed chunk cross attention masked attention Signed-off-by: Yi Dong <yidong@nvidia.com> * fix the padding Signed-off-by: Yi Dong <yidong@nvidia.com> * chunked cross attention layer inference passes unit test Signed-off-by: Yi Dong <yidong@nvidia.com> * working on the decoder Signed-off-by: Yi Dong <yidong@nvidia.com> * one mile stone for decoder Signed-off-by: Yi Dong <yidong@nvidia.com> * decoder is working Signed-off-by: Yi Dong <yidong@nvidia.com> * make encoder infer behave nicely Signed-off-by: Yi Dong <yidong@nvidia.com> * added encoder decoder inference unittest Signed-off-by: Yi Dong <yidong@nvidia.com> * remove bad imports Signed-off-by: Yi Dong <yidong@nvidia.com> * make training work Signed-off-by: Yi Dong <yidong@nvidia.com> * fix failed unit test Signed-off-by: Yi Dong <yidong@nvidia.com> * remove unused variables Signed-off-by: Yi Dong <yidong@nvidia.com> * added run on GPU for unittest Signed-off-by: Yi Dong <yidong@nvidia.com> * add pad id to the preprocessing script Signed-off-by: Yi Dong <yidong@nvidia.com> * added doc string for indexed dataset Signed-off-by: Yi Dong <yidong@nvidia.com> * efficient deduplicate doc Signed-off-by: Yi Dong <yidong@nvidia.com> * add test case data and retrieval use the same indexed data. handles no neighbor padding Signed-off-by: Yi Dong <yidong@nvidia.com> * address the reviewer comments Signed-off-by: Yi Dong <yidong@nvidia.com> * preserve some fraction data not used for retrieval index Signed-off-by: Yi Dong <yidong@nvidia.com> * added perplexity Signed-off-by: Yi Dong <yidong@nvidia.com> * remove the default batch limits Signed-off-by: Yi Dong <yidong@nvidia.com> * address review comment Signed-off-by: Yi Dong <yidong@nvidia.com> * added pad_id valdiation logics Signed-off-by: Yi Dong <yidong@nvidia.com> * fix no attention issue Signed-off-by: Yi Dong <yidong@nvidia.com> * comment the fix, waiting for the fix from apex Signed-off-by: Yi Dong <yidong@nvidia.com> * get rid of pre_decoder final layernorm Signed-off-by: Yi Dong <yidong@nvidia.com> * add index check Signed-off-by: Yi Dong <yidong@nvidia.com> * same implementation Signed-off-by: Yi Dong <yidong@nvidia.com> * fix merge error Signed-off-by: Yi Dong <yidong@nvidia.com> * fix style Signed-off-by: Yi Dong <yidong@nvidia.com> * use the dec num layers to encoder scaling Signed-off-by: Yi Dong <yidong@nvidia.com> * fix style Signed-off-by: Yi Dong <yidong@nvidia.com> * address reviewer comments Signed-off-by: Yi Dong <yidong@nvidia.com> * added the headscale option Signed-off-by: Yi Dong <yidong@nvidia.com> Co-authored-by: Eric Harper <complex451@gmail.com> Signed-off-by: Hainan Xu <hainanx@nvidia.com>

yidong72 added 30 commits May 16, 2022 18:35

added retrieval index dataset

c13a39a

Signed-off-by: Yi Dong <yidong@nvidia.com>

added retrieval db handling

c9d952e

Signed-off-by: Yi Dong <yidong@nvidia.com>

create and load data

05ce8b8

Signed-off-by: Yi Dong <yidong@nvidia.com>

working chunk retrieval

616ce5e

Signed-off-by: Yi Dong <yidong@nvidia.com>

retrieval fetch works

01efab5

Signed-off-by: Yi Dong <yidong@nvidia.com>

unit test passes

b4721ec

Signed-off-by: Yi Dong <yidong@nvidia.com>

add option to run retrieval preprocess

234d832

Signed-off-by: Yi Dong <yidong@nvidia.com>

slice into chunks

caa7f4e

Signed-off-by: Yi Dong <yidong@nvidia.com>

add script to build index

b0e2e73

Signed-off-by: Yi Dong <yidong@nvidia.com>

building faiss index works

cb708a3

Signed-off-by: Yi Dong <yidong@nvidia.com>

speed up the index building

ac0103b

Signed-off-by: Yi Dong <yidong@nvidia.com>

added knn map index file

563c8ef

Signed-off-by: Yi Dong <yidong@nvidia.com>

workign build knn map

b2bc340

Signed-off-by: Yi Dong <yidong@nvidia.com>

added docstring

8533624

Signed-off-by: Yi Dong <yidong@nvidia.com>

add retro dataset

b66f8ac

Signed-off-by: Yi Dong <yidong@nvidia.com>

dataset test passes

8428c0b

Signed-off-by: Yi Dong <yidong@nvidia.com>

added unittest asserts

af093e9

Signed-off-by: Yi Dong <yidong@nvidia.com>

added dataset

471ad5f

Signed-off-by: Yi Dong <yidong@nvidia.com>

training of retro is working

bf4305b

Signed-off-by: Yi Dong <yidong@nvidia.com>

remove unused imports

b605ef5

Signed-off-by: Yi Dong <yidong@nvidia.com>

Merge branch 'main' into feature_retrieval_idx

5578085

fix jenkins

5514c41

Signed-off-by: Yi Dong <yidong@nvidia.com>

added knn example data

c5f273c

Signed-off-by: Yi Dong <yidong@nvidia.com>

better print format

081b336

Signed-off-by: Yi Dong <yidong@nvidia.com>

configure the number of neighbors

59e3b24

Signed-off-by: Yi Dong <yidong@nvidia.com>

removed non-used cfg

48a4739

Signed-off-by: Yi Dong <yidong@nvidia.com>

turn on normliazaiton option

c41fb9b

Signed-off-by: Yi Dong <yidong@nvidia.com>

add layer number offset

029e952

Signed-off-by: Yi Dong <yidong@nvidia.com>

need to add one

9ddcf3c

Signed-off-by: Yi Dong <yidong@nvidia.com>

use at leaset one layer

ac05ddd

Signed-off-by: Yi Dong <yidong@nvidia.com>

yidong72 and others added 9 commits May 25, 2022 23:37

Merge branch 'main' into feature_retrieval_idx

533e606

Merge branch 'main' into feature_retrieval_idx

e0e3fef

fix no attention issue

e20c2cd

Signed-off-by: Yi Dong <yidong@nvidia.com>

comment the fix, waiting for the fix from apex

67b330d

Signed-off-by: Yi Dong <yidong@nvidia.com>

get rid of pre_decoder final layernorm

cb56645

Signed-off-by: Yi Dong <yidong@nvidia.com>

add index check

a150625

Signed-off-by: Yi Dong <yidong@nvidia.com>

same implementation

30c3ad5

Signed-off-by: Yi Dong <yidong@nvidia.com>

Merge branch 'main' into feature_retrieval_idx

6adc88d

fix merge error

cd0a84a

Signed-off-by: Yi Dong <yidong@nvidia.com>

fix style

e4809f7

Signed-off-by: Yi Dong <yidong@nvidia.com>

use the dec num layers to encoder scaling

ab14327

Signed-off-by: Yi Dong <yidong@nvidia.com>

yidong72 and others added 4 commits May 31, 2022 20:59

fix style

43e10bc

Signed-off-by: Yi Dong <yidong@nvidia.com>

Merge branch 'main' into feature_retrieval_idx

62c15db

Merge branch 'main' into feature_retrieval_idx

edee75d

Merge branch 'main' into feature_retrieval_idx

4b92346

michalivne reviewed Jun 1, 2022

View reviewed changes

yidong72 added 3 commits June 1, 2022 22:15

Merge branch 'main' into feature_retrieval_idx

8b9bd90

address reviewer comments

a51e135

Signed-off-by: Yi Dong <yidong@nvidia.com>

added the headscale option

9a1b017

Signed-off-by: Yi Dong <yidong@nvidia.com>

MaximumEntropy approved these changes Jun 2, 2022

View reviewed changes

Merge branch 'main' into feature_retrieval_idx

a4a35a8

yidong72 merged commit 75c2d82 into main Jun 2, 2022

yidong72 deleted the feature_retrieval_idx branch June 2, 2022 13:11

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add RETRO indexed dataset and `set_inference_key_value_memory` inference #4220

Add RETRO indexed dataset and `set_inference_key_value_memory` inference #4220

yidong72 commented May 20, 2022

lgtm-com bot commented May 30, 2022

lgtm-com bot commented May 30, 2022

lgtm-com bot commented May 31, 2022

okuchaiev commented Jun 1, 2022

michalivne left a comment

michalivne Jun 1, 2022

yidong72 Jun 1, 2022

michalivne Jun 1, 2022

yidong72 Jun 1, 2022

MaximumEntropy left a comment

yidong72 commented Jun 2, 2022

Add RETRO indexed dataset and set_inference_key_value_memory inference #4220

Add RETRO indexed dataset and set_inference_key_value_memory inference #4220

Conversation

yidong72 commented May 20, 2022

What does this PR do ?

lgtm-com bot commented May 30, 2022

lgtm-com bot commented May 30, 2022

lgtm-com bot commented May 31, 2022

okuchaiev commented Jun 1, 2022

michalivne left a comment

Choose a reason for hiding this comment

michalivne Jun 1, 2022

Choose a reason for hiding this comment

yidong72 Jun 1, 2022

Choose a reason for hiding this comment

michalivne Jun 1, 2022

Choose a reason for hiding this comment

yidong72 Jun 1, 2022

Choose a reason for hiding this comment

MaximumEntropy left a comment

Choose a reason for hiding this comment

yidong72 commented Jun 2, 2022

Add RETRO indexed dataset and `set_inference_key_value_memory` inference #4220

Add RETRO indexed dataset and `set_inference_key_value_memory` inference #4220