TFTrainer with TPUs: Here's a suggestion on getting it to work
Author: alexorona
Pull request: no
Labels: name: wontfix

## Environment info
- `transformers` version: 3.0.2
- Platform: Linux-4.19.112+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.6.0+cu101 (False)
- Tensorflow version (GPU?): 2.3.0 (False)
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: Yes

## Analysis and Temporary Solution
The approach in `TFTrainer` and `TFTrainingArguments` is really good, but it's not working right now on TPUs. It looks like we need to do some work on updating the trainer. There are a number of errors on this, the common being gradient accumulation (#6479) and `Unable to parse tensor proto`. Since Julien is on vacation, here's some things I did to get it to train on Colab with TPUs. It's hacky, but should be able to get it to work if you're anxious to use TPUs in until Julien has a fix:

- The `strategy` loading order in `TFTrainingArguments` and `TFTrainer` doesn't play well with a typical workflow (process data, create `training_args`, load model and pass to `TFTrainer`). The model needs to be loaded after the strategy has been initialized, and right now the strategy is being initialized inside of `TFTrainer.` 
- Shuffle, batch etc. need to be called prior to instantiating the strategy. I think this has something to do with the way the strategy is defined in `TFTrainingArguments`.
- Calling `training_args.training_batch_size` automatically calculates the number of TPU cores. Unfortunately, this causes the strategy to initialize, so this cannot be used to calculate `total_train_batch_size` with the current strategy implementation because it will prematurely initialize before shuffle, batch, etc. are done.
- To avoid the `Unable to parse tensor proto`, shuffle, batch etc. will need to be pulled from `TFTrainer`. They're handled by the `TFTrainer` method `get_train_tfdataset`. With the current strategy implementation in TFTrainingArguments, you'll need to do that after shuffle, batch and before loading the model.

## Example with GPT2
Here's a example implementing the above changes:

```
# Note: you'll need to build transformers from source

# Grab a temporary version of TFTrainer with get_train_tfdataset pulled out
git clone https://github.com/alexorona/lm_tf_trainer
from lm_tf_trainer import LMTFTrainer

# Pulled out of TFTrainer
def get_train_tfdataset(train_dataset, 
                        training_args,
                        train_batch_size, 
                        gradient_accumulation_steps, 
                        dataloader_drop_last = False, 
                        seed = 40):
    total_train_batch_size = train_batch_size * gradient_accumulation_steps
    num_train_examples = tf.data.experimental.cardinality(train_dataset).numpy()

    if num_train_examples < 0:
        raise ValueError("The training dataset must have an asserted cardinality")
    ds = (
        train_dataset.repeat()
        .shuffle(num_train_examples, seed=seed)
        .batch(total_train_batch_size, dataloader_drop_last)
        .prefetch(tf.data.experimental.AUTOTUNE)
    )

    return training_args.strategy.experimental_distribute_dataset(ds), num_train_examples

# Get Training Args
training_args, num_train_examples = TFTrainingArguments(...) # Create a normal training_args object

# Manual settings to avoid prematurely initializing the strategy
tpu_cores = 8
train_batch_size =  tpu_cores * training_args.per_device_train_batch_size

# Formatting tf dataset from lists of different kinds of inputs
input_ids = tf.convert_to_tensor(train_input_ids) #  train_input_ids[0] is a list of input ids and train_input_ids is a list of lists
attention_mask = tf.convert_to_tensor(attention_mask) #  as above
token_type_ids = tf.convert_to_tensor(token_type_ids) #  as above
labels = tf.convert_to_tensor(labels) #  as above
train_inputs = {'input_ids': input_ids, 'attention_mask': attention_mask, 'token_type_ids': token_type_ids}
train_dataset = tf.data.Dataset.from_tensor_slices((train_inputs,  train_labels))

# Now, call the function to do shuffle, batch and initialize the strategy
train_ds = get_train_tfdataset(train_dataset = train_dataset ,
                                    training_args = training_args,
                                    train_batch_size = train_batch_size ,
                                    gradient_accumulation_steps = training_args.gradient_accumulation_steps
                                    )

# Then, load the model with the strategy
with training_args.strategy.scope():
    model = TFGPT2LMHeadModel.from_pretrained('gpt2-medium')

# Now, train it
trainer = LMTFTrainer(args = training_args,
        model = model,
        num_train_examples = num_train_examples,
        total_train_batch_size = 8,
        train_dataset = train_ds)

trainer.train()
```


Comments:
cc @jplu for when he comes back

Thanks a lot @alexorona for this explicit issue.

Indeed the TF Trainer was working with usual TPU creation on GCP but not with Colab and investigating on this was one of my priotities when back from vacation. Apparently, TPU on Colab is more restrictive which is good to have a better implementation :+1: 

There are two small things we still have to think about from your proposal:

1) The preprocessing is not anymore in the Trainer. A solution might be to make the function static to be used without the need to instanciate the Trainer.
2) We have to be careful when calling `training_args.training_batch_size` only when not running on TPU and use the `tpu_num_cores` argument instead.

Can you train on TPU without loading tfrecords from a remote bucket?

I usually got an error `[local] not supported` and assumed that TPU does not support loading datasets directly.

You cannot load/save data on your local environment, everything must be in a GCS bucket.

There is some sample notebook on how to finetune TFGPT2LMHeadModel and about making that tfrecords. 

