Hyperparameter search function not working with Trainer and mlflow
Author: laveenaBachani
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.9.2
- Platform:
- Python version:
- PyTorch version (GPU?): 1.9.0
- Tensorflow version (GPU?):
- Using GPU in script?:Yes
- Using distributed or parallel set-up in script?:Yes

### Who can help

@sgugger  @noise-field @LysandreJik 
Models: bert

my own task or dataset: (give details below) : I am using my own dataset

## To reproduce

Steps to reproduce the behavior:

1.  Create Bert model for AutoModelForSequenceClassification and pre process with a dataset 
2. Use pytorch GPU distributed fashion
3. Use trainer.hyperparameter_search backend as optuna
3. Install the mlflow
4.  Use report_to = "none" in the training argument to stop the callback  


The Hyper parameter training stops in middle, there is a communication break between the master and worker nodes.
If I do not give report_to = "none", I am using mlflow in my own script. It given me error that a session with mlflow uuid already running. If I stop my mlflow session then the hyperparameter search happens successfully. 

I do not want to integrated mlflow with the Trainer and want to use mlflow for logging my params. But using the report_to = "none", breaks my code by breaking the communication with master and worker. 

This is my code:

`    def tune(self, train_df, val_df, test_df):
        logging.info("         tuning started")


        model_full_path = self.model_base_path + "/" + self.model_name
        logging.info("         model path["+model_full_path+"]")
        
        tokenizer = AutoTokenizer.from_pretrained(model_full_path, do_lower_case=True)
        
        train_dataset, val_dataset, test_dataset = self.pre_process(tokenizer, train_df, val_df, test_df)

        training_args = TrainingArguments(
            output_dir=self.train_out_dir,          # output directory
            logging_dir=self.train_log_dir,            # directory for storing logs
            num_train_epochs=self.train_param_nb_epochs,              # total # of training epochs
            per_device_train_batch_size=self.train_param_per_device_train_batch_size,  # batch size per device during training
            per_device_eval_batch_size=self.train_param_per_device_eval_batch_size,   # batch size for evaluation
            warmup_steps=self.train_param_warmup_steps,                # number of warmup steps for learning rate scheduler
            weight_decay=self.train_param_weight_decay,               # strength of weight decay
            learning_rate = self.train_param_learning_rate, # args.learning_rate - default is 5e-5, our notebook had 2e-5
            adam_epsilon = self.train_param_adam_epsilon,
            report_to = "none"   #no callbacks to mlflow
        )
        
        def model_init():
            print("-----------------")
            print(model_full_path)
            print(self.nb_labels)
            print("-----------------")

            return AutoModelForSequenceClassification.from_pretrained(model_full_path,
                num_labels = self.nb_labels, # The number of output labels--2 for binary classification.
                output_attentions = False, # Whether the model returns attentions weights.
                output_hidden_states = False, # Whether the model returns all hidden-states.
            )

        from ray.tune.examples.pbt_transformers import utils

        trainer = Trainer(
            model_init=model_init,                 # the instantiated ðŸ¤— Transformers model to be trained
            args=training_args,          # training arguments, defined above
            train_dataset=train_dataset, # training dataset
            eval_dataset=val_dataset,     # evaluation dataset
            compute_metrics=utils.build_compute_metrics_fn('rte')
        )

        def hp_space(trial):
            if self.backend == 'ray':
                from ray import tune
                return {
                    "learning_rate": tune.choice(self.tune_param_learning_rate),
                    "num_train_epochs": tune.choice(self.tune_param_nb_epochs),
                    "per_device_train_batch_size": tune.choice(self.tune_param_per_device_train_batch_size),
                    "per_device_eval_batch_size": tune.choice(self.tune_param_per_device_eval_batch_size),
                    "warmup_steps": tune.choice(self.tune_param_warmup_steps),
                    "weight_decay": tune.choice(self.tune_param_weight_decay),
                    "adam_epsilon": tune.choice(self.tune_param_adam_epsilon)
                }
            elif self.backend == 'optuna':
                return {
                    "learning_rate": trial.suggest_float("learning_rate", self.tune_param_learning_rate[0], self.tune_param_learning_rate[1]),
                    "num_train_epochs": trial.suggest_int("num_train_epochs", self.tune_param_nb_epochs[0], self.tune_param_nb_epochs[1]),
                    "per_device_train_batch_size": trial.suggest_discrete_uniform("per_device_train_batch_size", self.tune_param_per_device_train_batch_size[0], self.tune_param_per_device_train_batch_size[1],self.tune_param_per_device_train_batch_size[2]),
                    "per_device_eval_batch_size": trial.suggest_discrete_uniform("per_device_eval_batch_size", self.tune_param_per_device_eval_batch_size[0], self.tune_param_per_device_eval_batch_size[1], self.tune_param_per_device_eval_batch_size[2]),
                    "warmup_steps": trial.suggest_int("warmup_steps", self.tune_param_warmup_steps[0], self.tune_param_warmup_steps[1]),
                    "weight_decay": trial.suggest_float("weight_decay", self.tune_param_weight_decay[0], self.tune_param_weight_decay[1]),
                    "adam_epsilon": trial.suggest_float("adam_epsilon", self.tune_param_adam_epsilon[0], self.tune_param_adam_epsilon[1])
                }
        best_run = trainer.hyperparameter_search(
            direction="maximize",
            hp_space = hp_space,
            backend=self.backend,
            n_trials = self.tune_param_n_trials,
        )

        for n, v in best_run.hyperparameters.items():
            setattr(trainer.args, n, v)
            mlflow.log_param("best_"+n, v)

        trainer.train()

        # log the validation metrics with test data
        val_eval = trainer.evaluate(eval_dataset = val_dataset)
        self.log_metrics(val_eval, 'val')
        self.log_artifacts(val_df, tokenizer, trainer, "validation_set")

        train_eval = trainer.evaluate(eval_dataset = train_dataset)
        self.log_metrics(train_eval, 'train')
        self.log_artifacts(train_df, tokenizer, trainer, "train_set")

        test_eval = trainer.evaluate(eval_dataset = test_dataset)
        self.log_metrics(test_eval, 'test')
        self.log_artifacts(test_df, tokenizer, trainer, "test_set")
        now = dt.datetime.now()
        model_out_dir = self.model_out_dir+'_'+str(now.year)+'_'+str(now.month)+'_'+str(now.day)+'_'+str(now.hour)+'_'+str(now.minute)      
        trainer.save_model(output_dir = model_out_dir)      
        
        return trainer`


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