LongT5ForConditionalGeneration NAN losses with bf16
### System Info

transformers version: 4.23.0.dev0
torch version: 1.12.1
OS: Ubuntu 20
Cuda: 11.6

The problem is that LongT5 is supposed to work with bf16=True, but it doesn't. It is known that fp16 should fail in this, and I have tried it and effectively fails. However, Longt5 is supposed to be trained on bf16, therefore it would be expected that turning bf16 to True would work.
My training arguments look like this:

```python
{
        "evaluation_strategy": "epoch", 
        "num_train_epochs": 4,
        "do_train": True,
        "do_eval": False,
        "eval_steps": 2,
        "logging_strategy":"epoch",
        "save_strategy": "epoch",
        "save_total_limit": 4,
        "seed": 69,
        "bf16": True, 
        "dataloader_num_workers": 32,
        "adam_epsilon": 1e-8,
        "adam_beta1": 0.9,
        "adam_beta2": 0.999,
        "group_by_length": False,
        "gradient_checkpointing": False,
        "lr_scheduler_type": "linear",
        "learning_rate": 1e-4,
        "per_device_train_batch_size": 1,
        "per_device_eval_batch_size": 1,
        "gradient_accumulation_steps": 64,
        "warmup_ratio": 0.08
    }
```



### Who can help?

@patrickvonplaten 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

If you need a script for reproduction please let me know.

### Expected behavior

Longt5 (as I understand from the forum etc) should work with bf16.


Agree that it should work with `bf16` - when you mean it fails, how does it fail? Just bad training results ? Could you define this here? 

Also cc @ArthurZucker here<|||||>Sorry I did not specify this, you're right! What I meant is that losses are nans, therefore the model does not learn (it happens the same as when using fp16 with a bf16-pretrained model).<|||||>Okey cc @ArthurZucker and @stancld here<|||||>Thanks! If there is anything I can do to help please let me know, I need this kind of urgently so I'm very much interested in it working properly :) <|||||>Any update?<|||||>Hey! Sorry I will have a look tomorrow! 🤗<|||||>Okay! :) Thanks!<|||||>Could you give me a training script and the full error stack so that I can work on your issue? 🤗  Sorry for the delay <|||||>Yeah, I'll post it here as soon as I can so that you can reproduce it :) thank you very much!!<|||||>First run:

```bash
pip install transformers datasets rouge_score
```

Then, with this script you can replicate it. It is set with bf16=True.

```python



from transformers import AutoModelForSeq2SeqLM, Seq2SeqTrainer, Seq2SeqTrainingArguments, AutoTokenizer, DataCollatorForSeq2Seq
from datasets import load_dataset, load_metric
import nltk
import numpy as np

nltk.download('punkt')

model_str = "google/long-t5-tglobal-base"

metric = load_metric("rouge")

tokenizer = AutoTokenizer.from_pretrained(model_str)

dataset = load_dataset("IIC/sqac_tests")

max_input_length = 128
max_target_length = 16

def preprocess_function(examples):
    inputs = [doc for doc in examples["context"]]
    model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)

    # Setup the tokenizer for targets
    with tokenizer.as_target_tokenizer():
        labels = tokenizer(examples["title"], max_length=max_target_length, truncation=True)

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

tokenized_dataset = dataset.map(preprocess_function, batched=True)

model = AutoModelForSeq2SeqLM.from_pretrained(model_str)

batch_size = 1
args = Seq2SeqTrainingArguments(
    "fail_test",
    evaluation_strategy = "epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    save_total_limit=3,
    num_train_epochs=20,
    predict_with_generate=True,
    bf16=True, # change to fp16 in no Ampere GPU  available.
)

data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
    # Replace -100 in the labels as we can't decode them.
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    
    # Rouge expects a newline after each sentence
    decoded_preds = ["\n".join(nltk.sent_tokenize(pred.strip())) for pred in decoded_preds]
    decoded_labels = ["\n".join(nltk.sent_tokenize(label.strip())) for label in decoded_labels]
    
    result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
    # Extract a few results
    result = {key: value.mid.fmeasure * 100 for key, value in result.items()}
    
    # Add mean generated length
    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions]
    result["gen_len"] = np.mean(prediction_lens)
    
    return {k: round(v, 4) for k, v in result.items()}

trainer = Seq2SeqTrainer(
    model,
    args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics
)

trainer.train()
```

With this you can observe that both training and validation losses are Nans... Hope it helps to figure out what happens, if I can provide any more help please let me know :) @ArthurZucker <|||||>Did you have the time to try the script? @ArthurZucker @stancld <|||||>Hi, any updates? @ArthurZucker @stancld @patrickvonplaten I did not receive any answer after sending the reproduction script....<|||||>Hey sorry not yet ! <|||||>Okay thanks! Let me know if I can help in some way... :) @ArthurZucker <|||||>Hey! Just tested your script and both losses are not Nan. Since you seem to be using a dev version, where you on main? 
It seems that with most recent versions it works perfectly well. 
```python 
{'eval_loss': 2.1923375129699707, 'eval_rouge1': 26.7045, 'eval_rouge2': 12.0, 'eval_rougeL': 25.4545, 'eval_rougeLsum': 25.8144, 'eval_gen_len': 19.0, 'eval_runtime': 5.0999, 'eval_samples_per_second': 1.961, 'eval_steps_per_second': 1.961, 'epoch': 20.0}                                                                                                          
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [02:37<00:00,  3.10it/s]
                                                                                                                                                                                     
Training completed. Do not forget to share your model on huggingface.co/models =)


{'train_runtime': 157.0429, 'train_samples_per_second': 1.274, 'train_steps_per_second': 1.274, 'train_loss': 3.485249328613281, 'epoch': 20.0}                                      
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [02:37<00:00,  1.27it/s]
```

 I used both `transformers==4.22` and `4.23`. <|||||>Thank you so much for trying it out! Let me try it again with latest version (4.25.0.dev) to check if it is working now also in my machine. I was trying to 4.23.0.dev, so maybe if it doesn't work also in 4.25.0.dev I'll have to turn to 4.22 or 4.23 :) <|||||>Great, with the last version it does work!! Thank you very much for helping me ! @ArthurZucker 