Using `torch.compile` disables eval loss when using `Trainer`
Author: RoniGurvich
Pull request: no
Labels: name: solved

### System Info

python version 3.10.6
transformers version 4.30.2
linux machine
running in a poetry env

### Who can help?

@sgugger

### Information

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

### Tasks

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

### Reproduction

The following snippet is a minimal example of training a BERT model for masked language modeling (adopted from the examples). 
When calling it with a compiled model the `eval_loss` log is missing. (initially observed in Tensorboard)

```python
import tempfile

import torch
from torch.utils.data import Dataset
from transformers import (
    BertConfig,
    BertForMaskedLM,
    BertTokenizer,
    DataCollatorForLanguageModeling,
    Trainer,
    TrainingArguments,
)


def _setup_model(max_tokens, pad_token_id, vocab_size, device):
    bert_config = BertConfig(
        vocab_size=vocab_size,
        max_position_embeddings=max_tokens,
        pad_token_id=pad_token_id,
    )
    model = BertForMaskedLM(config=bert_config).to(device)

    return model


class MockDataset(Dataset):
    def __init__(self, sample_size, dataset_size, vocab_size, device):
        self.sample_size = sample_size
        self.n_samples = dataset_size
        self.vocab_size = vocab_size
        self.device = device

    def __len__(self):
        return self.n_samples

    def __getitem__(self, item):
        return torch.randint(0, self.vocab_size, (self.sample_size,)).to(self.device)


def run_training(compile_model: bool):
    batch_size = 2
    max_tokens = 128
    dataset_size = 6
    mlm_proba = 0.15
    device = "cpu"

    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
    vocab_size = tokenizer.vocab_size

    model = _setup_model(max_tokens, tokenizer.pad_token_id, vocab_size, device)

    if compile_model:
        model = torch.compile(model)

    dataset = MockDataset(
        sample_size=max_tokens,
        dataset_size=dataset_size,
        vocab_size=vocab_size,
        device=device,
    )

    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer, mlm=True, mlm_probability=mlm_proba
    )

    with tempfile.TemporaryDirectory() as temp_dir:
        trainer_args = TrainingArguments(
            output_dir=temp_dir,
            do_train=True,
            do_eval=True,
            num_train_epochs=1,
            evaluation_strategy="steps",
            eval_steps=1,
            logging_steps=1,
            per_device_train_batch_size=batch_size,
            per_device_eval_batch_size=batch_size,
        )
        trainer = Trainer(
            model=model,
            args=trainer_args,
            train_dataset=dataset,
            eval_dataset=dataset,
            data_collator=data_collator,
        )
        trainer.train()


if __name__ == "__main__":
    for c in [False, True]:
        run_training(compile_model=c)
```

Outputs: 
```
{'loss': 10.3661, 'learning_rate': 3.3333333333333335e-05, 'epoch': 0.33}                                                                 
{'eval_loss': 10.538872718811035, 'eval_runtime': 0.0573, 'eval_samples_per_second': 104.724, 'eval_steps_per_second': 52.362, 'epoch': 0.33}                                                                                                                                       
{'loss': 10.5957, 'learning_rate': 1.6666666666666667e-05, 'epoch': 0.67}                                                                 
{'eval_loss': 10.470108985900879, 'eval_runtime': 0.0562, 'eval_samples_per_second': 106.714, 'eval_steps_per_second': 53.357, 'epoch': 0.67}                                                                                                                                       
{'loss': 10.4629, 'learning_rate': 0.0, 'epoch': 1.0}                                                                                     
{'eval_loss': 10.48947811126709, 'eval_runtime': 0.0583, 'eval_samples_per_second': 102.897, 'eval_steps_per_second': 51.449, 'epoch': 1.0}                                                                                                                                         
{'train_runtime': 1.5504, 'train_samples_per_second': 3.87, 'train_steps_per_second': 1.935, 'train_loss': 10.474879264831543, 'epoch': 1.0}                                                                                                                                        
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00,  1.94it/s]

{'loss': 10.4365, 'learning_rate': 3.3333333333333335e-05, 'epoch': 0.33}                                                                 
{'eval_runtime': 9.5648, 'eval_samples_per_second': 0.627, 'eval_steps_per_second': 0.314, 'epoch': 0.33}                                 
{'loss': 10.3686, 'learning_rate': 1.6666666666666667e-05, 'epoch': 0.67}                                                                 
{'eval_runtime': 0.0502, 'eval_samples_per_second': 119.503, 'eval_steps_per_second': 59.751, 'epoch': 0.67}                              
{'loss': 10.2967, 'learning_rate': 0.0, 'epoch': 1.0}                                                                                     
{'eval_runtime': 0.0484, 'eval_samples_per_second': 123.961, 'eval_steps_per_second': 61.981, 'epoch': 1.0}                               
{'train_runtime': 37.8631, 'train_samples_per_second': 0.158, 'train_steps_per_second': 0.079, 'train_loss': 10.367273966471354, 'epoch': 1.0}                                                                                                                                      
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:37<00:00, 12.62s/it]
```

### Expected behavior

Expected identical outputs from both calls to `run_training`, specifically I expected `eval_loss` to be logged in the second run as well.

