Fix TF Trainer loss calculation
<!-- This line specifies which issue to close after the pull request is merged. -->
Fixes #6968

## Description
Issue #6968 is about the incorrect loss calculation due to dividing the per example losses by the number of sentences, rather than by the number of tokens (not ignored, i.e. label != -100) - if the task is a token level task.

## Implementation
Before a whole batch being distributed to replicas, we compute the number of instances in that batch. Depending on the task types (sentence level or token level), the word `instance` means different things:
- sentence level task: it means examples
- token level task: it means the tokens with label != -100

This information (number of instances) is injected into global batches. While each replica receives a small batch, it use this information to correctly computing the scaled losses.

If no information is provided in the dataset, the default behavior is to use the number of examples in a global batch.
This way, the code change is minimal.

## Test code

    import os
    import random
    import shutil
    
    
    shutil.rmtree("./tmp/", ignore_errors=True)
    os.mkdir("./tmp/")
    
    nb_sentences = 70
    words = ["i", "like", "dogs", "but", "you", "prefer", "cats"]
    with open("./tmp/train.txt", "w", encoding="UTF-8") as fp:
        for i in range(nb_sentences):
            if i == 0:
                for word in words:
                    fp.write(f"{word} O\n")
            else:
                fp.write(f" \n")
    
            fp.write("\n")
    
    os.system("cp ./tmp/train.txt ./tmp/dev.txt")
    
    labels = ["O", "B-MISC", "I-MISC", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]
    with open("./tmp/test.txt", "w", encoding="UTF-8") as fp:
        for i in range(nb_sentences):
            for word in words:
                label = random.choice(labels)
                fp.write(f"{word} {label}\n")
    
            fp.write("\n")
    
    command = (
        "python run_tf_ner.py "
        + "--model_name_or_path distilbert-base-uncased "
        + "--data_dir ./tmp/ --seed 2020 --output_dir ./tmp/ "
        + "--overwrite_output_dir --logging_steps 1 "
        + "--do_train --do_eval --do_predict "
        + "--num_train_epochs 1 "
        + f"--per_device_train_batch_size {nb_sentences} "
        + f"--per_device_eval_batch_size {nb_sentences} --labels '' "
        + "--max_seq_length 16"
    )
    
    print(command)
    
    os.system(command)

