Bug with finetuning Gemma 2 models
Author: sahandrez
Pull request: no
Labels: name: bug

### System Info

Hi, 

I am getting a loss of 0 and NaN outputs when finetuning Gemma 2 models with the `Trainer` class. Strangely, the model behaves normally during a single forward pass outside of the `Trainer`, so I suspect the problem lies within the `Trainer` itself. The 0 loss and NaN outputs appear immediately as training begins, even before the model's weights are updated. Note that as instructed [here](https://huggingface.co/google/gemma-2-2b-it#running-the-model-on-a-single--multi-gpu), I am using `bfloat16`. Please see below for the reproduction code which is adapted from the [Text Classification Tutorial](https://huggingface.co/docs/transformers/en/tasks/sequence_classification). 

Thanks in advance!

### Who can help?

@ArthurZucker @muellerz @SunMarc 

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

I am using Python 3.10.11 and CUDA 12.1.1 and the following versions: 
```
transformers             4.44.2
torch                    2.4.0
```

Here is a minimal code to reproduce it:
```python
import torch
from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import DataCollatorWithPadding
from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer


def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True)


device = torch.device("cuda")
id2label = {0: "NEGATIVE", 1: "POSITIVE"}
label2id = {"NEGATIVE": 0, "POSITIVE": 1}
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")

imdb = load_dataset("imdb", split="train")
tokenized_imdb = imdb.map(preprocess_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)


model = AutoModelForSequenceClassification.from_pretrained(
    "google/gemma-2-2b-it",
    num_labels=2,
    id2label=id2label,
    label2id=label2id,
    torch_dtype=torch.bfloat16,
    device_map=device,
)

# Sample forward pass
output = model(
    input_ids=torch.tensor(tokenized_imdb["input_ids"][0]).unsqueeze(0).to(device),
    attention_mask=torch.tensor(tokenized_imdb["attention_mask"][0]).unsqueeze(0).to(device),
    labels=torch.tensor(tokenized_imdb["label"][0]).unsqueeze(0).to(device),
    return_dict=True,
)
print(f"Logits: {output['logits']} Loss: {output['loss']}")

training_args = TrainingArguments(
    output_dir="logs/gemma_2_test",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_checkpointing=True,
    num_train_epochs=2,
    weight_decay=0.01,
    bf16=True,
    push_to_hub=False,
    logging_first_step=True,
    logging_steps=10,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_imdb,
    tokenizer=tokenizer,
    data_collator=data_collator,
)

trainer.train()
```

### Expected behavior

The forward pass outside of the `Trainer` class works fine, but in the `Trainer` class the model outputs NaN and the loss is 0.

Comments:
Hey! Can you make sure the model uses either the `eager` or the `flash_attention_2` modules? You should have a warning if it is not! 

Thank you, that fixed it! But as far as I can tell the only warning I was getting was this: 
```
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.4
 [WARNING]  using untested triton version (3.0.0), only 1.0.0 is known to be compatible
```

In main we did merge this: #32865