@jplu Great points, Julien. The proposal above is just a temporary work-around. From a user perspective, there really aren't any options in `get_train_tfdataset` that haven't already been declared elsewhere, so this is a routine task with no value in exposing it to the user. Therefore, it should be hidden _somewhere_. The question is whether that _somewhere_ is in the `TFTrainer` or in `TFTrainingArguments`. From a library management perspective, there are a lot of considerations, including how similar `TFTrainer` and `TFTrainingArguments` are to `Trainer` and `TrainingArguments` for pytorch. You want these classes to behave as similarly as possible. With that in mind, here are the options from best to worst:
1. See if there's a way to modify the current `TFTrainingArguments` tpu initialization procedure so that `get_train_tfdataset` can be left in `TFTrainer`. The model is still likely to be initialized outside of the scope, so a full-proof way of dealing with this is to re-initialize the model when `trainer.train()` is called by adding something like this in `train.train()`:
```
with args.strategy.scope():
     self.model = self.model
```
2. Barring that, it might be possible to initialize the strategy when `TFTrainingArguments` is first declared. In that case, `get_train_tfdataset` could be placed inside of `TFTrainingArguments`. We'd also need to know in the documentation that the model has to be loaded after `TFTrainingArguments` and with the clause `with training_args.strategy.scope():` coming before the line that loads the model.

@volker42maru I haven't had any problems with loading TF data records directly. Can you restructure so that the dataset is something like `tf.data.Dataset.from_tensor_slices((train_inputs,  train_labels))`? Are you sure your batch size is equal to at least the the number of tensor cores and you're calling` strategy.experimental_distribute_dataset(dataset)` somewhere? I've been able to load and transform data just fine on Colab. You can also connect to Google Drive too and use it as a disk with:
```
from google.colab import drive
drive.mount('/content/drive')
```

@jplu @alexorona 
Hey i want small help please see my colab notebook. i am trying to finetune gpt2 its showing training but not returning loss and please confirm me that its using tpu or not.
https://colab.research.google.com/drive/1IqXH0_VZ8LqgnbgjP3GqRXecChVitcea?usp=sharing
If i am doing something wrong please let me know.

Thanks @alexorona! I'm gonna investigate to be able to create the strategy in `trainer.train()` once all the datasets have been created and not in `TFTrainingArguments` anymore.

I tried a fix in the PR #6880.

@alexorona Can you try it and tell me if it works for your env?

@jplu  I'm getting that `Unable to parse tensor proto` error. Did you merge the changes into TF Trainer already?

No it is not merged, that's why if you could try the PR, it would be nice to have your feedback if it works or not.

@jplu Your approach looks great! I setup two dev notebook so you can see the remaining challenges:
- It looks like there's an expected `args.training_args` attribute that isn't there. Maybe the changes to `TFTrainingArguments` didn't make it to the fork? I had to revert most instances of `self.args.training_args` to `self.args` to get to the next step.
- `from_pt=bool(".bin" in self.args.model_name_or_path)` was throwing an error, but this might be due to `TFTrainingArguments` problem above.
- The current strategy implementation is running out of resources on models we know it can train on (gpt2-medium 1024 tokens). This is speculation, but it might be because `strategy.experimental_distribute_dataset(ds`) isn't being used in `get_train_tfdataset` anymore.
- `tb_writer` was also causing problems, so I had to comment that out too
- `self.model.ckpt_manager.save()` is throwing `File system scheme '[local]' not implemented` 
- Special tokens are sometimes added to the model, especially when using this for dialogue-style generation. Maybe add a parameter `tokenizer_length = None` to the class on `__init__` and then replace with this: 

```
        if not self.model:
            with self.strategy.scope():
                self.model = self.model_class.from_pretrained(
                    self.args.model_name_or_path,
                    # from_pt=bool(".bin" in self.args.model_name_or_path),
                    config=self.config,
                    cache_dir=self.argscache_dir,
                )
                if self.tokenizer_length:
                    self.model.resize_token_embeddings(self.tokenizer_length)
```

> * It looks like there's an expected `args.training_args` attribute that isn't there. Maybe the changes to `TFTrainingArguments` didn't make it to the fork? I had to revert most instances of `self.args.training_args` to `self.args` to get to the next step.
> * `from_pt=bool(".bin" in self.args.model_name_or_path)` was throwing an error, but this might be due to `TFTrainingArguments` problem above.

Yes, you have to modify your main file, look at the `run_tf_ner.py` example to see what has changed.

> * The current strategy implementation is running out of resources on models we know it can train on (gpt2-medium 1024 tokens). This is speculation, but it might be because `strategy.experimental_distribute_dataset(ds`) isn't being used in `get_train_tfdataset` anymore.

`strategy.experimental_distribute_dataset(ds)` are now respectively in the `train` and `prediction_loop` methods.

> * `tb_writer` was also causing problems, so I had to comment that out too
> * `self.model.ckpt_manager.save()` is throwing `File system scheme '[local]' not implemented`

You have to give the log/input/output directories as a GCS path and not a local path.

> * Special tokens are sometimes added to the model, especially when using this for dialogue-style generation. Maybe add a parameter `tokenizer_length = None` to the class on `__init__` and then replace with this:
> 
> ```
>         if not self.model:
>             with self.strategy.scope():
>                 self.model = self.model_class.from_pretrained(
>                     self.args.model_name_or_path,
>                     # from_pt=bool(".bin" in self.args.model_name_or_path),
>                     config=self.config,
>                     cache_dir=self.argscache_dir,
>                 )
>                 if self.tokenizer_length:
>                     self.model.resize_token_embeddings(self.tokenizer_length)
> ```

This is a temporary solution, the model will soon be created into a `model_init` closure that one has to provide to the Trainer as argument.

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.