[RAG] Add Ray implementation for distributed retrieval
Author: amogkam
Pull request: yes

# What does this PR do?
This PR adds a new distributed retriever implementation for RAG built on Ray, as an alternative to the current retriever implementation that uses torch.distributed. With Ray it's possible to load the index on multiple processes instead of just the rank 0 training worker, allowing fine tuning to scale out better to multiple GPUs, and also allowing the index to potentially be fit in GPU memory. This also removes a core dependency on Pytorch, allowing a Tensorflow implementation of `finetune.py`.

This PR also makes changes to support finetune.py with Pytorch Lightning >v1.0.

A benchmark of Pytorch distribtued retrieval vs. Ray distributed retrieval
![image](https://user-images.githubusercontent.com/8068268/99327675-4ae49380-282f-11eb-8126-6250558fd795.png)

## Implementation Details
In the current Pytorch retrieval implementation, the index is loaded once on just the rank 0 training workers. Training worker 0 gathers the inputs from all other workers, performs the index lookup, and scatters the results back to the other workers.
![image](https://user-images.githubusercontent.com/8068268/99328533-238ec600-2831-11eb-9d26-e153cd0de74f.png)

With the Ray implementation, the index is loaded on *separate* processes, which are referred to as Ray actors. Each training worker randomly selects a retrieval actor to query for documents and Ray handles all the communication between the processes. Because the index can be loaded in *multiple* processes, training can scale up since no synchronization needs to happen for the index lookup.
![image](https://user-images.githubusercontent.com/8068268/99328651-618bea00-2831-11eb-990b-f1727df15190.png)

Note that Pytorch Lightning is still handling distributed *training*, but Ray manages distributed *retrieval*. Because PTL calls the entire training script under the hood multiple times, we have to use Ray's named actors feature (https://docs.ray.io/en/master/actors.html?highlight=named%20actors#named-actors) allowing the retrieval actors to be referenced by all training processes. The use of named actors is necessitated by how PTL handles distributed training, and a simpler approach could probably be used for a Tensorflow implentation.

## Testing Strategy
Unit tests were added to `test_distributed_retriever.py`. Note that the local Ray cluster for the tests had to be started with `local_mode=True` because the test file modifies `sys.path` and these changes are not propagated to remote processes. See https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder for more info.
<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes # (issue)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to the it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors which may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

 albert, bert, XLM: @LysandreJik
 GPT2: @LysandreJik, @patrickvonplaten
 tokenizers: @mfuntowicz
 Trainer: @sgugger
 Benchmarks: @patrickvonplaten
 Model Cards: @julien-c
 Translation: @sshleifer
 Summarization: @sshleifer
 examples/distillation: @VictorSanh
 nlp datasets: [different repo](https://github.com/huggingface/nlp)
 rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
 Text Generation: @patrickvonplaten, @TevenLeScao
 Blenderbot, Bart, Marian, Pegasus: @sshleifer
 T5: @patrickvonplaten
 Rag: @patrickvonplaten, @lhoestq
 EncoderDecoder: @patrickvonplaten
 Longformer, Reformer: @patrickvonplaten
 TransfoXL, XLNet: @TevenLeScao, @patrickvonplaten
 examples/seq2seq: @sshleifer
 examples/bert-loses-patience: @JetRunner
 tensorflow: @jplu
 examples/token-classification: @stefan-it
 documentation: @sgugger
 -->


Comments:
Hi ! This looks awesome :)
I was about to create a PR that fixes the init_ddp_connection in finetune.py and that adds a test script to make sure the finetuning script works as expected. With minimal changes on my side I can easily reduce conflicts between our two changes to finetune.py (I guess I'll just reuse the CustomAccelerator). Does that sound good to you ?

@lhoestq yes that sounds great!

@amogkam 


Hi seems like finetune.sh is not working in multi gpu training.

@shamanez Hmm that's odd, I was able to get this working on a single node with 4 GPUs. Do you have a stack trace?

> @shamanez Hmm that's odd, I was able to get this working on a single node with 4 GPUs. Do you have a stack trace?

I tried to run without RAY , but with pytorch DDP. Here is the error I got.


![Selection_006](https://user-images.githubusercontent.com/16892570/99744800-4117a480-2b3d-11eb-8757-498cff2819c3.png)
![Selection_007](https://user-images.githubusercontent.com/16892570/99744801-42e16800-2b3d-11eb-8a84-7ad9fdefb703.png)


This is the command-line argument I used, 
`python examples/rag/finetune.py     --data_dir ./examples/rag/test_data/dummy_seq2seq     --output_dir ./examples/rag/outputs   --model_name_or_path facebook/rag-token-base      --model_type rag_sequence       --do_train     --do_predict     --n_val -1     --val_check_interval 0.25     --train_batch_size 1     --eval_batch_size 1     --max_source_length 128     --max_target_length 25     --val_max_target_length 25     --test_max_target_length 25     --label_smoothing 0.1     --dropout 0.1     --attention_dropout 0.1     --weight_decay 0.001     --adam_epsilon 1e-08     --max_grad_norm 0.1     --lr_scheduler polynomial     --learning_rate 3e-05     --num_train_epochs 100     --warmup_steps 500     --gradient_accumulation_steps 1    --index_name custom    --passages_path ./examples/rag/data/my_knowledge_dataset    --index_path  ./examples/rag/data/my_knowledge_dataset_hnsw_index.faiss --gpus 2
`


I think this has to do with using a custom index which I didn't try out. Can you try with just the wiki_dpr index to confirm? It seems like the training workers are expecting a tensor of type float, but a tensor of type double is being sent instead. I think the fix might just be to set an explicit target_type in line 137 of distributed_pytorch_retriever.py- @lhoestq does this seem right?

Ok I will also give it a try

On Fri, Nov 20, 2020, 15:30 Amog Kamsetty <notifications@github.com> wrote:

> I think this has to do with using a custom index which I didn't try out.
> Can you try with just the wiki_dpr index to confirm? It seems like the
> training workers are expecting a tensor of type float, but a tensor of type
> double is being sent instead. I think the fix might just be to set an
> explicit target_type in line 137 of distributed_pytorch_retriever.py-
> @lhoestq <https://github.com/lhoestq> does this seem right?
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/pull/8583#issuecomment-730805824>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AEA4FGVO47GE3WRYB2VWLLTSQXIFZANCNFSM4TYGZUWA>
> .
>


@lhoestq now that https://github.com/huggingface/transformers/pull/8585 is merged, should I mark this PR as ready for review?

Yes indeed ! Feel free to set this PR to ready for review

Also it looks like the CI fails because of a failed import of `ray`.
To fix that you need to move the import of ray into the test functions decorated with `require_distributed_retrieval `.

You should also add `ray` to the test dependencies, or the test will simply be ignored

@lhoestq CI is passing now!

@lhoestq any ETA on when this PR can get reviewed? Thanks

Hi ! I've already started to look at the changes and it looks pretty good so far :) I'll finish my review soon, probably tomorrow

Awesome thanks!

@sgugger it would be cool if you could review as this changes some things in the trainer/integrations.

Hi @lhoestq @sgugger I addressed the feedback you guys gave. Do you think you can take another look? Thanks

Hi there, sorry for the delay. Could you close and reopen your PR? Because of a bad force-push on our side, the diff has become unreadable. Also, the examples folder has slightly changed structure, so you might need to move the folder.

Ping me, @patrickvonplaten  and @LysandreJik on the PR you reopen and we'll look at it quickly.

Opened a new one here: https://github.com/huggingface/transformers/pull/9197!