Scheduler Not Pickleable
Author: iamNCJ
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.4.2
- Platform: Linux-3.10.0-1127.13.1.el7.x86_64-x86_64-with-glibc2.10
- Python version: 3.8.8 (Anaconda)
- PyTorch version (GPU): 1.8.0+cu111 (True)
- Tensorflow version (GPU): not installed (NA)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes

### Who can help

This seems to be a bug concerning the `Optimization` class.

## Information

Model I am using (Bert, XLNet ...): BertMultipleChoice

The problem arises when using:
* [ ] my own modified scripts: I'm using transformers with Pytorch Lightning, and the distributed training function is provided by PyTorch Lightening.

The tasks I am working on is:
* Reading Comprehensive on RACE Dataset

## To reproduce

Steps to reproduce the behavior:

1. load RACE into a datamodule
2. finetune BertMultipleChoice on this datamodule
3. start training with `gpus=-1`

Output:
```text
Traceback (most recent call last):
  File "train.local.py", line 35, in <module>
    trainer.fit(model, dm)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 498, in fit
    self.dispatch()
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 545, in dispatch
    self.accelerator.start_training(self)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 73, in start_training
    self.training_type_plugin.start_training(trainer)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp_spawn.py", line 106, in start_training
    mp.spawn(self.new_process, **self.mp_spawn_kwargs)  
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 230, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 179, in start_processes
    process.start()
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)  
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
AttributeError: Can't pickle local object 'get_linear_schedule_with_warmup.<locals>.lr_lambda'
```

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->

It should start training on all gpus.

Comments:
Succeed to reproduce with another model and dataset.

https://gist.github.com/iamNCJ/a30afcbac392f6036bed65198ce5295e

[gist](https://gist.github.com/iamNCJ/a30afcbac392f6036bed65198ce5295e)

This gist is derived from [an example provided by the pytorch lightening team](https://colab.research.google.com/github/PytorchLightning/pytorch-lightning/blob/master/notebooks/04-transformers-text-classification.ipynb), but it also causes this problem with multiple gpus.

Output:
```text
Traceback (most recent call last):  File "glue.py", line 272, in <module>
    trainer.fit(model, dm)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 498, in fit
    self.dispatch()
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 545, in dispatch
    self.accelerator.start_training(self)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 73, in start_training
    self.training_type_plugin.start_training(trainer)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp_spawn.py", line 106, in start_training
    mp.spawn(self.new_process, **self.mp_spawn_kwargs)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 230, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 179, in start_processes
    process.start()
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/home/<user>/.conda/envs/RACE/lib/python3.8/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
AttributeError: Can't pickle local object 'get_linear_schedule_with_warmup.<locals>.lr_lambda'
```

I succeed to start training after using DDP instead of DDP Spawn, since DDP Spawn forces the model to be pickleable but DDP doesn't, but I still wonder if it's possible to make the scheduler pickleable.

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.