Fine-tuning the Entire RAG Architecture (including DPR retriever)
Author: shamanez
Pull request: yes

# What does this PR do?

The original RAG implementation is able to end-to-end train the question encoder and the generator. 
This extension enables the end-to-end training of RAG including the context encoder in the retriever component. 
Please read the [accompanying blog post](https://shamanesiri.medium.com/how-to-finetune-the-entire-rag-architecture-including-dpr-retriever-4b4385322552) for details on this implementation.

The original RAG code has also been modified to work with the latest versions of PyTorch lightning (version 1.2.10) and RAY (version ). All other implementation details remain the same as the [original RAG code](https://github.com/huggingface/transformers/tree/master/examples/research_projects/rag).
Read more about RAG  at https://arxiv.org/abs/2005.11401.

This code can be modified to experiment with other research on retrieval augmented models that include training of the retriever such as [REALM](https://arxiv.org/abs/2002.08909) and [MARGE](https://arxiv.org/abs/2006.15020). 


Reviewers  @patrickvonplaten  @lhoestq 

Comments:
> Great, I think we are very close to merging this PR :-)
> 
> Could we add a test in both `test_modeling_rag` and `test_retrieval_rag` ? We should test there that the model behaves correctly when `set_context_encoder_for_training` is set

on it!

Hello, I tried running the code in this pull request because the methodology is something I'm very interested in and I ran into a few issues. These may be things you are aware of but I just wanted to mention them in case you hadn't run into them.
1.  When launching the code in distributed mode (CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch --nproc_per_node=8) I received the error:

     handle = worker.core_worker.get_named_actor_handle(name)
  File "python/ray/_raylet.pyx", line 1496, in ray._raylet.CoreWorker.get_named_actor_handle
  File "python/ray/_raylet.pyx", line 157, in ray._raylet.check_status
ValueError: Failed to look up actor with name 'retrieval_worker_0'. You are either trying to look up a named actor you didn't create, the named actor died, or the actor hasn't been created because named actor creation is asynchronous.

Based on the traceback it seems to be related to line 733 in finetune_rag.py 
named_actors = [ray.get_actor("retrieval_worker_{}".format(i)) for i in range(args.num_retrieval_workers)]
If I run it with the normal python launch it goes fine however. This is with ray version 1.3.0, pytorch-lightning 1.2.10 and transformers version 4.7.0.dev0 which matches the requirements.txt file. This is an issue because I don't believe the sharded_ddp plugin works without the distributed launch.

2. pytorch_lightning.utilities.exceptions.MisconfigurationException: ModelCheckpoint(monitor='val_em') not found in the returned metrics: ['val_loss', 'val_avg_loss', 'val_avg_em', 'val_avg_gen_time', 'val_avg_gen_len', 'step_count', 'loss']. HINT: Did you call self.log('val_em', value) in the LightningModule?

This happens at the first validation step after training epoch 0. I believe that the values being passed on line 421 in the finetune_rag.py script are not named correctly? The metrics all have "_avg_" in the name however the monitored metric doesn't seem to have that and is just "val_em".

3. It seems that in some locations the ctx_encoder_tokenizer is a required keyword argument in some locations and not in others. I had change line 367 in retrieval_rag.py to: def __init__(self, config, question_encoder_tokenizer, ctx_encoder_tokenizer, generator_tokenizer, index=None, init_retrieval=True): adding the ctx_encoder_tokenizer otherwise it said it was missing the keyword argument.

4. I had to change the line 528 in modeling_rag.py from "self.context_encoder_training=False" to "self.context_encoder_training=True" in order to get it to properly do the context_encoding. I could have messed something else up in the training to cause it to not properly set this to True when setting the context encoder but I couldn't get it to work without doing this (threw the error KeyError: 'tokenized_doc_ids')

5. I had to add "from transformers import PreTrainedTokenizer" to retrieval_rag.py also because that doesn't seem to be imported anywhere in the file but is used in the file on line 552. This could be an issue with my transformer version but I still believe it would have to be in the import statements anyways no?

Any or all of these could be issues with how I'm running it but I figured I'd bring them to your attention because these were all the things I had to change in order to get it to run. I can provide more info on any/all of these if you would like but I figured I would give you a list of things I ran into since it hasn't been merged yet so i'd imagine not many people have tried to run it end to end. Thanks for adding this feature though; it's definitely going to be a big upgrade for those of us who are using the model on different datasets and use cases.


@calderma 

You have done an amazing exploration. I am really sorry :( that apart from the first issue, all other things being already fixed, but I did not push the latest version since I am conducting some experiments to update the README (original RAG vs Updated RAG). 



In the latest version, I added a dummy training example which makes things a lot more clear. 

For the first issue, I actually do not have a distributed system to test. I think it is something to do with PyTorch-lightning initialization. Try to add distributed parameters to the lightning trainer. **(Trainer(gpus=8, accelerator='ddp', num_nodes=4))**

**named_actors = [ray.get_actor("retrieval_worker_{}".format(i)) for i in range(args.num_retrieval_workers)]**

During the initialization of RAY workers, we create them only on the master DDP process. So the above line is there to get the retrieval workers for other ddp processes (let's say cuda 1, 2, ..) that have already created during the master process [(check this line)](https://github.com/shamanez/transformers/blob/rag-retriever-end2end/examples/research_projects/rag-end2end-retriever/finetune_rag.py#L712). 

As the error shows in your issue, what has happened is the initialization across the nodes somehow hasn't [shared](https://github.com/shamanez/transformers/blob/rag-retriever-end2end/examples/research_projects/rag-end2end-retriever/finetune_rag.py#L712). So obviously ray gives an error saying, it can't find an already initialized worker by its name. I think this has something to do with how pytorch-lightning execute multi-node training. So I would suggest you to follow their guidelines.  Something as follows.



Please let me know if this works, So I can include it in the latest version. 




Sure I should be able to try it tomorrow in the early afternoon. I will post in this thread afterwards.

> Sure I should be able to try it tomorrow in the early afternoon. I will post in this thread afterwards.

Since I have only one HPC cluster, I just run the process with CUDA_VISIBLE_DEVICES-0,1,2,3... bash script. It works fine for me.  Can you please elaborate a little bit on this statement " I don't believe the sharded_ddp plugin works without the distributed launch." ?

Sure I was referencing the plugin for fairscale with pytorch lightning:
https://medium.com/pytorch/pytorch-lightning-1-1-model-parallelism-training-and-more-logging-options-7d1e47db7b0b
I was under the impression to use that plugin it had to be launched with the distributed pytorch launch but honestly I've never tried it with just launching it via python. When I trained on the old RAG model i used the distributed command and passed --plugins ddp_sharded but i suppose it might just work with regular python launching. I don't currently have the ability to test it until tomorrow though.

I believe I had to alter the original RAG code to work with pytorch lightning 1.1 as it was based on 1.0 but I needed to use fairscale to use a batch size larger than 1. Unfortunately I no longer have access to the repository I was pushing to at that time, which was a few months ago.

@calderma  Exactly, I also think it is something to do with the launching of the script with multiple nodes.  

I am asking this because PL has changed a lot compared to the version that used in the original RAG. Especially plugins.  In the upgraded RAG, I had to remove [this method with plugins](https://github.com/shamanez/transformers/blob/rag-retriever-end2end/examples/research_projects/rag/finetune_rag.py#L84) and use Callbacks since PL at the moment is not that much encouraging us to use custom pluggings. 




I'd imagine that the distributed launch isn't worth worrying about then but I can still test it tomorrow if you would like to see if it makes the ray workers behave.

Thanks a lot.

Hello again. I pulled your latest changes and tried running it again. I didn't get the distributed training to work with those changes but I'm by no means an expert at that so maybe someone better at distributed pytorch-lightning could take a look at it. I did notice a few other things with regards to training.
1. It seems the --data_cache_dir option doesn't exist in the code anymore? looking through finetune_rag.py i didn't see it anywhere but when i looked back at previous commits it was there in the arguments.
2. I had to manually create the kb_shards directory to get it to run. Could have easily been an issue on my end regarding permissions or something.
3. I got the error:
    callbacks_rag.py", line 44, in get_checkpoint_callback
    every_n_val_epochs=1,  # maybe save a checkpoint every time val is run, not just end of epoch.
TypeError: __init__() got an unexpected keyword argument 'every_n_val_epochs'

There was an issue with my cluster apparently that should be resolved soon so I'll try again but I don't believe it was related to these errors. 

@calderma 

1. I did swamp the data_cache_dir with cache_dir parameter given in lightning_base.py.
2. Now you do not need to create kb-shards directory manually, check [this line, which creates them automatically](https://github.com/shamanez/transformers/blob/rag-retriever-end2end/examples/research_projects/rag-end2end-retriever/finetune_rag.py#L689)
3. I think you need to update PL. Or you can run the code by [swapping those parameters with **period**](https://pytorch-lightning.readthedocs.io/en/latest/extensions/generated/pytorch_lightning.callbacks.ModelCheckpoint.html#pytorch_lightning.callbacks.ModelCheckpoint).


For the distributed running, I checked with PyTorch lightning examples. Seem like it depends on how you have constructed the cluster. Code-wise we do not need to change anything other than adding num_nodes and accelerator parameters to the trainer. 
**trainer = Trainer(gpus=8, num_nodes=4, accelerator='ddp')**, which I already did. [ See these examples, it seems pretty straightforward](https://pytorch-lightning.readthedocs.io/en/latest/clouds/cluster.html#general-purpose-cluster). 

Now you can also run the original [RAG with new PyTorch lightning.](https://github.com/huggingface/transformers/pull/11806)




great sounds good. The only thing that might need to be switched with regard to the first one is I believe the finetune_rag_ray_end2end.sh script still passes "--data_cache_dir". Regarding the PL version I was going by what was in the requirements file. my pip list shows pytorch-lightning == 1.2.10 which seems to be what's in the requirements. Thanks for the help with the distributed PL!

Omg yeah. My bad!  Thanks a lot highlighting them. I will quickly change it
and push it.

On Sat, May 22, 2021, 11:09 calderma ***@***.***> wrote:

> great sounds good. The only thing that might need to be switched with
> regard to the first one is I believe the finetune_rag_ray_end2end.sh script
> still passes "--data_cache_dir". Regarding the PL version I was going by
> what was in the requirements file. my pip list shows pytorch-lightning ==
> 1.2.10 which seems to be what's in the requirements. Thanks for the help
> with the distributed PL!
>
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/pull/11655#issuecomment-846303671>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AEA4FGT7TIRQFUPWMOCHBSTTO3R3DANCNFSM44QW2KSA>
> .
>


Great! I'm going to run it over the weekend so I will let you know if I hit any other roadblocks or if it finishes without issue.



Hi Patrick and Quentin ,
**I added the testing files in the test_run folder.** [test_rag_new_features.sh ](https://github.com/shamanez/transformers/blob/rag-retriever-end2end/examples/research_projects/rag-end2end-retriever/test_run/test_rag_new_features.sh#L6 )tests if the new functions are working and test_finetune.sh trains with a dummy dataset.

Additionally, We also did a small experiment with the SQuAD dataset using all context passages as the knowledge base. The increase in the EM-Scores was better than we expected. Users also can compare these two methods.
Cheers, @patrickvonplaten  @lhoestq 

Just to close the loop on this, my test run of the end2end went smoothly and I had no additional roadblocks. Thanks!

@calderma  , I had some problems with version control. So  I created a revamp pull request .. can you just run in and let me know :) 
https://github.com/huggingface/transformers/pull/11893


sure i'll run it today

In this version I added the following:

1. Added test functions as requested in [this comment](https://github.com/huggingface/transformers/pull/11655#pullrequestreview-658830866). 
2. Added results section to the README.



Closing in favor of #11893

@calderma 


> 1. When launching the code in distributed mode (CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch --nproc_per_node=8) I received the error:
>    handle = worker.core_worker.get_named_actor_handle(name)
>    File "python/ray/_raylet.pyx", line 1496, in ray._raylet.CoreWorker.get_named_actor_handle
>    File "python/ray/_raylet.pyx", line 157, in ray._raylet.check_status
>    ValueError: Failed to look up actor with name 'retrieval_worker_0'. You are either trying to look up a named actor you didn't create, the named actor died, or the actor hasn't been created because named actor creation is asynchronous.
> 


@calderma  I think We found the problem when running RAG with RAY on distributed systems. 

In some distributed systems, **os.environ["NODE_RANK"]**  is a string but not an integer. 

So basically the if the condition can get mess up. So please upgrade the if condition inf finetune.py as follows:

`
        if ("LOCAL_RANK" not in os.environ or os.environ["LOCAL_RANK"] == 0) and (  "NODE_RANK" not in os.environ or int(os.environ["NODE_RANK"] )== 0        ):`

Hope this solves the problem!