metric_for_best_model not found
Author: asaluja
Pull request: no

### System Info

Transformers version: 4.31.0
Python version: 3.10
Accelerate version: 0.21.0
TRL version: 0.7.2

I am not sure if this issue should be here or in the `trl` repo, but the stack trace showed an error in `transformers/trainer.py`.

I am using `trl.DPOTrainer`. I define `metric_for_best_model="eval_rewards/accuracies"` (a metric from `trl.DPOTrainer`), but then I get the following error:

```
KeyError: 'eval_rewards/accuracies'
    return inner_training_loop(
  File "/mnt/miniconda3/envs/synlm/lib/python3.10/site-packages/transformers/trainer.py", line 1901, in _inner_training_loop
    self._save_checkpoint(model, trial, metrics=metrics)
  File "/mnt/miniconda3/envs/synlm/lib/python3.10/site-packages/transformers/trainer.py", line 2356, in _save_checkpoint
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/mnt/miniconda3/envs/synlm/lib/python3.10/site-packages/transformers/trainer.py", line 2237, in _maybe_log_save_evaluate
    print(metrics['eval_rewards/accuracies'])
KeyError: 'eval_rewards/accuracies'
    self._save_checkpoint(model, trial, metrics=metrics)
  File "/mnt/miniconda3/envs/synlm/lib/python3.10/site-packages/transformers/trainer.py", line 2356, in _save_checkpoint
    print(metrics['eval_rewards/accuracies'])
KeyError: 'eval_rewards/accuracies'
```

I noticed that the metric I defined is only available on `cuda:0` (I am running on 8 40GB A100s). Is there something special I need to do to get this to work? If this is an issue for the `trl` repo let me know accordingly please. 



### Who can help?

@muellerzr @pacman100 

### 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

```
training_args = TrainingArguments(
        remove_unused_columns=False,
        per_device_train_batch_size=config.BATCH_SIZE_TRAIN,
        gradient_accumulation_steps=config.GRADIENT_ACC_STEPS,
        lr_scheduler_type="linear",
        warmup_steps=150,
        report_to="tensorboard",
        max_steps=args.max_steps,
        num_train_epochs=config.NUM_EPOCHS,
        learning_rate=config.LEARNING_RATE,
        bf16=True,
        logging_steps=config.SAVE_STEPS,
        optim="paged_adamw_8bit",
        max_grad_norm=0.3,
        evaluation_strategy="steps" if val_set_size > 0 else "no",
        save_strategy="steps",
        eval_steps=config.SAVE_STEPS if val_set_size > 0 else None,
        per_device_eval_batch_size=config.BATCH_SIZE_EVAL,
        eval_accumulation_steps=128,
        save_steps=config.SAVE_STEPS,
        output_dir=args.output_dir,
        load_best_model_at_end=True,
        metric_for_best_model="eval_rewards/accuracies", # if I remove this line and the following line it works
        greater_is_better=True,
        ddp_find_unused_parameters=False,
        group_by_length=False,
        dataloader_num_workers=accelerator.num_processes,
        save_total_limit=1,
    )
    lora_config = LoraConfig(
        r=config.LORA_R,
        lora_alpha=config.LORA_ALPHA,
        target_modules=config.LORA_TARGET_MODULES,
        lora_dropout=config.LORA_DROPOUT,
        bias="none",
        task_type="CAUSAL_LM",
    )
    trainer = DPOTrainer(
        model,
        model_ref,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        args=training_args,
        tokenizer=tokenizer,
        beta=config.DPO_BETA,
        max_length=config.EXAMPLE_MAX_LENGTH,
        max_prompt_length=config.EXAMPLE_MAX_LENGTH - synlm_config.BUFFER_TOKENS,
        truncation_mode="keep_start",
        peft_config=lora_config,
    )
    trainer.train()
```

### Expected behavior

Trainer will find best model as per `eval_rewards/accuracies` and save that to disk.

Comments:
cc @pacman100 

@ArthurZucker @pacman100 following up on this, am I doing something wrong here or is there a bug?

cc @muellerzr if you can have a look since @pacman100 is off! 😉 

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.

@muellerz can you please take a look?

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.

Gentle ping @muellerzr @pacman100 

Gently pinging @muellerzr @pacman100, if you can have a look as this has been stale for quite some time. 

Hello @asaluja,

I have gone through the implementation of DPOTrainer at https://github.com/huggingface/trl/blob/main/trl/trainer/dpo_trainer.py. 

it overrides the `compute_loss` and `prediction_step` methods to specifically compute and store DPO related metrics such as `rewards/chosen`, `rewards/accuracies` etc  which is then store in private dict attribute specific to DPOTrainer `_stored_metrics` and it is only used in `log` method to log to tensorboard/wandb ... 

https://github.com/huggingface/trl/blob/f12a1da74baef7aff6656bb48247ac109ce72744/trl/trainer/dpo_trainer.py#L1222-L1236

On the other hand, for you to use `metric_for_best_model`, you need to pass `compute_metrics` function when creating Trainer instance which should compute the metric you are interested in by taking as input `EvalPrediction`. Because you are not doing that, `eval_rewards/accuracies` isn't computed by the parent Trainer's evaluation loop thereby such metric being unavailable leading to key error. To resolve this, you need to pass the `compute_metrics` which mimics what DPOTrainer is doing internally for each prediction step and adapt it. 

This is not an issue with Trainer and it should be handled by the calling script/user.

thanks for clarifying @pacman100 ! It looks like I'll have to re-implement `get_batch_loss_metrics` [here](https://github.com/huggingface/trl/blob/f12a1da74baef7aff6656bb48247ac109ce72744/trl/trainer/dpo_trainer.py#L1000) in `compute_metrics`. I will close this issue now.