Wandb sweeps integraition: custom objective
Author: BramVanroy
Pull request: no
Labels: name: WIP

name: Integrations

### System Info

- `transformers` version: 4.29.2
- Platform: Linux-5.14.0-162.6.1.el9_1.0.1.x86_64-x86_64-with-glibc2.34
- Python version: 3.10.10
- Huggingface_hub version: 0.14.1
- Safetensors version: not installed
- PyTorch version (GPU?): 2.0.1+cu117 (True) -- using torch in my experiments
- Tensorflow version (GPU?): 2.12.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.6.10 (cpu)
- Jax version: 0.4.10
- JaxLib version: 0.4.10


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

I followed [this guide](https://huggingface.co/docs/transformers/hpo_train) to use wandb sweeps with the trainer by modifying the summarization scripts slightly. Below are the hp_space, model_init, and hyperparameter_search commands that I use.

Most notably, the objective is to maximize the sum of sari and rougeLsum.

```python
    def wandb_hp_space(trial):
        return {
            "method": "bayes",
            "metric": {
                "name": "objective",
                "goal": "minimize" if hyperopt_args.hparam_optimize_for_loss else "maximize"
            },
            "parameters": {
                "num_train_epochs": {"min": hyperopt_args.hparam_epoch_min, "max": hyperopt_args.hparam_epoch_max}
            },
            "run_cap": hyperopt_args.hparam_max_trials
        }

    def model_init(trial):
        return AutoModelForSeq2SeqLM.from_pretrained(
            model_args.model_name_or_path,
            from_tf=bool(".ckpt" in model_args.model_name_or_path),
            config=config,
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            use_auth_token=True if model_args.use_auth_token else None,
        )

    def hparam_objective(metrics: Dict[str, float]) -> float:
        metrics = copy.deepcopy(metrics)

        if hyperopt_args.hparam_optimize_for_loss:
            return metrics["eval_loss"]

        return metrics["eval_rougeLsum"] + metrics["eval_sari"]

    best_trial = trainer.hyperparameter_search(
        compute_objective=hparam_objective,
        backend="wandb",
        # I think that this is only used to set the column in the sweep chart but does not mean that we use
        # this metric only for optimization. That is what the hparam_objective is for?
        metric="eval/sari",
        hp_space=wandb_hp_space,
        n_trials=hyperopt_args.hparam_max_trials,
        direction="minimize" if hyperopt_args.hparam_optimize_for_loss else "maximize",
    )
```

The problem is that when I look at the wandb interface at the generated sweep config, it looks like this:

```yaml
method: bayes
metric:
  goal: maximize
  name: eval/sari
parameters:
  num_train_epochs:
    distribution: int_uniform
    max: 30
    min: 2
run_cap: 16
```

So the generated sweep config includes `eval/sari` as the metric name because I passed it in the `hyperparameter_search`. But as you can read in the comment, I thought this was only for the wandb visualization but now I am not so sure. When I leave out the `metric` keyword (as in the example), wandb seems to fallback to `eval/loss`.

```yaml
method: bayes
metric:
  goal: maximize
  name: eval/loss
parameters:
  num_train_epochs:
    distribution: int_uniform
    max: 30
    min: 2
run_cap: 16
```

My worry here is the disconnect between the generated sweep config and the custom objective function. What is wandb optimizing now? My custom objective function (sari+rougeLsum) or `metric` that is passed to `hyperparameter_search` for `direction`?

### Expected behavior

More clarity about the relationship between a wandb sweep configuration and how the trainer uses wandb as a backend, and the importance of `metric` and `direciton` arguments vs. `compute_objective`.

Comments:
We are not maintaining the Wandb integration, so you should ping the Wandb team here :-) 

Hey @BramVanroy thanks for raising this issue. Here's my response from top of my head:

Scenario 1:
>  the objective is to maximize the sum of sari and rougeLsum

