Modify the Trainer class to handle simultaneous execution of Ray Tune and Weights & Biases

[ruanchaves]

What does this PR do?

The proper way to integrate Ray Tune and Weights & Biases is to pass a wandb parameter to tune.run.

However, this parameter is handled as a dictionary inside the config argument, and there is no distinction between wandb parameters and standard model optimization parameters. The following code comes from their docs:

from ray.tune.logger import DEFAULT_LOGGERS
from ray.tune.integration.wandb import WandbLogger
tune.run(
    train_fn,
    config={
        # define search space here
        "parameter_1": tune.choice([1, 2, 3]),
        "parameter_2": tune.choice([4, 5, 6]),
        # wandb configuration
        "wandb": {
            "project": "Optimization_Project",
            "api_key_file": "/path/to/file",
            "log_config": True
        }
    },
    loggers=DEFAULT_LOGGERS + (WandbLogger, ))

This is not a problem for Ray Tune. However, it is a problem for the transformers integration because it treats wandb as a model parameter, and therefore configuring wandb in this way will raise an error message claiming that wandb is not a training argument.

The following code will raise such an error:

    # Initialize our Trainer
    trainer = Trainer(
        model_init=model_init,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        compute_metrics=compute_metrics,
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    # Hyperparameter Search

    def hp_space_fn(empty_arg):
        config = {
                    "warmup_steps": tune.choice([50, 100, 500, 1000]),
                    "learning_rate": tune.choice([1.5e-5, 2e-5, 3e-5, 4e-5]),
                    "num_train_epochs": tune.quniform(0.0, 10.0, 0.5),
        }
        wandb_config = {
                "wandb": {
                        "project": os.environ.get(
                            'WANDB_PROJECT',
                            'wandb_project'),
                        "api_key": os.environ.get('API_KEY'),
                        "log_config": True
                        }
        }
        config.update(wandb_config)
        return config

    best_run = trainer.hyperparameter_search(
            direction="maximize",
            backend="ray",
            scheduler=PopulationBasedTraining(
                        time_attr='time_total_s',
                        metric='eval_f1_thr_0',
                        mode='max',
                        perturbation_interval=600.0
                    ),
            hp_space=hp_space_fn,
            loggers=DEFAULT_LOGGERS + (WandbLogger, ),
    )

One way to work around this is to instantiate a subclass based on the Trainer:

    class CustomTrainer(Trainer):

        def __init__(self, *args, **kwargs):
            super(CustomTrainer, self).__init__(*args, **kwargs)

        def _hp_search_setup(self, trial: Any):
            try:
                trial.pop('wandb', None)
            except AttributeError:
                pass
            super(CustomTrainer, self)._hp_search_setup(trial)

However, this looks like a hack because throwing away wandb arguments in model config on _hp_search_setup should be standard Trainer behavior.

That's why I'm submitting a PR that directly modifies the _hp_search_setup of the Trainer class to ignore wandb arguments if Ray is chosen as a backend.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

I'm tagging @richardliaw and @amogkam as they're directly involved in Ray Tune.

[sgugger]

Thanks for fixing! Waiting for @richardliaw or @amogkam approval before merging.

[amogkam]

LGTM! Thanks