"total_flos" showing much bigger number than expected
Author: minh-h-ng
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.14.1
- Platform: Ubuntu 20.04.1
- Python version: 3.8.5
- PyTorch version (GPU?): 1.8.0 (True)
- Tensorflow version (GPU?): n/a
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: DDP

### Who can help

Models:
- TAPAS, LayoutLM, LayoutLMv2, LUKE, ViT, BEiT, DEiT, DETR, CANINE: @NielsRogge

If the model isn't in the list, ping @LysandreJik who will redirect you to the correct contributor.

Library:
- Trainer: @sgugger

Documentation: @sgugger

## Information

Model I am using (Bert, XLNet ...): ViT

The problem arises when using:
* [o ] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [o ] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. download transformers with git clone and pip install
2. run example script in transformers/examples/pytorch/image-classification/run_image_classification.py
3. disable pretrained weights loading to enable start pretraining of ViT-L/16
4. pretraining ViT-L/16 on 1 A100 GPU for 2 images (1 epoch only). In the below result, dividing "total_flos" by 2 would result in 273,934,171,766,784 floating point operations, or around 274 teraFLOPs. This is much bigger than the reported 122.9 GFLOPs in   the "Scaling Vision Transformer" paper (https://arxiv.org/pdf/2106.04560.pdf).
```
{
  "best_metric": null,
  "best_model_checkpoint": null,
  "epoch": 1.0,
  "global_step": 2,
  "is_hyper_param_search": false,
  "is_local_process_zero": true,
  "is_world_process_zero": true,
  "log_history": [
    {
      "epoch": 1.0,
      "step": 2,
      "total_flos": 547868343533568.0,
      "train_loss": 0.06623733788728714,
      "train_runtime": 2.6909,
      "train_samples_per_second": 0.743,
      "train_steps_per_second": 0.743
    }
  ],
  "max_steps": 2,
  "num_train_epochs": 1,
  "total_flos": 547868343533568.0,
  "trial_name": null,
  "trial_params": null
}
```

## Expected behavior

"total_flos" divided by number of train images would match the reported model's GFLOPs in "Scaling Vision Transformer" paper.

Comments:
cc @TevenLeScao since he implemented the feature.

Hey, sorry for the incorrect number - this function was originally written for language models, and might not transfer to vision models, for example. Do you need this number or is it just for checking?

Hello, I was trying to benchmark my A100 server, so having this FLOS number would be helpful for me to calculate the % of peak GPU FLOPS attained by ViT models.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

~~Hi @TevenLeScao, could you confirm whether the `total_flos` computation is indeed accurate for decoder-type LMs?~~

For instance, training a freshly initialized 106M parameter Llama model on 52M tokens results in total_flos=2.7e16, whereas C= 6 * N * D = 3.3e16 FLOPS, which is quite a substantial difference for the purposes of deriving scaling laws etc

**Update**: the discrepancy arises because currently the calculation does not take into account the embedding parameters following [Kaplan et al](https://arxiv.org/pdf/2001.08361.pdf). More recent works, most notably Chinchilla, do take into account the embedding layer in the FLOP computation.