Reset loss to zero on logging in Trainer to avoid bfloat16 issues
Author: bminixhofer
Pull request: yes

# What does this PR do?

I have had a weird issue with the logged loss going to zero after some training steps training with `bfloat16` on v3 TPUs:

![W B Chart 11_16_2020, 12_00_44 PM](https://user-images.githubusercontent.com/13353204/99245133-63d74f80-2803-11eb-8cc0-ce326b2a82b0.png)

while it works correctly using 32bit precision:

![W B Chart 11_16_2020, 12_01_54 PM](https://user-images.githubusercontent.com/13353204/99245229-89645900-2803-11eb-988c-324e4e6d6267.png)

After some investigation I found that the `tr_loss` variable in `Trainer.train` seems to overflow after reaching 1024 (?). 
I did not track this down more closely because it is easily fixed by making `tr_loss` a regular Python float instead of a tensor. It doesn't actually need to be a tensor as it is only ever accessed by `.item()`.

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to the it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors which may be interested in your PR.

@sgugger 

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

 albert, bert, XLM: @LysandreJik
 GPT2: @LysandreJik, @patrickvonplaten
 tokenizers: @mfuntowicz
 Trainer: @sgugger
 Benchmarks: @patrickvonplaten
 Model Cards: @julien-c
 examples/distillation: @VictorSanh
 nlp datasets: [different repo](https://github.com/huggingface/nlp)
 rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
 Text Generation: @patrickvonplaten, @TevenLeScao
 Blenderbot, Bart, Marian, Pegasus: @patrickvonplaten
 T5: @patrickvonplaten
 Rag: @patrickvonplaten, @lhoestq
 EncoderDecoder: @patrickvonplaten
 Longformer, Reformer: @patrickvonplaten
 TransfoXL, XLNet: @TevenLeScao, @patrickvonplaten
 examples/seq2seq: @patil-suraj
 examples/bert-loses-patience: @JetRunner
 tensorflow: @jplu
 examples/token-classification: @stefan-it
 documentation: @sgugger
 FSTM: @stas00
 -->


Comments:
No, we don't want to do `loss.item()` at each step since it slows down a lot the training on TPUs. We can cast this scalar to Float if it overflows. Not super familiar with `bfloat16` on TPUs so we may have overlooked something there (on mixed precision for GPUs, the loss is always a float).

Oh ok, I hadn't considered that. I'll take a closer look then to check what exactly causes the overflow.

And just curious, do you have any metrics regarding slowdown from `loss.item()` on TPUs? I'm currently using the code from this PR and see good TPU utilization and training time.

Alright, this might be a problem with XLA? I'm not at all familiar with how TPUs and XLA work internally but here is a minimal example of the problem:

```bfloat_demo.py```
```py
import torch

import torch_xla
import torch_xla.core.xla_model as xm


device = xm.xla_device()

loss_tensor = torch.tensor(0.0).to(device)
loss_float = 0.0

print(f"loss_tensor is on device {loss_tensor.device} with dtype {loss_tensor.dtype}")

to_add = torch.tensor(10.0).to(device)

for i in range(10):
    for _ in range(100):
        loss_tensor += to_add
        loss_float += to_add.item()

    print(loss_tensor, loss_float)
```

Running this regularly:

```
(torch-xla-1.7) bminixhofer@gerpt2:~$ python bfloat_demo.py
0.0 is on device xla:1 with dtype torch.float32
tensor(1000., device='xla:1') 1000.0
tensor(2000., device='xla:1') 2000.0
tensor(3000., device='xla:1') 3000.0
tensor(4000., device='xla:1') 4000.0
tensor(5000., device='xla:1') 5000.0
tensor(6000., device='xla:1') 6000.0
tensor(7000., device='xla:1') 7000.0
tensor(8000., device='xla:1') 8000.0
tensor(9000., device='xla:1') 9000.0
tensor(10000., device='xla:1') 10000.0
```

and with `bfloat16`:

```
(torch-xla-1.7) bminixhofer@gerpt2:~$ XLA_USE_BF16=1 python bfloat_demo.py 
2020-11-16 14:52:12.488382: I    1663 torch_xla/csrc/tensor_util.cpp:28] Using BF16 data type for floating point values
0.0 is on device xla:1 with dtype torch.float32
tensor(904., device='xla:1') 1000.0
tensor(1704., device='xla:1') 2000.0
tensor(2960., device='xla:1') 3000.0
tensor(4096., device='xla:1') 4000.0
tensor(4096., device='xla:1') 5000.0
tensor(4096., device='xla:1') 6000.0
tensor(4096., device='xla:1') 7000.0
tensor(4096., device='xla:1') 8000.0
tensor(4096., device='xla:1') 9000.0
tensor(4096., device='xla:1') 10000.0
```

but notably the issue doesn't seem to be the magnitude at all but rather how often a value is added:

(setting `to_add = torch.tensor(0.1).to(device)`)
```
(torch-xla-1.7) bminixhofer@gerpt2:~$ XLA_USE_BF16=1 python bfloat_demo.py 
2020-11-16 14:57:57.844438: I    1860 torch_xla/csrc/tensor_util.cpp:28] Using BF16 data type for floating point values
loss_tensor is on device xla:1 with dtype torch.float32
tensor(10.0625, device='xla:1') 10.009765625
tensor(22.5000, device='xla:1') 20.01953125
tensor(32., device='xla:1') 30.029296875
tensor(32., device='xla:1') 40.0390625
tensor(32., device='xla:1') 50.048828125
tensor(32., device='xla:1') 60.05859375
tensor(32., device='xla:1') 70.068359375
tensor(32., device='xla:1') 80.078125
tensor(32., device='xla:1') 90.087890625
tensor(32., device='xla:1') 100.09765625
```

So the dtype does not seem to be the problem, the problem seems to be something along the lines of not enough operations being tracked by XLA, but as I said I really don't know the internals at all so I don't want to go off on speculation here :)
Any ideas how to proceed?

And if you do:
```python
import torch

import torch_xla
import torch_xla.core.xla_model as xm


device = xm.xla_device()

loss_tensor = torch.tensor(0.0).to(device)
loss_float = 0.0

print(f"loss_tensor is on device {loss_tensor.device} with dtype {loss_tensor.dtype}")

to_add = torch.tensor(10.0).to(device)

for i in range(10):
    for _ in range(100):
        loss_tensor += to_add.float()
        loss_float += to_add.item()

    print(loss_tensor, loss_float)
```
does this solve the issue?

Trying to avoid the `.item` as it triggers a synchronization of TPUs :-)

No, still the same :(

`bfloat_demo.py`
```python
import torch

import torch_xla
import torch_xla.core.xla_model as xm


device = xm.xla_device()

loss_tensor = torch.tensor(0.0).to(device)
loss_float = 0.0

print(f"loss_tensor is on device {loss_tensor.device} with dtype {loss_tensor.dtype}")

to_add = torch.tensor(10.0).to(device)

for i in range(10):
    for _ in range(100):
        loss_tensor += to_add.float()
        loss_float += to_add.item()

    print(loss_tensor, loss_float)

```

```
(torch-xla-1.7) bminixhofer@gerpt2:~$ XLA_USE_BF16=1 python bfloat_demo.py 
2020-11-16 16:00:35.131065: I    1197 torch_xla/csrc/tensor_util.cpp:28] Using BF16 data type for floating point values
loss_tensor is on device xla:1 with dtype torch.float32
tensor(904., device='xla:1') 1000.0
tensor(1704., device='xla:1') 2000.0
tensor(2960., device='xla:1') 3000.0
tensor(4096., device='xla:1') 4000.0
tensor(4096., device='xla:1') 5000.0
tensor(4096., device='xla:1') 6000.0
tensor(4096., device='xla:1') 7000.0
tensor(4096., device='xla:1') 8000.0
tensor(4096., device='xla:1') 9000.0
tensor(4096., device='xla:1') 10000.0
```

Ok investigated a bit more by asking some TPU experts and there is no way around this. So the proper fix will be to reset the loss to 0 each time we log it (instead of summing everything from the beginning). I can work on that later this week after the changes needed for v4, or you can work on it if it interests you.

Ok thanks for the quick followup!

> So the proper fix will be to reset the loss to 0 each time we log it (instead of summing everything from the beginning)

Wouldn't that still have the same issue if `logging_steps` is sufficiently large?

I think something like this could be a solution:

```python
import torch

import torch_xla
import torch_xla.core.xla_model as xm


device = xm.xla_device()

loss = 0.0

loss_agg_steps = 52
loss_agg = torch.tensor(0.0).to(device)
zero = torch.tensor(0.0).to(device)

to_add = torch.tensor(10.0).to(device)

for i in range(10):
    for j in range(100):
        loss_agg += to_add

        if (j + 1) % loss_agg_steps == 0:
            loss += loss_agg.item()
            loss_agg.copy_(zero)

    loss += loss_agg.item()
    loss_agg.copy_(zero)

    print(loss)
```

updating the loss tensor for `n` steps and syncing with a Python float loss afterwards (and always syncing before logging and after the last batch in an epoch). `n = 52` is the highest that worked in this demo but maybe a more informed decision could be taken about that value.

The `agg_step` would be the `logging_step` we have as training argument, the user can then tune it to their need.

Ok, I thought it would be good to have a second parameter as the logging step is tied to other things as well but in that case I can give implementing it the way you described earlier a shot.

The latest commit should do the trick. Just not sure if `tr_loss -= tr_loss` is the best way to reset `tr_loss` to zero.

I think you also need to do something for the final reported training loss.

Oh right, thanks. I introduced a variable `_total_loss_scalar` which handles that (comparable to `loss` in the demo code).