Reformer
Author: patrickvonplaten
Pull request: yes

## Add the Reformer 
Paper: (https://arxiv.org/pdf/2001.04451.pdf)

### First steps to take: 

- [x] Copy Bert PT code to Reformer PT file.
- [x] Replace self-attention with LSH attention
- [x] Make forward pass work for Bert Layer

### Forward-Pass: Get 1-to-1 same outputs as original Flax code for forward pass
   - [x] for LSH attention layer
   - [x] for Bert Layer with RevNet
   - [x] for different attention masks
   - [x] for feed forward chunking layer
   - [x] for whole Reformer model
   - [x] for sinusoidal position encodings
   - [x] for axial position encodings
   - [x] for local blocking attention (chunked attention)
   - [x] for pretrained weights from official reformer model: ReformerLM model was trained using https://colab.research.google.com/github/google/trax/blob/master/trax/models/reformer/text_generation.ipynb and weights were loaded into https://huggingface.co/patrickvonplaten/reformer-crime-and-punish and checked that a single forward pass is identical. `predict_mem_len` had to be adapted to make functions equal. 
   - [x] Add optional attention mask 
   - [x] Add support for fp16
   - [ ] Speed up incremental generation. This is needed for generation and will not be trivial since the buckets have to ordered correctly and there is a chunk length parameter.


### Backpropagation:
   - [x] Make backpropagation work
   - [x] Check that backpropagation works with chunked feed forward layers
   - [x] Implement RevResLayers for backprop
      - [x] Implement code using  and 
      - [x] Get identical results for forward pass
      - [x] Make sure backprop works
      - [x] Implement bucket caching
      - [x] Implement random seed caching to have deterministic dropout for backward pass: https://github.com/RobinBruegger/RevTorch/pull/4
      - [ ] Make rev resnet work for multi-gpu training
   - [x] Check that RevReslayers backprop works on CPU
   - [x] Check that RevReslayers backprop works on GPU
   - [x] Get same gradients as original trax code
   - [x] Train model on crime-and-punishment text and check that model performs reasonable afterwards 

### Tokenizer
   - [x] Copy sentence piece tokenizer from T5
   - [x] Add vanilla sentence piece tokenizer for crime-and-punishment pretrained tokenizer: https://console.cloud.google.com/storage/browser/_details/trax-ml/reformer/cp.320.model
   - [ ] Check how many tokenizers are needed 
   - [ ] Get pretrained tokenizers

### Optimize time and memory efficiency
   - [x] Compare memory & time complexity to standard Bert: check https://github.com/huggingface/transformers/pull/3186
   - [x] Check and improve memory and speed when training
   - [ ] Move "on-the-fly" created masks in LSHSelfAttention to using them as an input
   - [ ] Optimizie all unnecessary calculations

### Pretrained Models
   - [ ] Check if pretrained model on C4 is added soon: https://github.com/google/trax/commit/b1f0c176a281d35e285137a45ff117b8c5495173

- [ ] Add Reformer / Bert in trax

Useful code resources: 
- Original trax code: https://github.com/google/trax/tree/master/trax/models/reformer
- Working trax notebook: https://github.com/google/trax/blob/master/trax/models/reformer/machine_translation.ipynb
- Working PyTorch implementation: https://github.com/lucidrains/reformer-pytorch
- Great to implement for backprop. https://github.com/lucidrains/reformer-pytorch/blob/master/reformer_pytorch/reversible.py
- Pretrained weights: https://console.cloud.google.com/storage/browser/trax-ml/reformer

Useful blog/paper resources:
- Original paper: https://arxiv.org/pdf/2001.04451.pdf
- Google AI blog: https://ai.googleblog.com/2020/01/reformer-efficient-transformer.html
- Good blog post 1: https://www.pragmatic.ml/reformer-deep-dive/
- Good blog post 2: https://towardsdatascience.com/illustrating-the-reformer-393575ac6ba0

Previous Discussions:
- #2341 

## Update 

The code is clean and ready for review now.
Small ToDos before merging:
- [x] Fill in TODOs in docs 
- [ ] Check whether more pre-trained weights can be used
- [ ] Train on fp16 once
- [ ] Update notebook showing how to use Reformer

### Review

I added quite some docstrings to explain the new methods introduced by the Reformer (Axial Position Encoding, LSH Attention, Local Attention, Feed Forward chunking), so it might be better to first go through the doctsrings. Docstrings are easier to read when switching to this branch and creating the docs locally.

Comments:
Memory complexity ReformerLayer
 vs BertLayer: 
![Figure_Memory_GPU_bs_1](https://user-images.githubusercontent.com/23423619/78665693-7fc25c80-78d6-11ea-9795-03fd86e9d709.png)




Time complexity ReformerLayer vs. BertLayer:

![Figure_Memory_GPU_bs_1_time](https://user-images.githubusercontent.com/23423619/78665764-9963a400-78d6-11ea-94b4-c2e2591bf579.png)


## Experiment

I tested training the Reformer model on 0.5M tokens per sample on the novel "Crime and Punishment" using conventional LM training. I essentially translated the official trax notebook: https://colab.research.google.com/github/google/trax/blob/master/trax/models/reformer/text_generation.ipynb into hugging face code: https://colab.research.google.com/drive/1jR6hA2CQXDbucJXdiDXhmxmyoQmM2Pws

The only differences to the official notebook are:

- The gradient is accumulated over 8 samples and then updated whereas in the official notebook 8 TPUs are used and the gradient is calculated in parallel and then averaged together.

- The learning rate is 0.005 instead of 0.01 (because already at 0.005, the gradient seems to become too big).

## Results

My training starts similarly around **6.2** and goes down smoothly in the beginning.
At some point though the gradient seem to explode and the loss goes up again and that even at a learning rate of "only" 0.05.

The attached plots are here:

### Loss
![eval_loss](https://user-images.githubusercontent.com/23423619/80461705-f979c300-8935-11ea-8fc7-9e97c2ac8de6.png)

### Accuracy
![eval_accuracy](https://user-images.githubusercontent.com/23423619/80461748-07c7df00-8936-11ea-9c7d-02d65de9cee3.png)

### Learning rate (cosine scheduler)
![learning_rate](https://user-images.githubusercontent.com/23423619/80461760-0e565680-8936-11ea-8c1f-3a9aec8bcc39.png)

When lowering the learning rate more, to **0.0005** e.g. the loss keeps going down but only reaches something around 2.3 in the end. 

**Comparison**

The training in the official trax notebook is very smooth. 
Loss starts at **6.2** something and goes down smoothly to **0.8** while the accuracy reaches **>80%** in the end for a learning rate of **0.01**.

## Analysis

- It is confirmed that the forward pass is identical with the trax implementation thanks to integration tests. Things that are not fully tested for the backward pass are:

    - **Dropout**: the dropout used in the official trax library does not seem to correspond to the "usual" `nn.Dropout` used in PyTorch but sometimes drop specific dimensions only or whole matrices. It is tested though that the dropout used here is deterministic for both the "normal" forward pass and the forward pass used in the backward pass to recalculate the activations, by means of setting the random seed used for the first forward pass. Nevertheless, there could still be small bugs.

    - **Reversible Layers**: Because Reformer uses reversible layers, I had to fiddle with a customized backward function here. This is IMO quite prone to errors. I checked multiple times that from a logical point of view everything is correct and compared my code with: https://github.com/RobinBruegger/RevTorch and https://github.com/lucidrains/reformer-pytorch which do similar / the same architecture. IMO, it quite hard to test this for correctness. One could also write the whole code without having reversible layers and then see whether the gradient is the same (Seems actually not like a bad idea to me). 

   - **Attention mask**: The official trax code does not seem to use a user-specific attention mask for the LSH Attn Layer, but only for the Local Attn Layer. I tested that the attn mask is correct for the local attention task by integration tests and checked that the attn mask for the LSH layer works correctly (input with mask gives the same result as input without mask), but maybe the LSH Attn mask has to be removed. But don't really see a reason why ?! 

  - **Initialization**: The initialization scheme used in the trax library is different from what is normally done in `transformers`, so there are small changes in my code. But I doubt that this is an issue, especially since the training looks very similar in the beginning.

   - **Training parameters**: It might also be simply due to different training / optimization parameters. Maybe there are some under-the-hood training parameters that I didn't notice (special gradient clipping, ...)

> https://colab.research.google.com/drive/1jR6hA2CQXDbucJXdiDXhmxmyoQmM2Pws

Tried to train model over longer time, but getting [error](http://prntscr.com/s79e04)

> Forward got unexcepted keyword "lm_labels" after calling trainer.train()

P: Fixed the typo. I will change the model into half-precision soon so that the memory will be sufficient :-) 

I get some good results with the following parameters: https://gist.github.com/flozi00/b491b41a9865733e5f8bb4032c313540

the best eval loss is about 1.654, but is increasing now again the same as yours
will have a look in a few hours again


![Screenshot_1](https://user-images.githubusercontent.com/47894090/80512980-a415d400-897e-11ea-94c8-67b3cea5cbd2.png)


> I get some good results with the following parameters: https://gist.github.com/flozi00/b491b41a9865733e5f8bb4032c313540
> 
> the best eval loss is about 1.654, but is increasing now again the same as yours
> will have a look in a few hours again
> 
> ![Screenshot_1](https://user-images.githubusercontent.com/47894090/80512980-a415d400-897e-11ea-94c8-67b3cea5cbd2.png)

Awesome that's already much better than what I got! If you manage to get it under 1 (loss) / >75% (accuracy) that would be great. Also feel free to change the hyper-parameters as you wish! Especially the adam betas and co.

I also added support for fp16 - so the notebook now only needs 8GB of RAM.

(You might have to reset the environment and re-install the github branch though)

Sounds very great.
Trying to decrease sequence length, cause while increasing number of hashes or heads getting memory error.
Training on 24GB GPU

Read that 4 hashes are good and 8 brings the best quality.

Trained on some configurations now and everytime the loss goes to ~1 but then increases to 4 very fast and keeps on there for minimum 1000 steps.
Any idea about it ?

> Sounds very great.
> Trying to decrease sequence length, cause while increasing number of hashes or heads getting memory error.
> Training on 24GB GPU
> 
> Read that 4 hashes are good and 8 brings the best quality.
> 
> Trained on some configurations now and everytime the loss goes to ~1 but then increases to 4 very fast and keeps on there for minimum 1000 steps.
> Any idea about it ?

My guess is that since it's such a small dataset (0.5M tokens is tiny) the model needs very well-calibrated hyperparameter tuning. When the learning rate is low enough, this actually does not happen anymore but also the loss only gets to about ~2. But I ran very few experiments and didn't do any hyperparameter search.
Also, I use slightly different dropouts, then were used in the official code so maybe using weight decay instead of dropout could work better. 

Will check that the gradients are correct in the next days and then hopefully be ready soon. 

@patrickvonplaten I'm excited to see a lot of progress here!

The loss curves above could be due to poor hyperparameter choice, but they're also very similar to what you see when the reverse pass of the network doesn't match the forward pass. For example, failing to cache hash bucket assignments (for exact re-use in the backward pass) leads to a failure mode with loss rebounds very similar to the figures you posted above. I also once had a bug where the wrong random seed was used for dropout in the backward pass, which IIRC manifested itself in the same way.

> @patrickvonplaten I'm excited to see a lot of progress here!
> 
> The loss curves above could be due to poor hyperparameter choice, but they're also very similar to what you see when the reverse pass of the network doesn't match the forward pass. For example, failing to cache hash bucket assignments (for exact re-use in the backward pass) leads to a failure mode with loss rebounds very similar to the figures you posted above. I also once had a bug where the wrong random seed was used for dropout in the backward pass, which IIRC manifested itself in the same way.

Thanks for taking a look @nkitaev. I just found a bug in the `backward()`. I now have 1-to-1 the same gradients as your trax code. Will retrain tonight and should get better results :-) 

# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/3351?src=pr&el=h1) Report
> Merging [#3351](https://codecov.io/gh/huggingface/transformers/pull/3351?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/7b5bec373ca2b442a7ac8ac46f8eac6e8003e2ae&el=desc) will **not change** coverage.
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/3351/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/3351?src=pr&el=tree)

```diff
@@           Coverage Diff           @@
##           master    #3351   +/-   ##
=======================================
  Coverage   79.13%   79.13%           
=======================================
  Files         117      117           
  Lines       19517    19517           
=======================================
  Hits        15444    15444           
  Misses       4073     4073           
```



------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/3351?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/3351?src=pr&el=footer). Last update [7b5bec3...7b5bec3](https://codecov.io/gh/huggingface/transformers/pull/3351?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Training looks good now on Crime and Punishment. To verify that training works, the model was trained on over 1200 steps and with little regularization.

### Eval loss 

![eval_loss-_1_](https://user-images.githubusercontent.com/23423619/80860924-69a48380-8c6b-11ea-80d7-d76f4d20974e.png)

### Eval accuracy

![eval_accuracy-_1_](https://user-images.githubusercontent.com/23423619/80860927-70cb9180-8c6b-11ea-8a16-827f65d38704.png)


https://gist.github.com/flozi00/b491b41a9865733e5f8bb4032c313540

This gist contains two notebooks, one of them with trainings batch = 2 --> error
in the other I tried to train model with pre-configured parameters, sequence length 4096 --> error

Is it mistake by me ?

> https://gist.github.com/flozi00/b491b41a9865733e5f8bb4032c313540
> 
> This gist contains two notebooks, one of them with trainings batch = 2 --> error
> in the other I tried to train model with pre-configured parameters, sequence length 4096 --> error
> 
> Is it mistake by me ?

how did you get past 
```
# get a pretrained tokenizer
tokenizer = ReformerTokenizer.from_pretrained("patrickvonplaten/reformer-crime-and-punish")
```

@lapolonio
I just used the notebook posted by patrickvonplaten four days ago.

https://colab.research.google.com/drive/1jR6hA2CQXDbucJXdiDXhmxmyoQmM2Pws

The notebook is still under construction, so I would not waste too much time playing around with it at the moment @lapolonio @flozi00. 

Thanks a lot for your good comments and remarks @lapolonio and @flozi00 :-) 

This is looking awesome, thanks! is there plans to add an encoder-decoder version? 

> This is looking awesome, thanks! is there plans to add an encoder-decoder version?

Yes, this should soon be possible with the encoder-decoder framework

@patrickvonplaten awesome! is there an issue or a PR I can follow for that? 

> @patrickvonplaten awesome! is there an issue or a PR I can follow for that?

Not yet, this will probably still need 1,2 weeks :-) 

Is it possible or are there any plans to implement reformer for question answering too ?
seq2seq and QA could be very great tasks for it

> Is it possible or are there any plans to implement reformer for question answering too ?
> seq2seq and QA could be very great tasks for it

Yeah, I will add a cross attention layer in another PR and then the Reformer can be used as a seq-2-seq model with our Encoder-Decoder framework: https://huggingface.co/transformers/model_doc/encoderdecoder.html

I'm not familiar with the Encoder-Decoder framework after the cross attention layer is added can the decoder be BertForSequenceClassification? Where do I ask questions like this?


@patrickvonplaten  Based on your merge, it seems like the input size for each batch is fixed in order to match the product of axial position embedding size? I am correct?

> @patrickvonplaten Based on your merge, it seems like the input size for each batch is fixed in order to match the product of axial position embedding size? I am correct?

For training, yes that's correct. For inference the input_size can also be smaller. Also check out: https://huggingface.co/transformers/model_doc/reformer.html

@patrickvonplaten , I wanted to train a language model for reformers on a custom dataset.
What are the steps, and any sample notebooks available for the same

Hi @prajwal-PHAI, there are a lot of [community notebooks covering T5 finetuning](https://github.com/huggingface/transformers/tree/master/notebooks#community-notebooks).

Thanks @LysandreJik 
I was running into error loading other datasets, which were not there in the nlp library.

hey. thanks for your amazing work!
I'm running into error while trying the colab example:
https://colab.research.google.com/drive/1jR6hA2CQXDbucJXdiDXhmxmyoQmM2Pws#scrollTo=WskGtnXsnWdu

the problem is that it doesn't recognize the apex package:

ImportError                               Traceback (most recent call last)
<ipython-input-29-30584d4c4987> in <module>()
     11 
     12 # train
---> 13 trainer.train()

/usr/local/lib/python3.6/dist-packages/transformers/trainer.py in train(self, model_path)
    384         if self.args.fp16:
    385             if not is_apex_available():
--> 386                 raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
    387             model, optimizer = amp.initialize(model, optimizer, opt_level=self.args.fp16_opt_level)
    388 

ImportError: Please install apex from https://www.github.com/nvidia/apex to use fp16 training.

though I installed it...anyone know what to do?


Linking a related git issue #16972. cc @patrickvonplaten