Invalid argument error with TFRoberta on GLUE
Author: h4ste
Pull request: no
Labels: name: wontfix

## 🐛 Bug

I am unable to get TFRoberta working on the GLUE benchmark.

Model I am using (Bert, XLNet....): TFRoberta

Language I am using the model on (English, Chinese....): English

The problem arise when using:
* [ ] the official example scripts: (give details)
* [x] my own modified scripts: (give details)

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: *MRPC*
* [ ] my own task or dataset: (give details)

## To Reproduce

Steps to reproduce the behavior:

1. Copy the following script into a file name `run_glue_task_tf.py`:
```python
import tensorflow as tf
import tensorflow.keras as keras

import tensorflow_datasets.public_api as tfds

import transformers

from absl import app
from absl import flags
from absl import logging

FLAGS = flags.FLAGS
flags.DEFINE_multi_string("tasks", None, "One or more tasks to be used for pretraining")

flags.DEFINE_integer('num_epochs', 3, 'Number of epochs to train')
flags.DEFINE_integer('batch_size', 32, 'Batch size to use for training')
flags.DEFINE_integer('eval_batch_size', 64, 'Batch size to use when evaluating validation/test sets')
flags.DEFINE_boolean('use_xla', False, 'Enable XLA optimization')
flags.DEFINE_boolean('use_amp', False, 'Enable AMP optimization')
flags.DEFINE_integer('max_seq_len', 128, 'Maximum sequence length')
flags.DEFINE_string('model_name', 'bert-base-cased', 'Name of pretrained transformer model to load')


def configure_tf():
    logging.info(('Enabling' if FLAGS.use_xla else 'Disabling') + ' XLA optimization')
    tf.config.optimizer.set_jit(FLAGS.use_xla)
    logging.info(('Enabling' if FLAGS.use_xla else 'Disabling') + ' auto mixed precision (AMP)')
    tf.config.optimizer.set_experimental_options({'auto_mixed_precision': FLAGS.use_amp})


def get_model():
    logging.info('Loading pre-trained TF model from %s', FLAGS.model_name)
    model = transformers.TFAutoModelForSequenceClassification.from_pretrained(FLAGS.model_name)

    opt = keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08)
    if FLAGS.use_amp:
        logging.debug('Enabling loss scaling')
        opt = keras.mixed_precision.experimental.LossScaleOptimizer(opt, 'dynamic')

    loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

    metric = keras.metrics.SparseCategoricalAccuracy('accuracy')

    model.compile(optimizer=opt, loss=loss, metrics=[metric])

    return model


def run_task(task: str, model: keras.Model, tokenizer):
    data, info = tfds.load(task, with_info=True)

    if task.startswith('glue'):
        glue_task = task[len('glue/'):]

        def load_features(split: tfds.core.splits.NamedSplit):
            logging.debug('Converting %s.%s to features', task, split)
            is_xlnet: bool = 'xlnet' in model.name.lower()
            return transformers.glue_convert_examples_to_features(examples=data[split],
                                                                  tokenizer=tokenizer,
                                                                  max_length=FLAGS.max_seq_len,
                                                                  output_mode=transformers.glue_output_modes[glue_task],
                                                                  pad_on_left=is_xlnet,  # Pad on the left for XLNet
                                                                  pad_token=tokenizer.convert_tokens_to_ids(
                                                                      [tokenizer.pad_token])[0],
                                                                  pad_token_segment_id=4 if is_xlnet else 0,
                                                                  task=glue_task)

        train = load_features(tfds.Split.TRAIN)
        valid = load_features(tfds.Split.VALIDATION)
    else:
        raise ValueError('Unsupported task: %s' % task)

    train = train.shuffle(128).batch(FLAGS.batch_size).repeat(FLAGS.num_epochs)
    valid = valid.batch(FLAGS.eval_batch_size).take(FLAGS.eval_batch_size)

    logging.info('Training %s on %s...', model.name, task)
    history = model.fit(x=train, validation_data=valid)

    print(task, 'Performance:')
    print(history)


def main(argv):
    del argv  # Unused.

    logging.debug('Loading tokenizer from %s', FLAGS.model_name)
    tokenizer = transformers.AutoTokenizer.from_pretrained(FLAGS.model_name)
    model = get_model()

    for task in FLAGS.tasks:
        print('-' * 20, task, '-' * 20)
        run_task(task, model, tokenizer)


if __name__ == '__main__':
    app.run(main)
```
2. Run the following command:
```bash
$ python run_glue_task_tf.py --tasks glue/mrpc --model_name roberta-base
```
3. Experience the following error:
```
I1101 14:04:14.134916 46912496418432 run_pretraining.py:76] Training tf_roberta_for_sequence_classification on glue/mrpc...
/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/framework/indexed_slices.py:424: UserWarning: Converting sparse IndexedSlices to a dense Tensor of unknown shape. This may consume a large amount of memory.
  "Converting sparse IndexedSlices to a dense Tensor of unknown shape. "
WARNING:tensorflow:Gradients do not exist for variables ['tf_roberta_for_sequence_classification/roberta/pooler/dense/kernel:0', 'tf_roberta_for_sequence_classification/roberta/pooler/dense/bias:0'] when minimizing the loss.
W1101 14:04:25.996359 46912496418432 optimizer_v2.py:1029] Gradients do not exist for variables ['tf_roberta_for_sequence_classification/roberta/pooler/dense/kernel:0', 'tf_roberta_for_sequence_classification/roberta/pooler/dense/bias:0'] when minimizing the loss.
/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/framework/indexed_slices.py:424: UserWarning: Converting sparse IndexedSlices to a dense Tensor of unknown shape. This may consume a large amount of memory.
  "Converting sparse IndexedSlices to a dense Tensor of unknown shape. "
WARNING:tensorflow:Gradients do not exist for variables ['tf_roberta_for_sequence_classification/roberta/pooler/dense/kernel:0', 'tf_roberta_for_sequence_classification/roberta/pooler/dense/bias:0'] when minimizing the loss.
W1101 14:04:36.157998 46912496418432 optimizer_v2.py:1029] Gradients do not exist for variables ['tf_roberta_for_sequence_classification/roberta/pooler/dense/kernel:0', 'tf_roberta_for_sequence_classification/roberta/pooler/dense/bias:0'] when minimizing the loss.
2019-11-01 14:04:45.689286: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAbort Invalid argument: indices[28,20] = 1 is not in [0, 1)
         [[{{node tf_roberta_for_sequence_classification/roberta/embeddings/token_type_embeddings/embedding_lookup}}]]
2019-11-01 14:04:45.900194: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAbort Invalid argument: indices[28,20] = 1 is not in [0, 1)
         [[{{node tf_roberta_for_sequence_classification/roberta/embeddings/token_type_embeddings/embedding_lookup}}]]
         [[Shape/_8]]
      1/Unknown - 28s 28s/stepTraceback (most recent call last):
  File "run_pretraining.py", line 96, in <module>
    app.run(main)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/absl/app.py", line 299, in run
    _run_main(main, args)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/absl/app.py", line 250, in _run_main
    sys.exit(main(argv))
  File "run_pretraining.py", line 92, in main
    run_task(task, model, tokenizer)
  File "run_pretraining.py", line 77, in run_task
    history = model.fit(x=train, validation_data=valid)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training.py", line 728, in fit
    use_multiprocessing=use_multiprocessing)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 324, in fit
    total_epochs=epochs)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 123, in run_one_epoch
    batch_outs = execution_function(iterator)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2_utils.py", line 86, in execution_function
    distributed_function(input_fn))
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py", line 457, in __call__
    result = self._call(*args, **kwds)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py", line 520, in _call
    return self._stateless_fn(*args, **kwds)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1823, in __call__
    return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1141, in _filtered_call
    self.captured_inputs)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1224, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 511, in call
    ctx=ctx)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/execute.py", line 67, in quick_execute
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError:  indices[28,20] = 1 is not in [0, 1)
         [[node tf_roberta_for_sequence_classification/roberta/embeddings/token_type_embeddings/embedding_lookup (defined at /data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:1751) ]] [Op:__inference_distributed_function_58202]

Function call stack:
distributed_function

2019-11-01 14:04:47.519198: W tensorflow/core/kernels/data/generator_dataset_op.cc:102] Error occurred when finalizing GeneratorDataset iterator: Failed precondition: Python interpreter state is not initialized. The process may be terminated.
         [[{{node PyFunc}}]]
```

## Expected behavior

The script should train on GLUE when using TFRoberta just as it trains when using TFBert.

## Environment

OS: CentOS-7
Python version: 3.6.9
PyTorch version: 1.2.0
PyTorch Transformers version (or branch): 2.1.1 (master from git)
Using GPU ? Yes
Distributed of parallel setup ? No
Any other relevant information:

## Additional context


Comments:
I have encountered this embedding error and I found that wrapping the model in a tf.distribute strategy fixes it (extremely odd...)

You can wrap the model in `tf.distribute.OneDeviceStrategy` and it will work.

Did you ever figure it out? 

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.