Ray pickling issue when running hp search
Author: mvacaporale
Pull request: no

## Environment info

- `transformers` version: 4.5.1
- Platform: Amazon Linux 2
- Python version: 3.7.7
- PyTorch version (GPU?): 1.8.1
- Tensorflow version (GPU?): 2.0.0
- Using GPU in script?:  Yes
- Using distributed or parallel set-up in script?: No

### Who can help
Trainer / Ray
@sgugger, @richardliaw, @amogkam

## Information

I'm using BERT with my own modified training scripts for hp searches on pre-training; however the issue is reproduced on a simple example as shown below.

## To reproduce

The following code snippet is a slight modification from the blog post [here](https://huggingface.co/blog/ray-tune)

```
from datasets import load_dataset, load_metric
from transformers import (AutoModelForSequenceClassification, AutoTokenizer,
                          Trainer, TrainingArguments)

tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
dataset = load_dataset('glue', 'mrpc')
metric = load_metric('glue', 'mrpc')

def encode(examples):
    outputs = tokenizer(
        examples['sentence1'], examples['sentence2'], truncation=True)
    return outputs

encoded_dataset = dataset.map(encode, batched=True)

def model_init():
    return AutoModelForSequenceClassification.from_pretrained(
        'distilbert-base-uncased', return_dict=True)

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = predictions.argmax(axis=-1)
    return metric.compute(predictions=predictions, references=labels)

# Evaluate during training and a bit more often
# than the default to be able to prune bad trials early.
# Disabling tqdm is a matter of preference.
training_args = TrainingArguments(
    "test", eval_steps=500, disable_tqdm=True)
trainer = Trainer(
    args=training_args,
    tokenizer=tokenizer,
    train_dataset=encoded_dataset["train"],
    eval_dataset=encoded_dataset["validation"],
    model_init=model_init,
    compute_metrics=compute_metrics,
)

# Defaut objective is the sum of all metrics
# when metrics are provided, so we have to maximize it.
trainer.hyperparameter_search(
    direction="maximize", 
    backend="ray", 
    n_trials=10, # number of trials
)
```

The call to `trainer.hyperparameter_search` creates the following error:
```
TypeError: can't pickle _thread.RLock objects
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-3-af3e9d5e18dd> in <module>
     42     direction="maximize",
     43     backend="ray",
---> 44     n_trials=10, # number of trials
     45 )
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/transformers/trainer.py in hyperparameter_search(self, hp_space, compute_objective, n_trials, direction, backend, hp_name, **kwargs)
   1457 
   1458         run_hp_search = run_hp_search_optuna if backend == HPSearchBackend.OPTUNA else run_hp_search_ray
-> 1459         best_run = run_hp_search(self, n_trials, direction, **kwargs)
   1460 
   1461         self.hp_search_backend = None
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/transformers/integrations.py in run_hp_search_ray(trainer, n_trials, direction, **kwargs)
    233         config=trainer.hp_space(None),
    234         num_samples=n_trials,
--> 235         **kwargs,
    236     )
    237     best_trial = analysis.get_best_trial(metric="objective", mode=direction[:3])
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/tune/tune.py in run(run_or_experiment, name, metric, mode, stop, time_budget_s, config, resources_per_trial, num_samples, local_dir, search_alg, scheduler, keep_checkpoints_num, checkpoint_score_attr, checkpoint_freq, checkpoint_at_end, verbose, progress_reporter, log_to_file, trial_name_creator, trial_dirname_creator, sync_config, export_formats, max_failures, fail_fast, restore, server_port, resume, queue_trials, reuse_actors, trial_executor, raise_on_failed_trial, callbacks, loggers, ray_auto_init, run_errored_only, global_checkpoint_period, with_server, upload_dir, sync_to_cloud, sync_to_driver, sync_on_checkpoint)
    296 
    297     trial_executor = trial_executor or RayTrialExecutor(
--> 298         reuse_actors=reuse_actors, queue_trials=queue_trials)
    299     if isinstance(run_or_experiment, list):
    300         experiments = run_or_experiment
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/tune/ray_trial_executor.py in __init__(self, queue_trials, reuse_actors, ray_auto_init, refresh_period)
    198                         "For cluster usage or custom Ray initialization, "
    199                         "call `ray.init(...)` before `tune.run`.")
--> 200             ray.init()
    201 
    202         if ray.is_initialized():
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/_private/client_mode_hook.py in wrapper(*args, **kwargs)
     45         if client_mode_enabled and _client_hook_enabled:
     46             return getattr(ray, func.__name__)(*args, **kwargs)
---> 47         return func(*args, **kwargs)
     48 
     49     return wrapper
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/worker.py in init(address, num_cpus, num_gpus, resources, object_store_memory, local_mode, ignore_reinit_error, include_dashboard, dashboard_host, dashboard_port, job_config, configure_logging, logging_level, logging_format, log_to_driver, _enable_object_reconstruction, _redis_max_memory, _plasma_directory, _node_ip_address, _driver_object_store_memory, _memory, _redis_password, _java_worker_options, _temp_dir, _lru_evict, _metrics_export_port, _system_config)
    771 
    772     for hook in _post_init_hooks:
--> 773         hook()
    774 
    775     node_id = global_worker.core_worker.get_current_node_id()
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/tune/registry.py in flush(self)
    169     def flush(self):
    170         for k, v in self.to_flush.items():
--> 171             self.references[k] = ray.put(v)
    172         self.to_flush.clear()
    173 
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/_private/client_mode_hook.py in wrapper(*args, **kwargs)
     45         if client_mode_enabled and _client_hook_enabled:
     46             return getattr(ray, func.__name__)(*args, **kwargs)
---> 47         return func(*args, **kwargs)
     48 
     49     return wrapper
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/worker.py in put(value)
   1487     with profiling.profile("ray.put"):
   1488         try:
-> 1489             object_ref = worker.put_object(value)
   1490         except ObjectStoreFullError:
   1491             logger.info(
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/worker.py in put_object(self, value, object_ref)
    267                                         "inserting with an ObjectRef")
    268 
--> 269         serialized_value = self.get_serialization_context().serialize(value)
    270         # This *must* be the first place that we construct this python
    271         # ObjectRef because an entry with 0 local references is created when
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/serialization.py in serialize(self, value)
    317             return RawSerializedObject(value)
    318         else:
--> 319             return self._serialize_to_msgpack(value)
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/serialization.py in _serialize_to_msgpack(self, value)
    297             metadata = ray_constants.OBJECT_METADATA_TYPE_PYTHON
    298             pickle5_serialized_object = \
--> 299                 self._serialize_to_pickle5(metadata, python_objects)
    300         else:
    301             pickle5_serialized_object = None
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/serialization.py in _serialize_to_pickle5(self, metadata, value)
    257         except Exception as e:
    258             self.get_and_clear_contained_object_refs()
--> 259             raise e
    260         finally:
    261             self.set_out_of_band_serialization()
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/serialization.py in _serialize_to_pickle5(self, metadata, value)
    254             self.set_in_band_serialization()
    255             inband = pickle.dumps(
--> 256                 value, protocol=5, buffer_callback=writer.buffer_callback)
    257         except Exception as e:
    258             self.get_and_clear_contained_object_refs()
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/cloudpickle/cloudpickle_fast.py in dumps(obj, protocol, buffer_callback)
     71                 file, protocol=protocol, buffer_callback=buffer_callback
     72             )
---> 73             cp.dump(obj)
     74             return file.getvalue()
     75 
~/miniconda3/envs/nupic.research/lib/python3.7/site-packages/ray/cloudpickle/cloudpickle_fast.py in dump(self, obj)
    572     def dump(self, obj):
    573         try:
--> 574             return Pickler.dump(self, obj)
    575         except RuntimeError as e:
    576             if "recursion" in e.args[0]:
TypeError: can't pickle _thread.RLock objects
```

## Expected behavior
Ray seems to be pickling the Trainer although it can't. On an earlier version `transformers==4.4.2` this was not an issue. I believe some update in 4.5 is causing this and the above script should be able to complete without any pickling error.




Comments:
Same error as this: https://github.com/huggingface/transformers/issues/11249. We are currently looking into this. In the meantime, can you use an earlier version as you suggested? Thanks.

Will do. Didn't see that other issue there. Thanks for the help.