Testing it against master, you will see the loss values is smaller (~0.3) than testing against this PR code, you will (1.0 ~ 2.0), because on master, the denominator is `70` (but 69 of them having only ignored tokens) while on this PR, the denominator is `7` (the number of tokens).
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/6998?src=pr&el=h1) Report
> Merging [#6998](https://codecov.io/gh/huggingface/transformers/pull/6998?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/0054a48cdd64e7309184a64b399ab2c58d75d4e5?el=desc) will **increase** coverage by `0.19%`.
> The diff coverage is `20.00%`.

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

```diff
@@            Coverage Diff             @@
##           master    #6998      +/-   ##
==========================================
+ Coverage   80.53%   80.72%   +0.19%     
==========================================
  Files         168      168              
  Lines       32179    32197      +18     
==========================================
+ Hits        25915    25991      +76     
+ Misses       6264     6206      -58     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/6998?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/trainer\_tf.py](https://codecov.io/gh/huggingface/transformers/pull/6998/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyX3RmLnB5) | `12.46% <20.00%> (+0.24%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_flaubert.py](https://codecov.io/gh/huggingface/transformers/pull/6998/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9mbGF1YmVydC5weQ==) | `24.53% <0.00%> (-63.81%)` | :arrow_down: |
| [src/transformers/modeling\_marian.py](https://codecov.io/gh/huggingface/transformers/pull/6998/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19tYXJpYW4ucHk=) | `60.00% <0.00%> (-30.00%)` | :arrow_down: |
| [src/transformers/modeling\_openai.py](https://codecov.io/gh/huggingface/transformers/pull/6998/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19vcGVuYWkucHk=) | `72.25% <0.00%> (-10.00%)` | :arrow_down: |
| [src/transformers/activations.py](https://codecov.io/gh/huggingface/transformers/pull/6998/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9hY3RpdmF0aW9ucy5weQ==) | `85.00% <0.00%> (-5.00%)` | :arrow_down: |
| [src/transformers/modeling\_bart.py](https://codecov.io/gh/huggingface/transformers/pull/6998/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19iYXJ0LnB5) | `93.93% <0.00%> (-0.51%)` | :arrow_down: |
| [src/transformers/generation\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6998/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9nZW5lcmF0aW9uX3V0aWxzLnB5) | `96.92% <0.00%> (-0.28%)` | :arrow_down: |
| [src/transformers/tokenization\_utils\_base.py](https://codecov.io/gh/huggingface/transformers/pull/6998/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdXRpbHNfYmFzZS5weQ==) | `93.77% <0.00%> (-0.28%)` | :arrow_down: |
| [src/transformers/file\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6998/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9maWxlX3V0aWxzLnB5) | `82.41% <0.00%> (-0.26%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_xlm.py](https://codecov.io/gh/huggingface/transformers/pull/6998/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl94bG0ucHk=) | `93.26% <0.00%> (+4.84%)` | :arrow_up: |
| ... and [2 more](https://codecov.io/gh/huggingface/transformers/pull/6998/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/6998?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/6998?src=pr&el=footer). Last update [0054a48...b5c5bdc](https://codecov.io/gh/huggingface/transformers/pull/6998?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).
<|||||>> 
> 
> Thanks! It still needs some refactoring but in overall it looks ok. Can you also try the examples scripts for sequence classification and multiple choices to see if there is a drop in performance or not and be sure to do not bring any inconsistency.
> 
> Optionaly if you have the possibility to run it, also the examples scripts, on a multi-gpu environment and check the same thing it would be appreciated, otherwise, no worries I will do it before merging :)

I will try to (and learn to) run the other scripts later today, but I can only run on a single-gpu environment.<|||||>No problem I can do it in multiple gpu env.<|||||>The trainer will still be task-agnostic, the goal is just to add a new parameter to the training_step function (or possibly a class field) to handle the value that will be used to compute the scaled loss instead. It should work for all the tasks.

This computation cannot be done in the loss, because the loss computation is done over a per replica batch size and not over the global batch size.

I'm not in favor of having different trainers. I don't mind having few differences between the two trainers as long as the external usage is the same.<|||||>I don't understand your last comment. Having users subclass Trainer for specific behavior is what is indicated in the documentation. We cannot offer everything any user can think of in the training loop so this the way of customizing one. There is a `Seq2SeqTrainer` in preparation on the PyTorch side for instance, for code that is relevant to this only.

Opening the door to have task-specific components in the main tf_trainer file will make it unreadable in a few months, when every user will have added their own, and then users that rely on a custom Trainers won't use that class anymore, because they won't understand it.

Tagging @julien-c for his advice.<|||||>@sgugger , if fact, I also doubt that the pytorch trainer, while working with token level tasks, have some inaccurate loss computation - when we have distributed training and/or gradient accumulation. For example, [in DistilBertForTokenClassification](https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_distilbert.py#L820), we use

    loss_fct = CrossEntropyLoss()

and 
   
     loss = loss_fct(active_logits, active_labels)

[CrossEntropyLoss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#crossentropyloss) has default reduction `mean`, so basically we compute the averaged loss over the tokens (label !=-100) on each single batch. Then we accumulate it `gradient_accumulation_steps`, and finally average again by dividing `gradient_accumulation_steps`, see [here](https://github.com/huggingface/transformers/blob/master/src/transformers/trainer.py#L1041).

However, this is no longer the same as the `per example losses over the global batch divided the number of tokens on that global batch`. In practice, it may not harm the training, but theoretically, it is not what exactly the gradient accumulation mean to be.

But this should be discussed in another thread, not here.<|||||>> pening the door to have task-specific components in the main tf_trainer file will make it unreadable in a few months, when every user will have added their own, and then users that rely on a custom Trainers won't use that class anymore, because they won't understand it.

The current draft code requires a task to compute the number of `instances` in each `example`. For token level tasks, it will be the number of tokens with labels != -100. Fore sentence level, in general, it will just be the number of sentences. The tf trainer will just use this information for further calculation, but it doesn't need to know what the tasks are.<|||||>@jplu , this is not a finalized version. However, I want to have some feedback from you about if this approach is OK. Thanks.<|||||>For the sake of simplicity, I suspect this would be best left as a user-space extension/custom loss (via a subclass of TFTrainer for instance)

If needed we can add/improve extension points in TFTrainer to make it easier.

What do you think?<|||||>@julien-c 

My thoughts are:

- Hugging Face's ` transformers` has `run_ner_tf.py` which is assumed to give the correct training/evaluation losses
- (maybe I misunderstand the roles of scripts in examples dir?)
- If they are meant to be correct, and let's say the right way for token level tasks is to count the number of tokens (not ignored), then the current script won't give the correct results. In this case, the majority of users who use it won't know they are supposed to subclass TFTrainer class
- Even they are aware of the necessity of counting tokens rather than examples, in TF, it is not easy to make it right - if it has to work correctly in a distributed strategy and/or gradient accumulation. It is not just about modifying the loss calculation - the token counting has to be in a global batch, not a batch already distributed to a single replica. Then in a single replica, the per example losses (on that small batch) has to be divided by the number of tokens computed in the big batch before being distributed.
- If the scripts in examples dir have only the purpose of demonstration of the library's usage - leaving users to customize trainers is fine, and adding/improving extension points seems good to me (although I don't know what it looks like for now). Maybe it is also good to have some warning and a brief tutorial to let users know how to do things.

Please let me know the team's decision about this issue (how to continue or if to close it). Thanks.<|||||>@chiapas Now it looks great! I really like it. Did you try it in context of single and multi-replica?

@julien-c @sgugger I think there is a misunderstanding here, this PR is not to add a new feature or any refactoring, this is a **bugfix** in the loss computation, that means that anybody that are currently using the current version of the trainer get a wrong loss value for token classifications class. And as @chiapas said the way it is computed in the PyTorch trainer might need a fix as well. There are only few changes and do not impact the readability for TF users.<|||||>> 
> 
> @chiapas Now it looks great! I really like it. Did you try it in context of single and multi-replica?

@jplu Thanks. I haven't tried testing yet. I preferred to have your feedback about this new way of fix before finalizing it and testing. By the way, for multi-replica, I can only run on Kaggle or Google colab. I will let you know once the test is done.

<|||||>This is ok no worries, with TPU is fine as well 👍 <|||||>@chiapas As far as I can say this PR should also fix the issue #6969 right? As we compute the total number of example at every step.<|||||>> 
> 
> @chiapas As far as I can say this PR should also fix the issue #6969 right? As we compute the total number of example at every step.

Yes. However, after I opened up that issue and work on this PR, I found that we have

        ds = (
            self.train_dataset.repeat()
            .shuffle(self.num_train_examples, seed=self.args.seed)
            .batch(self.total_train_batch_size, drop_remainder=self.args.dataloader_drop_last)
            .prefetch(tf.data.experimental.AUTOTUNE)
        )

Since the dataset is repeated first, it has no ending, and the drop_remainder has no actual effect (other than set the batch shpae from None to a fixed number along batch dimension), so issue #6969 will never occur. However, in this case, `args.dataloader_drop_last` is somehow confusing.<|||||>Indeed, having the `repeat` has the advantage to avoid the potential last partial batch in each epoch, so users don't need to think about scaling the gradients based on the actual batch size and makes the `dataloader_drop_last` useless.<|||||>> 
> 
> Indeed, having the `repeat` has the advantage to avoid the potential last partial batch in each epoch, so users don't need to think about scaling the gradients based on the actual batch size and makes the `dataloader_drop_last` useless.

Yes, I am ok with this. BTW, `dataloader_drop_last` might still have an effect - if `True`, the batch dimension will be fixed in the compiled graph. If it is set to `False`, even we repeat the dataset first, the batch dimension will still be `None` .  In this case, while working with TPU with the way we do gradient accumulation int trainer_tf.py, we will get an error message that complains TPU can't handle slice or shape - I don't remember the exact message, but I can reproduce one quickly.<|||||>The default value of `drop_reminder` is False, which result in an unknown batch size because the last batch may not be full, this is exactly why `drop_reminder` on TPU has to be set to `True` if no repeat is applied otherwise we can let it to False.

Small example:

```python
>>> dataset = tf.data.Dataset.range(100)
>>> dataset.batch(4)
<BatchDataset shapes: (None,), types: tf.int64>
>>> dataset = tf.data.Dataset.range(100)
>>> dataset.batch(4, drop_remainder=True)
<BatchDataset shapes: (4,), types: tf.int64><|||||>> 
> 
> The default value of `drop_reminder` is False, which result in an unknown batch size because the last batch may not be full, this is exactly why `drop_reminder` on TPU has to be set to `True` if no repeat is applied otherwise we can let it to False.

    otherwise we can let it to False

What I am saying is: We can't let it to `False` even repeat is applied in `trainer_tf.py`. In general, if `repeat` is used, we don't have to drop. But due to the gradient accumulation implementation, if TPU is used and, if we set `drop_remainde=False`, even the `repeat` is applied, we will still get

	<PrefetchDataset shapes: ((None, 512, 512, 3), (None,)), types: (tf.float32, tf.int32)>


	NotFoundError: 3 root error(s) found.
	  (0) Not found: {{function_node __inference_train_step_1_epoch_192920}} No proto found for key <<NO PROGRAM AS COMPILATION FAILED>>
		 [[{{node TPUVariableReshard/reshard/_16819633198340046116/_31}}]]
	  (1) Not found: {{function_node __inference_train_step_1_epoch_192920}} No proto found for key <<NO PROGRAM AS COMPILATION FAILED>>
		 [[{{node TPUVariableReshard/reshard/_17949385379616849075/_19}}]]
	  (2) Unimplemented: {{function_node __inference_train_step_1_epoch_192920}} Compilation failure: Dynamic input dimension to reshape that is both splitted and combined is not supported: output: f32[0,512,512,3], input: f32[<=0,512,512,3], input_dim: 0
		 [[{{node strided_slice_2}}]]
		 [[while/body/_1/while]]
		TPU compilation failed
		 [[tpu_compile_succeeded_assert/_17625000101377989734/_5]]
	0 successful operations.
	6 derived errors ignored.

You can test it if you want on

https://www.kaggle.com/yihdarshieh/tpu-gradient-accumulation?scriptVersionId=41324202

by changing one line in 

    def get_training_dataset(batch_size):<|||||>Hummm nice catch, I haven't tested this case with gradient accumulation, thanks!<|||||>@chiapas for me it looks ok, do you want to add anything else? If no can you switch the PR to open, in order to be able for us to merge it.<|||||>> 
> 
> @chiapas for me it looks ok, do you want to add anything else? If no can you switch the PR to open, in order to be able for us to merge it.

@jplu , I haven't done any test yet - I pushed the code immediately after writting it (in order to have your feedback), so from my side, I feel more comfortable to check a few things before merged to master. Unless you have done some testings and are eager to merge, maybe wait a bit please? I think tomorrow at some point would be ready. <|||||>I have tested it on single and 4 GPU training with the example script for NER and was ok, but take the time you need the more we can test the better. I will be happy to know how it works on TPU as well, even through a colab.<|||||>@jplu I had to fix bugs, and now I have a working version - I checked some intermediate values to make sure the calculation of number of `instances` is correct and sent to replica(s).

Due to the bugs I found - I think there might be a chance that your test yesterday didn't use my code - probably forgot to uninstall the original `transformers` and reinstall `transformers` that is based on my version? I have this doubt because those bugs would throw errors, and the training wouldn't be successful. Sorry about this. If possible, could you test it again with the latest version on multiple GPU env.? Thanks.

Otherwise, I tested with CPU and 1 GPU on colab. It works fine. You can check here

    https://colab.research.google.com/drive/148whpTObbF53qU_ec0bVkyVWUsLlwppn?usp=sharing

However, using TPU with tf 2.2 or 2.3, I had different errors, for which I think irrelevant to this PR code  See below for the error messages. We can probably open a bug report and fix it later.

Also, the CI is not green because of some problems from the master branch. From my side, the code is ready to be merged (if you can test on multiple GPU again would be better). Thanks.

        Traceback (most recent call last):
      File "run_tf_ner.py", line 299, in <module>
        main()
      File "run_tf_ner.py", line 128, in main
        training_args.n_replicas,
      File "/content/transformers/examples/token-classification/transformers/src/transformers/file_utils.py", line 926, in wrapper
        return func(*args, **kwargs)
      File "/content/transformers/examples/token-classification/transformers/src/transformers/training_args_tf.py", line 161, in n_replicas
        return self._setup_strategy.num_replicas_in_sync
      File "/content/transformers/examples/token-classification/transformers/src/transformers/file_utils.py", line 904, in __get__
        cached = self.fget(obj)
      File "/content/transformers/examples/token-classification/transformers/src/transformers/file_utils.py", line 926, in wrapper
        return func(*args, **kwargs)
      File "/content/transformers/examples/token-classification/transformers/src/transformers/training_args_tf.py", line 132, in _setup_strategy
        tf.tpu.experimental.initialize_tpu_system(tpu)
      File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/tpu/tpu_strategy_util.py", line 103, in initialize_tpu_system
        serialized_topology = output.numpy()
      File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py", line 961, in numpy
        maybe_arr = self._numpy()  # pylint: disable=protected-access
      File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py", line 929, in _numpy
        six.raise_from(core._status_to_exception(e.code, e.message), None)
      File "<string>", line 3, in raise_from
    tensorflow.python.framework.errors_impl.InvalidArgumentError: NodeDef expected inputs 'string' do not match 0 inputs specified; Op<name=_Send; signature=tensor:T -> ; attr=T:type; attr=tensor_name:string; attr=send_device:string ..........

and TPU with tf 2.3 gives different error

        Traceback (most recent call last):
      File "run_tf_ner.py", line 299, in <module>
        main()
      File "run_tf_ner.py", line 230, in main
        trainer.train()
      File "/content/transformers/src/transformers/trainer_tf.py", line 474, in train
        train_ds = self.get_train_tfdataset()
      File "/content/transformers/src/transformers/trainer_tf.py", line 137, in get_train_tfdataset
        self.num_train_examples = tf.data.experimental.cardinality(self.train_dataset).numpy()
      File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py", line 1063, in numpy
        maybe_arr = self._numpy()  # pylint: disable=protected-access
      File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py", line 1031, in _numpy
        six.raise_from(core._status_to_exception(e.code, e.message), None)  # pylint: disable=protected-access
      File "<string>", line 3, in raise_from
    tensorflow.python.framework.errors_impl.UnimplementedError: File system scheme '[local]' not implemented (file: 'runs/Sep11_12-45-28_1d2f5ee8ee35')
        Encountered when executing an operation using EagerExecutor. This error cancels all future operations and poisons their output tensors.
    Error in atexit._run_exitfuncs:
    Traceback (most recent call last):
      File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/tpu_strategy.py", line 540, in async_wait
        context.async_wait()
      File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/context.py", line 2319, in async_wait
        context().sync_executors()
      File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/context.py", line 658, in sync_executors
        pywrap_tfe.TFE_ContextSyncExecutors(self._context_handle)
    tensorflow.python.framework.errors_impl.UnimplementedError: File system scheme '[local]' not implemented (file: 'runs/Sep11_12-45-28_1d2f5ee8ee35')
        Encountered when executing an operation using EagerExecutor. This error cancels all future operations and poisons their output tensors.
    2020-09-11 12:46:28.864320: W ./tensorflow/core/distributed_runtime/eager/destroy_tensor_handle_node.h:57] Ignoring an error encountered when deleting remote tensors handles: Invalid argument: Unable to find the relevant tensor remote_handle: Op ID: 9416, Output num: 0
    Additional GRPC error information from remote target /job:worker/replica:0/task:0:
    :{"created":"@1599828388.860894650","description":"Error received from peer ipv4:10.42.193.26:8470","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Unable to find the relevant tensor remote_handle: Op ID: 9416, Output num: 0","grpc_status":3}<|||||>>Due to the bugs I found - I think there might be a chance that your test yesterday didn't use my code - probably forgot to uninstall the original transformers and reinstall transformers that is based on my version? I have this doubt because those bugs would throw errors, and the training wouldn't be successful. Sorry about this. If possible, could you test it again with the latest version on multiple GPU env.? Thanks.

Ah! Might be possible I forgot to install your version. I will re-test it, to be sure.

> However, using TPU with tf 2.2 or 2.3, I had different errors, for which I think irrelevant to this PR code See below for the error messages. We can probably open a bug report and fix it later.

The second error you get means that you cannot load data from localhost, the files have to be hosted on a GCS to make it works.

> Also, the CI is not green because of some problems from the master branch. 

Can you try to rebase on the current master and see if the CI error still occurs?

> From my side, the code is ready to be merged (if you can test on multiple GPU again would be better). Thanks.

I will test that ASAP today and will let you know.<|||||>I have been able to run a NER task over 4 GPUs. Without gradient accumulation:

```
***** Running training *****
  Num examples = 24000
  Num Epochs = 3
  Instantaneous batch size per device = 32
  Total train batch size (w. parallel, distributed & accumulation) = 128
  Gradient Accumulation steps = 1
  Steps per epoch = 188
  Total optimization steps = 564
{'loss': 18.6387, 'learning_rate': 4.9113474e-05, 'epoch': 0.05851063829787234, 'step': 10}
{'loss': 14.222743, 'learning_rate': 4.822695e-05, 'epoch': 0.11170212765957446, 'step': 20}
{'loss': 12.079448, 'learning_rate': 4.734042e-05, 'epoch': 0.16489361702127658, 'step': 30}
{'loss': 10.455857, 'learning_rate': 4.6453897e-05, 'epoch': 0.21808510638297873, 'step': 40}
{'loss': 9.295527, 'learning_rate': 4.5567373e-05, 'epoch': 0.2712765957446808, 'step': 50}
{'loss': 8.374706, 'learning_rate': 4.468085e-05, 'epoch': 0.324468085106383, 'step': 60}
{'loss': 7.674021, 'learning_rate': 4.379432e-05, 'epoch': 0.3776595744680851, 'step': 70}
{'loss': 7.106626, 'learning_rate': 4.29078e-05, 'epoch': 0.4308510638297872, 'step': 80}
{'loss': 6.6170573, 'learning_rate': 4.202128e-05, 'epoch': 0.48404255319148937, 'step': 90}
{'loss': 6.226298, 'learning_rate': 4.113475e-05, 'epoch': 0.5372340425531915, 'step': 100}
{'loss': 5.859087, 'learning_rate': 4.0248226e-05, 'epoch': 0.5904255319148937, 'step': 110}
{'loss': 5.560567, 'learning_rate': 3.93617e-05, 'epoch': 0.6436170212765957, 'step': 120}
{'loss': 5.2810636, 'learning_rate': 3.8475173e-05, 'epoch': 0.6968085106382979, 'step': 130}
{'loss': 5.040142, 'learning_rate': 3.758865e-05, 'epoch': 0.75, 'step': 140}
{'loss': 4.830164, 'learning_rate': 3.6702124e-05, 'epoch': 0.8031914893617021, 'step': 150}
{'loss': 4.6353145, 'learning_rate': 3.5815603e-05, 'epoch': 0.8563829787234043, 'step': 160}
{'loss': 4.446635, 'learning_rate': 3.492908e-05, 'epoch': 0.9095744680851063, 'step': 170}
{'loss': 4.300565, 'learning_rate': 3.4042554e-05, 'epoch': 0.9627659574468085, 'step': 180}
2020-09-11 23:08:59.187632: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:172] Filling up shuffle buffer (this may take a while): 23119 of 24000
2020-09-11 23:08:59.564647: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:221] Shuffle buffer filled.
{'loss': 1.8294506, 'learning_rate': 3.3156026e-05, 'epoch': 1.0106382978723405, 'step': 190}
{'loss': 1.576384, 'learning_rate': 3.22695e-05, 'epoch': 1.0638297872340425, 'step': 200}
{'loss': 1.4572238, 'learning_rate': 3.1382977e-05, 'epoch': 1.1170212765957448, 'step': 210}
{'loss': 1.4322406, 'learning_rate': 3.0496454e-05, 'epoch': 1.1702127659574468, 'step': 220}
{'loss': 1.3880556, 'learning_rate': 2.960993e-05, 'epoch': 1.2234042553191489, 'step': 230}
{'loss': 1.3613675, 'learning_rate': 2.8723403e-05, 'epoch': 1.2765957446808511, 'step': 240}
{'loss': 1.3514798, 'learning_rate': 2.7836879e-05, 'epoch': 1.3297872340425532, 'step': 250}
{'loss': 1.3266419, 'learning_rate': 2.6950353e-05, 'epoch': 1.3829787234042552, 'step': 260}
{'loss': 1.3012911, 'learning_rate': 2.606383e-05, 'epoch': 1.4361702127659575, 'step': 270}
{'loss': 1.2993147, 'learning_rate': 2.5177305e-05, 'epoch': 1.4893617021276595, 'step': 280}
{'loss': 1.2913059, 'learning_rate': 2.4290777e-05, 'epoch': 1.5425531914893615, 'step': 290}
{'loss': 1.2822802, 'learning_rate': 2.3404255e-05, 'epoch': 1.5957446808510638, 'step': 300}
{'loss': 1.2839314, 'learning_rate': 2.2517732e-05, 'epoch': 1.648936170212766, 'step': 310}
{'loss': 1.2641081, 'learning_rate': 2.1631204e-05, 'epoch': 1.702127659574468, 'step': 320}
{'loss': 1.2524884, 'learning_rate': 2.0744681e-05, 'epoch': 1.7553191489361701, 'step': 330}
{'loss': 1.2450953, 'learning_rate': 1.9858155e-05, 'epoch': 1.8085106382978724, 'step': 340}
{'loss': 1.2448001, 'learning_rate': 1.897163e-05, 'epoch': 1.8617021276595744, 'step': 350}
{'loss': 1.2407304, 'learning_rate': 1.8085108e-05, 'epoch': 1.9148936170212765, 'step': 360}
{'loss': 1.2282307, 'learning_rate': 1.719858e-05, 'epoch': 1.9680851063829787, 'step': 370}
2020-09-11 23:14:15.677977: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:172] Filling up shuffle buffer (this may take a while): 23228 of 24000
2020-09-11 23:14:16.010139: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:221] Shuffle buffer filled.
{'loss': 0.85539556, 'learning_rate': 1.6312057e-05, 'epoch': 2.021276595744681, 'step': 380}
{'loss': 0.94433624, 'learning_rate': 1.542553e-05, 'epoch': 2.074468085106383, 'step': 390}
{'loss': 0.93659353, 'learning_rate': 1.4539006e-05, 'epoch': 2.127659574468085, 'step': 400}
{'loss': 0.8978215, 'learning_rate': 1.365248e-05, 'epoch': 2.1808510638297873, 'step': 410}
{'loss': 0.9126247, 'learning_rate': 1.27659605e-05, 'epoch': 2.2340425531914896, 'step': 420}
{'loss': 0.9438525, 'learning_rate': 1.1879432e-05, 'epoch': 2.2872340425531914, 'step': 430}
{'loss': 0.9655043, 'learning_rate': 1.0992908e-05, 'epoch': 2.3404255319148937, 'step': 440}
{'loss': 0.9694119, 'learning_rate': 1.0106382e-05, 'epoch': 2.393617021276596, 'step': 450}
{'loss': 0.95613927, 'learning_rate': 9.219858e-06, 'epoch': 2.4468085106382977, 'step': 460}
{'loss': 0.9483009, 'learning_rate': 8.333333e-06, 'epoch': 2.5, 'step': 470}
{'loss': 0.93453395, 'learning_rate': 7.4468076e-06, 'epoch': 2.5531914893617023, 'step': 480}
{'loss': 0.92573655, 'learning_rate': 6.5602835e-06, 'epoch': 2.6063829787234045, 'step': 490}
{'loss': 0.9156919, 'learning_rate': 5.67376e-06, 'epoch': 2.6595744680851063, 'step': 500}
{'loss': 0.9098517, 'learning_rate': 4.787233e-06, 'epoch': 2.7127659574468086, 'step': 510}
{'loss': 0.9110744, 'learning_rate': 3.9007095e-06, 'epoch': 2.7659574468085104, 'step': 520}
{'loss': 0.89820355, 'learning_rate': 3.014183e-06, 'epoch': 2.8191489361702127, 'step': 530}
{'loss': 0.8982555, 'learning_rate': 2.1276592e-06, 'epoch': 2.872340425531915, 'step': 540}
{'loss': 0.8997822, 'learning_rate': 1.2411356e-06, 'epoch': 2.925531914893617, 'step': 550}
{'loss': 0.89611685, 'learning_rate': 3.5460886e-07, 'epoch': 2.978723404255319, 'step': 560}
Training took: 0:17:32.846350
Saving model in /home/jplu/model
Configuration saved in /home/jplu/model/config.json
Model weights saved in /home/jplu/model/tf_model.h5
***** Running Evaluation *****
  Num examples = 2200
  Batch size = 32
{'eval_loss': 1.4825085626132246, 'eval_precision': 0.8298914945747288, 'eval_recall': 0.8713708195516354, 'eval_f1': 0.8501254930082466, 'epoch': 3.0, 'step': 564}
```

With gradient accumulation:
```
***** Running training *****
  Num examples = 24000
  Num Epochs = 3
  Instantaneous batch size per device = 32
  Total train batch size (w. parallel, distributed & accumulation) = 256
  Gradient Accumulation steps = 2
  Steps per epoch = 94
  Total optimization steps = 282
{'loss': 17.029345, 'learning_rate': 4.822695e-05, 'epoch': 0.11702127659574468, 'step': 10}
{'loss': 13.583568, 'learning_rate': 4.6453897e-05, 'epoch': 0.22340425531914893, 'step': 20}
{'loss': 11.413903, 'learning_rate': 4.468085e-05, 'epoch': 0.32978723404255317, 'step': 30}
{'loss': 9.977048, 'learning_rate': 4.29078e-05, 'epoch': 0.43617021276595747, 'step': 40}
{'loss': 8.904137, 'learning_rate': 4.113475e-05, 'epoch': 0.5425531914893617, 'step': 50}
{'loss': 8.056796, 'learning_rate': 3.93617e-05, 'epoch': 0.648936170212766, 'step': 60}
{'loss': 7.339738, 'learning_rate': 3.758865e-05, 'epoch': 0.7553191489361702, 'step': 70}
{'loss': 6.7678766, 'learning_rate': 3.5815603e-05, 'epoch': 0.8617021276595744, 'step': 80}
{'loss': 6.2809086, 'learning_rate': 3.4042554e-05, 'epoch': 0.9680851063829787, 'step': 90}
2020-09-11 23:36:37.142279: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:172] Filling up shuffle buffer (this may take a while): 23596 of 24000
2020-09-11 23:36:37.311979: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:221] Shuffle buffer filled.
{'loss': 2.2556372, 'learning_rate': 3.22695e-05, 'epoch': 1.0638297872340425, 'step': 100}
{'loss': 2.0573573, 'learning_rate': 3.0496454e-05, 'epoch': 1.1702127659574468, 'step': 110}
{'loss': 1.9544038, 'learning_rate': 2.8723403e-05, 'epoch': 1.2765957446808511, 'step': 120}
{'loss': 1.8848253, 'learning_rate': 2.6950353e-05, 'epoch': 1.3829787234042552, 'step': 130}
{'loss': 1.837054, 'learning_rate': 2.5177305e-05, 'epoch': 1.4893617021276595, 'step': 140}
{'loss': 1.7885295, 'learning_rate': 2.3404255e-05, 'epoch': 1.5957446808510638, 'step': 150}
{'loss': 1.7535, 'learning_rate': 2.1631204e-05, 'epoch': 1.702127659574468, 'step': 160}
{'loss': 1.7068337, 'learning_rate': 1.9858155e-05, 'epoch': 1.8085106382978724, 'step': 170}
{'loss': 1.6874169, 'learning_rate': 1.8085108e-05, 'epoch': 1.9148936170212765, 'step': 180}
2020-09-11 23:41:05.123380: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:172] Filling up shuffle buffer (this may take a while): 23596 of 24000
2020-09-11 23:36:05.332080: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:221] Shuffle buffer filled.
{'loss': 1.189175, 'learning_rate': 1.6312057e-05, 'epoch': 2.021276595744681, 'step': 190}
{'loss': 1.2820004, 'learning_rate': 1.4539006e-05, 'epoch': 2.127659574468085, 'step': 200}
{'loss': 1.2558761, 'learning_rate': 1.27659605e-05, 'epoch': 2.2340425531914896, 'step': 210}
{'loss': 1.2865627, 'learning_rate': 1.0992908e-05, 'epoch': 2.3404255319148937, 'step': 220}
{'loss': 1.2853887, 'learning_rate': 9.219858e-06, 'epoch': 2.4468085106382977, 'step': 230}
{'loss': 1.2650005, 'learning_rate': 7.4468076e-06, 'epoch': 2.5531914893617023, 'step': 240}
{'loss': 1.2414478, 'learning_rate': 5.67376e-06, 'epoch': 2.6595744680851063, 'step': 250}
{'loss': 1.243986, 'learning_rate': 3.9007095e-06, 'epoch': 2.7659574468085104, 'step': 260}
{'loss': 1.2295542, 'learning_rate': 2.1276592e-06, 'epoch': 2.872340425531915, 'step': 270}
{'loss': 1.2296557, 'learning_rate': 3.5460886e-07, 'epoch': 2.978723404255319, 'step': 280}
Training took: 0:19:10.070496
Saving model in /home/jplu/model
Configuration saved in /home/jplu/model/config.json
Model weights saved in /home/jplu/model/tf_model.h5
***** Running Evaluation *****
  Num examples = 2200
  Batch size = 32
{'eval_loss': 1.5663623533387114, 'eval_precision': 0.8107354478912513, 'eval_recall': 0.8548327820654171, 'eval_f1': 0.8322003577817532, 'epoch': 3.0, 'step': 282}
```
Looks ok to me. I have also tested for text classification and question answering without any error.
<|||||>> 
> 
> I have been able to run a NER task over 4 GPUs. Without gradient accumulation:
> 
> ```
> ***** Running training *****
>   Num examples = 24000
>   Num Epochs = 3
>   Instantaneous batch size per device = 32
>   Total train batch size (w. parallel, distributed & accumulation) = 128
>   Gradient Accumulation steps = 1
>   Steps per epoch = 188
>   Total optimization steps = 564
> {'loss': 18.6387, 'learning_rate': 4.9113474e-05, 'epoch': 0.05851063829787234, 'step': 10}
> {'loss': 14.222743, 'learning_rate': 4.822695e-05, 'epoch': 0.11170212765957446, 'step': 20}
> {'loss': 12.079448, 'learning_rate': 4.734042e-05, 'epoch': 0.16489361702127658, 'step': 30}
> {'loss': 10.455857, 'learning_rate': 4.6453897e-05, 'epoch': 0.21808510638297873, 'step': 40}
> {'loss': 9.295527, 'learning_rate': 4.5567373e-05, 'epoch': 0.2712765957446808, 'step': 50}
> {'loss': 8.374706, 'learning_rate': 4.468085e-05, 'epoch': 0.324468085106383, 'step': 60}
> {'loss': 7.674021, 'learning_rate': 4.379432e-05, 'epoch': 0.3776595744680851, 'step': 70}
> {'loss': 7.106626, 'learning_rate': 4.29078e-05, 'epoch': 0.4308510638297872, 'step': 80}
> {'loss': 6.6170573, 'learning_rate': 4.202128e-05, 'epoch': 0.48404255319148937, 'step': 90}
> {'loss': 6.226298, 'learning_rate': 4.113475e-05, 'epoch': 0.5372340425531915, 'step': 100}
> {'loss': 5.859087, 'learning_rate': 4.0248226e-05, 'epoch': 0.5904255319148937, 'step': 110}
> {'loss': 5.560567, 'learning_rate': 3.93617e-05, 'epoch': 0.6436170212765957, 'step': 120}
> {'loss': 5.2810636, 'learning_rate': 3.8475173e-05, 'epoch': 0.6968085106382979, 'step': 130}
> {'loss': 5.040142, 'learning_rate': 3.758865e-05, 'epoch': 0.75, 'step': 140}
> {'loss': 4.830164, 'learning_rate': 3.6702124e-05, 'epoch': 0.8031914893617021, 'step': 150}
> {'loss': 4.6353145, 'learning_rate': 3.5815603e-05, 'epoch': 0.8563829787234043, 'step': 160}
> {'loss': 4.446635, 'learning_rate': 3.492908e-05, 'epoch': 0.9095744680851063, 'step': 170}
> {'loss': 4.300565, 'learning_rate': 3.4042554e-05, 'epoch': 0.9627659574468085, 'step': 180}
> 2020-09-11 23:08:59.187632: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:172] Filling up shuffle buffer (this may take a while): 23119 of 24000
> 2020-09-11 23:08:59.564647: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:221] Shuffle buffer filled.
> {'loss': 1.8294506, 'learning_rate': 3.3156026e-05, 'epoch': 1.0106382978723405, 'step': 190}
> {'loss': 1.576384, 'learning_rate': 3.22695e-05, 'epoch': 1.0638297872340425, 'step': 200}
> {'loss': 1.4572238, 'learning_rate': 3.1382977e-05, 'epoch': 1.1170212765957448, 'step': 210}
> {'loss': 1.4322406, 'learning_rate': 3.0496454e-05, 'epoch': 1.1702127659574468, 'step': 220}
> {'loss': 1.3880556, 'learning_rate': 2.960993e-05, 'epoch': 1.2234042553191489, 'step': 230}
> {'loss': 1.3613675, 'learning_rate': 2.8723403e-05, 'epoch': 1.2765957446808511, 'step': 240}
> {'loss': 1.3514798, 'learning_rate': 2.7836879e-05, 'epoch': 1.3297872340425532, 'step': 250}
> {'loss': 1.3266419, 'learning_rate': 2.6950353e-05, 'epoch': 1.3829787234042552, 'step': 260}
> {'loss': 1.3012911, 'learning_rate': 2.606383e-05, 'epoch': 1.4361702127659575, 'step': 270}
> {'loss': 1.2993147, 'learning_rate': 2.5177305e-05, 'epoch': 1.4893617021276595, 'step': 280}
> {'loss': 1.2913059, 'learning_rate': 2.4290777e-05, 'epoch': 1.5425531914893615, 'step': 290}
> {'loss': 1.2822802, 'learning_rate': 2.3404255e-05, 'epoch': 1.5957446808510638, 'step': 300}
> {'loss': 1.2839314, 'learning_rate': 2.2517732e-05, 'epoch': 1.648936170212766, 'step': 310}
> {'loss': 1.2641081, 'learning_rate': 2.1631204e-05, 'epoch': 1.702127659574468, 'step': 320}
> {'loss': 1.2524884, 'learning_rate': 2.0744681e-05, 'epoch': 1.7553191489361701, 'step': 330}
> {'loss': 1.2450953, 'learning_rate': 1.9858155e-05, 'epoch': 1.8085106382978724, 'step': 340}
> {'loss': 1.2448001, 'learning_rate': 1.897163e-05, 'epoch': 1.8617021276595744, 'step': 350}
> {'loss': 1.2407304, 'learning_rate': 1.8085108e-05, 'epoch': 1.9148936170212765, 'step': 360}
> {'loss': 1.2282307, 'learning_rate': 1.719858e-05, 'epoch': 1.9680851063829787, 'step': 370}
> 2020-09-11 23:14:15.677977: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:172] Filling up shuffle buffer (this may take a while): 23228 of 24000
> 2020-09-11 23:14:16.010139: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:221] Shuffle buffer filled.
> {'loss': 0.85539556, 'learning_rate': 1.6312057e-05, 'epoch': 2.021276595744681, 'step': 380}
> {'loss': 0.94433624, 'learning_rate': 1.542553e-05, 'epoch': 2.074468085106383, 'step': 390}
> {'loss': 0.93659353, 'learning_rate': 1.4539006e-05, 'epoch': 2.127659574468085, 'step': 400}
> {'loss': 0.8978215, 'learning_rate': 1.365248e-05, 'epoch': 2.1808510638297873, 'step': 410}
> {'loss': 0.9126247, 'learning_rate': 1.27659605e-05, 'epoch': 2.2340425531914896, 'step': 420}
> {'loss': 0.9438525, 'learning_rate': 1.1879432e-05, 'epoch': 2.2872340425531914, 'step': 430}
> {'loss': 0.9655043, 'learning_rate': 1.0992908e-05, 'epoch': 2.3404255319148937, 'step': 440}
> {'loss': 0.9694119, 'learning_rate': 1.0106382e-05, 'epoch': 2.393617021276596, 'step': 450}
> {'loss': 0.95613927, 'learning_rate': 9.219858e-06, 'epoch': 2.4468085106382977, 'step': 460}
> {'loss': 0.9483009, 'learning_rate': 8.333333e-06, 'epoch': 2.5, 'step': 470}
> {'loss': 0.93453395, 'learning_rate': 7.4468076e-06, 'epoch': 2.5531914893617023, 'step': 480}
> {'loss': 0.92573655, 'learning_rate': 6.5602835e-06, 'epoch': 2.6063829787234045, 'step': 490}
> {'loss': 0.9156919, 'learning_rate': 5.67376e-06, 'epoch': 2.6595744680851063, 'step': 500}
> {'loss': 0.9098517, 'learning_rate': 4.787233e-06, 'epoch': 2.7127659574468086, 'step': 510}
> {'loss': 0.9110744, 'learning_rate': 3.9007095e-06, 'epoch': 2.7659574468085104, 'step': 520}
> {'loss': 0.89820355, 'learning_rate': 3.014183e-06, 'epoch': 2.8191489361702127, 'step': 530}
> {'loss': 0.8982555, 'learning_rate': 2.1276592e-06, 'epoch': 2.872340425531915, 'step': 540}
> {'loss': 0.8997822, 'learning_rate': 1.2411356e-06, 'epoch': 2.925531914893617, 'step': 550}
> {'loss': 0.89611685, 'learning_rate': 3.5460886e-07, 'epoch': 2.978723404255319, 'step': 560}
> Training took: 0:17:32.846350
> Saving model in /home/jplu/model
> Configuration saved in /home/jplu/model/config.json
> Model weights saved in /home/jplu/model/tf_model.h5
> ***** Running Evaluation *****
>   Num examples = 2200
>   Batch size = 32
> {'eval_loss': 1.4825085626132246, 'eval_precision': 0.8298914945747288, 'eval_recall': 0.8713708195516354, 'eval_f1': 0.8501254930082466, 'epoch': 3.0, 'step': 564}
> ```
> 
> With gradient accumulation:
> 
> ```
> ***** Running training *****
>   Num examples = 24000
>   Num Epochs = 3
>   Instantaneous batch size per device = 32
>   Total train batch size (w. parallel, distributed & accumulation) = 256
>   Gradient Accumulation steps = 2
>   Steps per epoch = 94
>   Total optimization steps = 282
> {'loss': 17.029345, 'learning_rate': 4.822695e-05, 'epoch': 0.11702127659574468, 'step': 10}
> {'loss': 13.583568, 'learning_rate': 4.6453897e-05, 'epoch': 0.22340425531914893, 'step': 20}
> {'loss': 11.413903, 'learning_rate': 4.468085e-05, 'epoch': 0.32978723404255317, 'step': 30}
> {'loss': 9.977048, 'learning_rate': 4.29078e-05, 'epoch': 0.43617021276595747, 'step': 40}
> {'loss': 8.904137, 'learning_rate': 4.113475e-05, 'epoch': 0.5425531914893617, 'step': 50}
> {'loss': 8.056796, 'learning_rate': 3.93617e-05, 'epoch': 0.648936170212766, 'step': 60}
> {'loss': 7.339738, 'learning_rate': 3.758865e-05, 'epoch': 0.7553191489361702, 'step': 70}
> {'loss': 6.7678766, 'learning_rate': 3.5815603e-05, 'epoch': 0.8617021276595744, 'step': 80}
> {'loss': 6.2809086, 'learning_rate': 3.4042554e-05, 'epoch': 0.9680851063829787, 'step': 90}
> 2020-09-11 23:36:37.142279: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:172] Filling up shuffle buffer (this may take a while): 23596 of 24000
> 2020-09-11 23:36:37.311979: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:221] Shuffle buffer filled.
> {'loss': 2.2556372, 'learning_rate': 3.22695e-05, 'epoch': 1.0638297872340425, 'step': 100}
> {'loss': 2.0573573, 'learning_rate': 3.0496454e-05, 'epoch': 1.1702127659574468, 'step': 110}
> {'loss': 1.9544038, 'learning_rate': 2.8723403e-05, 'epoch': 1.2765957446808511, 'step': 120}
> {'loss': 1.8848253, 'learning_rate': 2.6950353e-05, 'epoch': 1.3829787234042552, 'step': 130}
> {'loss': 1.837054, 'learning_rate': 2.5177305e-05, 'epoch': 1.4893617021276595, 'step': 140}
> {'loss': 1.7885295, 'learning_rate': 2.3404255e-05, 'epoch': 1.5957446808510638, 'step': 150}
> {'loss': 1.7535, 'learning_rate': 2.1631204e-05, 'epoch': 1.702127659574468, 'step': 160}
> {'loss': 1.7068337, 'learning_rate': 1.9858155e-05, 'epoch': 1.8085106382978724, 'step': 170}
> {'loss': 1.6874169, 'learning_rate': 1.8085108e-05, 'epoch': 1.9148936170212765, 'step': 180}
> 2020-09-11 23:41:05.123380: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:172] Filling up shuffle buffer (this may take a while): 23596 of 24000
> 2020-09-11 23:36:05.332080: I tensorflow/core/kernels/data/shuffle_dataset_op.cc:221] Shuffle buffer filled.
> {'loss': 1.189175, 'learning_rate': 1.6312057e-05, 'epoch': 2.021276595744681, 'step': 190}
> {'loss': 1.2820004, 'learning_rate': 1.4539006e-05, 'epoch': 2.127659574468085, 'step': 200}
> {'loss': 1.2558761, 'learning_rate': 1.27659605e-05, 'epoch': 2.2340425531914896, 'step': 210}
> {'loss': 1.2865627, 'learning_rate': 1.0992908e-05, 'epoch': 2.3404255319148937, 'step': 220}
> {'loss': 1.2853887, 'learning_rate': 9.219858e-06, 'epoch': 2.4468085106382977, 'step': 230}
> {'loss': 1.2650005, 'learning_rate': 7.4468076e-06, 'epoch': 2.5531914893617023, 'step': 240}
> {'loss': 1.2414478, 'learning_rate': 5.67376e-06, 'epoch': 2.6595744680851063, 'step': 250}
> {'loss': 1.243986, 'learning_rate': 3.9007095e-06, 'epoch': 2.7659574468085104, 'step': 260}
> {'loss': 1.2295542, 'learning_rate': 2.1276592e-06, 'epoch': 2.872340425531915, 'step': 270}
> {'loss': 1.2296557, 'learning_rate': 3.5460886e-07, 'epoch': 2.978723404255319, 'step': 280}
> Training took: 0:19:10.070496
> Saving model in /home/jplu/model
> Configuration saved in /home/jplu/model/config.json
> Model weights saved in /home/jplu/model/tf_model.h5
> ***** Running Evaluation *****
>   Num examples = 2200
>   Batch size = 32
> {'eval_loss': 1.5663623533387114, 'eval_precision': 0.8107354478912513, 'eval_recall': 0.8548327820654171, 'eval_f1': 0.8322003577817532, 'epoch': 3.0, 'step': 282}
> ```
> 
> Looks ok to me. I have also tested for text classification and question answering without any error.

Wow! Thank you @jplu, a test and a reply on Friday night :) 
I might need to get some gpu though if I continue to contribute - can't let you do all such tests all the time.
Great to see it works.<|||||>Ahah no worries it is ok, not everybody can have such setup.

@LysandreJik looks ok to merge.