Comments:
Could you try again on the latest release please?

Sure, I get the same output with `4.31.0`

cc @muellerzr 

@RoniGurvich can you give me some more information about libraries being used?

What version of `accelerate` and `torch` specifically.

I ask as running your code gives me an error when using the latest `main` on accelerate and PyTorch 2.0.1:
```
Exception: Please convert all Tensors to FakeTensors first or instantiate FakeTensorMode with 'allow_non_fake_inputs'. Found in aten.embedding.default(*(tensor([[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
        [-0.0095, -0.0085,  0.0011,  ..., -0.0015, -0.0191,  0.0116],
        [ 0.0036,  0.0067,  0.0181,  ...,  0.0131, -0.0042,  0.0615],
        ...,
        [-0.0080, -0.0559, -0.0439,  ...,  0.0162,  0.0055, -0.0214],
        [ 0.0052, -0.0169, -0.0228,  ...,  0.0079,  0.0002,  0.0317],
        [-0.0319, -0.0168, -0.0156,  ..., -0.0213, -0.0202,  0.0351]],
       device='cuda:0', grad_fn=<BroadcastBackward>), FakeTensor(FakeTensor(..., device='meta', size=(2, 128), dtype=torch.int64), cuda:0), 0), **{}) 

While executing %self_embeddings_word_embeddings : [#users=1] = call_module[target=self_embeddings_word_embeddings](args = (%input_ids,), kwargs = {})
Original traceback:
  File "/home/zach_mueller_huggingface_co/transformers/src/transformers/models/bert/modeling_bert.py", line 232, in forward
    inputs_embeds = self.word_embeddings(input_ids)
 |   File "/home/zach_mueller_huggingface_co/transformers/src/transformers/models/bert/modeling_bert.py", line 1015, in <graph break in forward>
    embedding_output = self.embeddings(
```

```
accelerate                0.20.3
torch                     2.0.1+cu118
```

Is there anything else I can do to help with this?

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/main/CONTRIBUTING.md) are likely to be ignored.

@RoniGurvich sorry for the delay, I recommend using the latest `transformers` and `accelerate` versions for this, as I was able to report the `eval_loss` just fine:

```
{'loss': 10.5259, 'learning_rate': 3.3333333333333335e-05, 'epoch': 0.33}                                                                                             
{'eval_loss': 10.493117332458496, 'eval_runtime': 0.1002, 'eval_samples_per_second': 59.863, 'eval_steps_per_second': 29.931, 'epoch': 0.33}                          
{'loss': 10.5693, 'learning_rate': 1.6666666666666667e-05, 'epoch': 0.67}                                                                                             
{'eval_loss': 10.469869613647461, 'eval_runtime': 0.0749, 'eval_samples_per_second': 80.146, 'eval_steps_per_second': 40.073, 'epoch': 0.67}                          
{'loss': 10.5174, 'learning_rate': 0.0, 'epoch': 1.0}                                                                                                                 
{'eval_loss': 10.559216499328613, 'eval_runtime': 0.0741, 'eval_samples_per_second': 80.94, 'eval_steps_per_second': 40.47, 'epoch': 1.0}                             
{'train_runtime': 1.5694, 'train_samples_per_second': 3.823, 'train_steps_per_second': 1.912, 'train_loss': 10.537548383076986, 'epoch': 1.0}                         
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00,  1.91it/s]
  0%|                                                                                                                                           | 0/3 [00:00<?, ?it/s][2023-10-11 15:45:26,168] torch._inductor.utils: [WARNING] using triton random, expect difference from eager
{'loss': 10.4365, 'learning_rate': 3.3333333333333335e-05, 'epoch': 0.33}                                                                                             
{'eval_runtime': 8.5381, 'eval_samples_per_second': 0.703, 'eval_steps_per_second': 0.351, 'epoch': 0.33}                                                             
{'loss': 10.5766, 'learning_rate': 1.6666666666666667e-05, 'epoch': 0.67}                                                                                             
{'eval_runtime': 0.0631, 'eval_samples_per_second': 95.152, 'eval_steps_per_second': 47.576, 'epoch': 0.67}                                                           
{'loss': 10.4973, 'learning_rate': 0.0, 'epoch': 1.0}                                                                                                                 
{'eval_runtime': 0.0613, 'eval_samples_per_second': 97.835, 'eval_steps_per_second': 48.917, 'epoch': 1.0}                                                            
{'train_runtime': 66.7434, 'train_samples_per_second': 0.09, 'train_steps_per_second': 0.045, 'train_loss': 10.50345547993978, 'epoch': 1.0}                          
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [01:06<00:00, 22.25s/it]
```

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/main/CONTRIBUTING.md) are likely to be ignored.

I encountered the same problem with
python 3.12
transformers==4.44.1
torch==2.4.0
accelerate==0.33.0

How did you solve it?

Instead of compiling the model myself I'm passing `torch_compile=True` in `TrainingArguments` to tell `huggingface Trainer` to compile it for me.

[Documentation](https://huggingface.co/docs/transformers/v4.44.2/en/main_classes/trainer#transformers.TrainingArguments.torch_compile)