[WIP] RAG end-to-end retriever training (with ray workers)
Author: shamanez
Pull request: yes

# What does this PR do?

 As mentioned in this [issue](https://github.com/huggingface/transformers/issues/9646), this PR adds the fine-tuning ability the Retriever in the original RAG implementation. 

This PR first updates the ctx_encoder and then initializes the index using RAY workers.


@lhoestq 
@amogkam

Comments:
@lhoestq  @patrickvonplaten 

I have already started doing the above changes you have mentioned.  Apart from that, I changed the following elements of the codebase. 

1. I used a dedicated RAY worker to compute embeddings for the dataset with an updated ctx encoder in this version.  However, the process of add_faiss_index gets very slow when running inside a ray worker (I think it is due to the need for multiprocessing threads). I tried to increase the number of CPU cores, but it is still very slow.  The computing of embeddings is an embarrassingly parallel task, where we can share the dataset between GPUs and compute them very fast. Nevertheless, it is hard to work with RAY when it comes to multiple GPUs. So I utilize the DDP process to compute embeddings using N number of dedicated GPUs that only do the embeddings calculation task.

2. Then I did a minor thing. Pytorch lightning has removed the DDP accelerators in their latest installation. Nevertheless, we can easily use the **on_sanity_check_start** callback to initialize the index when using RAY. I feel it is a lot cleaner.  

__________________________________________________________________________________________________

As per my experiments, at the moment end-to-end training process is stable.  I would love to double-check the following parts with your help. 


1. Re-loading of the updated index for the workers.
2. Re-initialization of the retrieval index.

Apart from them, I see **add_faiss_index** can get hours when the dataset consists of more than a million passages. My custom dataset has 8 million datasets. Is it normal, or should we able to improve it?   If we can improve its speed, this whole process can be very engineering friendly. 

Please let me know your ideas. I will quickly do the updated PR. 




For the records, we are discussing the indexing speed difference here: https://github.com/huggingface/datasets/issues/2046

@lhoestq 

I and @elliott-wen updated the codebase. Now the embedding update happens with a parallel process and we use the stale gradients to update the entire model (pretty-much similar to REALM).    

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

@shamanez do you need any help to get the test working?

Hi, @patrickvonplaten  I did change the code and got a stable end-to-end trainable RAG. I also did the changes you have mentioned. But I had to update the code with new pytorch lightning, especially they do not use plugging now.  So is that ok to upload the code with a new end-to-end fine-tune script and lightning base?

@patrickvonplaten  @lhoestq 

I also added all the details to a blog and I am happy to share it with you two. It includes all the changes I did in the RAG. (I also included your names since you guys helped me a lot :) )

Hi, that's good news !

> is that ok to upload the code with a new end-to-end fine-tune script and lightning base?

I think it could be a good idea to make the code compatible with the latest pytorch-lightning yes :)
Especially since many things we used in lightning_base don't work anymore, and that we can now hope pytorch-lightning to not do such radical changes again.
pinging @patrickvonplaten to confirm it's ok

@lhoestq  Thanks. 

BTW I read this new paper named [Retrieval Augmentation Reduces Hallucination in Conversation](https://arxiv.org/abs/2104.07567), which kinds of highlights the importance of RAG-like models in LM modeling.  So I do believe end-to-end fine-tuning can allow users to experiment with different components of the RAG architecture.


Since you guys helped me a lot in this process, is it okay to include your names in the blog? Here's a link to the unpublished draft blog post.



https://medium.com/@shamanesiriwardhana/end-to-end-rag-fine-tuning-with-huggingface-pytorch-lightning-and-ray-4b4385322552



Please let me know your thoughts.



Good job with this blog post draft ! Sure it's fine to mention us, thanks

Thanks.

On Sat, May 8, 2021, 04:42 Quentin Lhoest ***@***.***> wrote:

> Good job with this blog post draft ! Sure it's fine to mention us, thanks
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/pull/10410#issuecomment-834610009>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AEA4FGVY24QYGUQ6SR5T473TMQJ67ANCNFSM4YID3MKQ>
> .
>


> Hi, that's good news !
> 
> > is that ok to upload the code with a new end-to-end fine-tune script and lightning base?
> 
> I think it could be a good idea to make the code compatible with the latest pytorch-lightning yes :)
> Especially since many things we used in lightning_base don't work anymore, and that we can now hope pytorch-lightning to not do such radical changes again.
> pinging @patrickvonplaten to confirm it's ok

So is that ok we create a folder in researc_project naming, RAG-end-to-end-Retriever training


closing this with a new pull request. 

https://github.com/huggingface/transformers/pull/11655

@lhoestq  @patrickvonplaten  could you please let me know if there is anything to change in the recent pull request. 

Hey @shamanez,

Sorry for being so inactive here! I reviewed your newly opened PR :-)

Hey, it is totally fine. Thanks million times :).