Are you logging this sum to wandb? Say you are logging it to wandb with name `eval/sari_rougelsum` then passing `metric: eval/sari_rougelsum` should work.

Scenario 2:
You are logging the two metrics separately. Sweeps doesn't support multi-objective optimization today. Say you are using the grid or random search you can use the parallel coordinate plot generated in the sweeps dashboard to find a set of hyperparameters that can optimize two objectives:


In this proxy example say `eval/val_accuracy` and `eval_loss` are two objectives. I can then use this parallel coordinate plot to find range of hparams that maximise the `val_accuracy` and minimise `eval_loss`.
<img width="1355" alt="image" src="https://github.com/huggingface/transformers/assets/31141479/1b8b62a4-4be0-4540-b331-f08801708e67">

> More clarity about the relationship between a wandb sweep configuration and how the trainer uses wandb as a backend, and the importance of metric and direciton arguments vs. compute_objective.

I have shared this request internally. Thanks.

Hi @ayulockin, thanks for the quick reply!

> Are you logging this sum to wandb? Say you are logging it to wandb with name eval/sari_rougelsum then passing metric: eval/sari_rougelsum should work.

Can you tell me how to log an extra metric in the trainer or directly in this training script (I'm using a modified version of this): https://github.com/huggingface/transformers/blob/main/examples/pytorch/summarization/run_summarization.py

> Sweeps doesn't support multi-objective optimization today.

Does that mean that the custom objective function that I pass to the huggingface trainer is NOT used by `wandb`?

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.

@ayulockin Any update on 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.

Kind reminder :) Bump @ayulockin 

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.

Final attempt @ayulockin 

Hey @BramVanroy, serious apologies for not responding to your pings. Past few weeks were crazy and I couldn't get to the GitHub notifications. Apologies again. :(

I will be working on it today and will let you know in a couple of hours. Thank you for the patience and apologies again. 🙏 

Hey @BramVanroy, will it be possible for you to share a colab notebook (if you have one handy) with your code. I am not able to reproduce your exact configuration.

@ayulockin I do my training on the cluster with full scripts so I can't really share those. But can you tell me how we should typically add a custom metric (like in the case above `metrics["eval_rougeLsum"] + metrics["eval_sari"]`) and then use that as a parameter to maximize in the sweep?

Say this is the `wandb_hp_space`.

```
wandb_hp_space = {
    'method': 'bayes',
    'name': 'sweep',
    'metric': {
        'goal': 'maximize', # notice here
        'name': 'custom_metric' # notice here
        },
    'parameters': {
        'batch_size': {'values': [16, 32, 64]},
        'epochs': {'values': [5, 10, 15]},
        'lr': {'max': 0.1, 'min': 0.0001}
     }
}
```

You will have to log the `custom_metric` somewhere in your training script - `wandb.log({"custom_metric": metrics["eval_rougeLsum"] + metrics["eval_sari"]})`.

Will this help?

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.

Replying so that the issue is not closed.

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.

@ayulockin Sorry from my end, I lost track of this. I will try this now and see what comes out!

Hi @BramVanroy, apologies for the very late reply, we're catching up on the backlog.

Your worry is correct, and it's by design: with the wandb backend the actual hyperparameter selection is driven by wandb's own sweep controller (wandb.agent), which optimizes whatever metric name you pass. The metric= kwarg (default eval/loss) and direction become the sweep's metric.name / metric.goal. Your compute_objective only sets trainer.objective, which Transformers uses solely to pick which run to return as BestRun at the end.
So to make wandb *optimize* your sari + rougeLsum sum, log it during training (e.g. a small callback doing `wandb.log({"eval/custom_objective": ...})`) and pass `metric="eval/custom_objective"` to `hyperparameter_search`, exactly as @ayulockin suggested.

As Sylvain noted early on, this backend is maintained by the wandb team rather than us, so this is expected behavior rather than a Transformers bug; closing on that basis