[Examples] TPU-based training of a language model using TensorFlow
Author: sayakpaul
Pull request: yes
Labels: name: TensorFlow

name: Examples

name: TPU

This PR adds an example of performing (masked) language model training using TensorFlow and TPUs. The example is meant to act as a reference for the community on this topic. The following are the main components of the PR: 

* Tokenizer training script (for completeness) 
* TFRecords preparation script (recommended practice when using TPUs)
* Training script
* Evaluation / inference

The purpose of this separation (as opposed to having everything in a single script) is to allow the community to have isolated reference points for performing TPU-based training of our models, which I think is beneficial. 

The artifacts produced during this project can be found here: https://huggingface.co/tf-tpu.

* Tokenizer (trained from scratch): https://huggingface.co/tf-tpu/unigram-tokenizer-wikitext
* Model: https://huggingface.co/tf-tpu/roberta-base-epochs-500-no-wd

Cc: @Rocketknight1 @gante @amyeroberts 

Comments:
_The documentation is not available anymore as the PR was closed or merged._

@Rocketknight1 I incorporated the `group_texts()` utility that we discussed over Slack. Let me know if the changes look good to you. Most of it is copy-pasted from [here](https://github.com/huggingface/notebooks/blob/main/examples/language_modeling-tf.ipynb). 

[Here's](https://colab.research.google.com/gist/sayakpaul/adfaa9b45c9b56f222487995d0971645/scratchpad.ipynb) Colab Notebook where I verified these. 

@Rocketknight1 I took a deeper look into the TFRecord preparation script. I don't understand why there's a discrepancy in the following. 

While serializing the TFRecords, I am making each TFRecord shard has got a specific number of samples. When there are lesser samples for a TFRecord shard than the specified amount, that's fine. 

But when I load the TFRecords back and create a `tf.data.Dataset` out of them, the number of entries in the dataset (before batching) is much lesser. 


Here is a minimal Colab Notebook that demonstrates the issue: https://colab.research.google.com/gist/sayakpaul/b4b02f3f656c0041c93f6ba78c8e65fd/scratchpad.ipynb.

When you get a moment, could you take a look? 

Thanks @Rocketknight1 for your help in debugging https://github.com/huggingface/transformers/pull/21657#issuecomment-1468086926 (discussed internally via Slack). I am currently regenerating the TFRecord shards. I will update here once that's done.

@Rocketknight1 corrected TFRecord shards have been pushed to `gs://tf-tpu-training-resources`.

Here are the record counts per split:

* Train: 300917
* Validation: 626
* Test: 722

The TFRecords were generated with a block size of 512. 

@Rocketknight1 the training code looks good to me, except for a few things:

* Maybe we should scale the LR with the batch size? 
* Take `mlm_probability` as a CLI arg? 
* Modularize the dataset preparation code a bit? 

But all these are non-blockers. Let's do 4 - 5 training runs varying the number of epochs and the learning rate. 

@sayakpaul MLM probability added as an arg and I modularized the loading!

@Rocketknight1 started a training run with:

```bash
python3 train_model.py \
  --tokenizer tf-tpu/unigram-tokenizer-wikitext \
  --per_replica_batch_size 64 \
  --tpu_name local --tpu_zone us-central1 --gcp_project huggingface-ml --bfloat16 \
  --train_dataset gs://tf-tpu-training-resources/train --eval_dataset gs://tf-tpu-training-resources/validation \
  --num_epochs 100 \
  --output_dir roberta-base-epochs-100 --hub_model_id tf-tpu/roberta-base-epochs-100
```

@Rocketknight1 here's the final model trained with the command from [here](https://github.com/huggingface/transformers/pull/21657#issuecomment-1483729424):

https://huggingface.co/tf-tpu/roberta-base-epochs-100

When you try out examples in the widget of the model page ^, pass `[MASK]` instead of the default `<mask>`. The results are far from perfect (evident from the validation accuracy), though. 

@Rocketknight1 could you review [this PR](https://huggingface.co/tf-tpu/roberta-base-epochs-500-no-wd/discussions/1)? 

@sgugger thanks!

I addressed your comments. For https://github.com/huggingface/transformers/pull/21657#discussion_r1164017322, I will defer to @Rocketknight1. 

Merging since the failing tests are unrelated.