Make schedulers picklable
Author: viktor-shcherb
Pull request: no

### Feature request

Change lambda functions passed to `LambdaLR` in `get_constant_schedule`, `get_constant_schedule_with_warmup`, `get_linear_schedule_with_warmup`, `get_cosine_schedule_with_warmup`, `get_cosine_with_hard_restarts_schedule_with_warmup` and `get_polynomial_decay_schedule_with_warmup` to callable objects.

### Motivation

Python cannot serialize lambda and local functions. Torch created a workaround around this in their `state_dict` method of `LambdaLR` by not returning any non-picklable functions:
```python
        ...
        for idx, fn in enumerate(self.lr_lambdas):
            if not isinstance(fn, types.FunctionType):
                state_dict['lr_lambdas'][idx] = fn.__dict__.copy()

        return state_dict
```
While this approach is fine when LR schedule is constant and deterministic, it makes it impossible to change the schedule mid training dynamically using lambda functions since any changes will not be saved to checkpoints.

In my particular case I wanted to implement a dynamic LR schedule based on evaluation metrics. I've implemented a wrapper around `LambdaLR` that applies transformation `fn: float -> float` to existing LR schedule:
```python
class LambdaWrapper:

    def __init__(self, lr_lamda: Callable[[Union[float, int]], float], wrapper_function: Callable[[float], float]):
        self._wrapper_function = wrapper_function
        self._lr_lambda = lr_lamda

    def __call__(self, x: Union[float, int]):
        return self._wrapper_function(self._lr_lambda(x))


class DynamicScheduler:

    def __init__(self, lr_scheduler: LambdaLR):
        self._scheduler = lr_scheduler

    def __getattr__(self, item):
        # Calling the super class to avoid recursion
        return getattr(super(DynamicScheduler, self).__getattribute__('_scheduler'), item)

    def wrap_schedule(self, fn: Callable[[float], float]):
        """If you want this object to be picklable, pass only picklable callable objects as `fn`!"""
        wrappers_builder = partial(LambdaWrapper, wrapper_function=fn)  # wrap in callable object to preserve picklability
        self._scheduler.lr_lambdas = list(map(wrappers_builder, self._scheduler.lr_lambdas))
```
I've taken special care to preserve picklability, however, since `LambdaLR` instances created by `transformers` library hold lambda and local functions in them, pickling of `DynamicScheduler` (as well as it's state, which is the same as the wrapped `LambdaLR` state) fails.

While reimplementing dynamic scheduling with lambda functions will allow the `torch` workaround that handles lambda functions in scheduler, the whole point of dynamic scheduling will be lost since the complex dynamically constructed lambdas: `f_n(f_n-1(...f_1(schedule(x))...))` will fall back to their default state: `schedule(x)`.

Here is the callback I use to track evaluation metrics for anyone interested:
```python
def get_warmup_steps(args: TrainingArguments, state: TrainerState) -> int:
    return (
        args.warmup_steps
        if args.warmup_steps > 0
        else math.ceil(state.max_steps * args.warmup_ratio)
    )


class DecreaseLRTransformer:
    def __init__(self, decrease_ratio: float):
        if decrease_ratio < 0.0 or decrease_ratio > 1.0:
            raise ValueError('Decrease ratio should be within [1.0, 0.0]')
        self._decrease_ratio = decrease_ratio

    def __call__(self, lr: float):
        return self._decrease_ratio * lr


# Developer notice (may change in the future versions of transformers):
# All kwargs have the following fields set: model, tokenizer, optimizer, lr_scheduler, train_dataloader, eval_dataloader

class LRDecreaseCallback(TrainerCallback):
    """
    A [`TrainerCallback`] that handles learning rate decrease based on evaluation metrics.
    """

    def __init__(self, decrease_ratio: float, patience: int, *, decrease_on_warmup: bool = False, decrease_threshold: float = 0.0):
        self._transformer = DecreaseLRTransformer(decrease_ratio)
        self._patience = patience
        self._decrease_on_warmup = decrease_on_warmup
        self._decrease_threshold = decrease_threshold

        self._failed_checks = 0

    def _metric_improved(self, new_metric: float, old_metric: float, *, greater_is_better: bool = True) -> bool:
        operator = np.greater if greater_is_better else np.less
        return operator(new_metric, old_metric) and abs(new_metric - old_metric) > self._decrease_threshold

    def check_metric_value(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metric_value: float):
        # best_metric is set by code for load_best_model
        no_metric = (state.best_metric is None)

        warmup_steps = get_warmup_steps(args, state)
        skip_warmup = (self._decrease_on_warmup and warmup_steps >= state.global_step)

        if skip_warmup:
            return

        if no_metric or self._metric_improved(metric_value, state.best_metric, greater_is_better=args.greater_is_better):
            self._failed_checks = 0
            control.should_save = True
        else:
            self._failed_checks += 1

    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        if args.metric_for_best_model is None:
            raise ValueError(f"{self.__class__.__name__} requires metric_for_best_model to be defined defined")

        if args.evaluation_strategy == IntervalStrategy.NO:
            raise ValueError(f"{self.__class__.__name__} requires IntervalStrategy of steps or epoch")

    def on_evaluate(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        metrics: Dict[str, float] = kwargs['metrics']
        lr_scheduler = kwargs['lr_scheduler']
        if not isinstance(lr_scheduler, DynamicScheduler):
            logger.warning(f'{self.__class__.__name__} is not compatible with {lr_scheduler.__class__.__name__} scheduler! '
                           f'Wrap your scheduler with {DynamicScheduler.__class__.__name__} to change LR dynamically. '
                           f'{self.__class__.__name__} is disabled!')
            return

        metric_to_check = args.metric_for_best_model
        if not metric_to_check.startswith("eval_"):
            metric_to_check = f"eval_{metric_to_check}"
        metric_value = metrics.get(metric_to_check)

        if metric_value is None:
            logger.warning(f"{self.__class__.__name__} required metric_for_best_model, "
                           f"but did not find {metric_to_check} in evaluation metrics. {self.__class__.__name__} is disabled!")
            return

        self.check_metric_value(args, state, control, metric_value)
        if self._failed_checks >= self._patience:
            lr_scheduler.wrap_schedule(self._transformer)
            self._failed_checks = 0

    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        logs: Dict[str, float] = kwargs['logs']
        logs['lr_decrease_patience'] = (self._patience - self._failed_checks) / self._patience
``` 

### Your contribution

The simplest and the cleanest workaround would be to make the local functions global:

Intead of:
```python
def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return max(
            0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
        )

    return LambdaLR(optimizer, lr_lambda, last_epoch)
```

Do this:
```python
def _linear_schedule_with_warmup_step(current_step: int, *, num_warmup_steps: int, num_training_steps: int) -> float:
      if current_step < num_warmup_steps:
          return float(current_step) / float(max(1, num_warmup_steps))
      return max(
          0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
      )

def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
    schedule = partial(_linear_schedule_with_warmup_step, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)
    return LambdaLR(optimizer, schedule, last_epoch)
```
When created with global functions, partial function are picklable:
```python
>>>from functools import partial

>>>import pickle

>>>def f(x):
...    print(x)
    
>>>with open('f.pkl', 'wb') as file:
...    pickle.dump(partial(f, x='Dog'), file)
    
>>>with open('f.pkl', 'rb') as file:
...    unpickled_f = pickle.load(file)
    
>>>unpickled_f()
Dog
```

The fix is straightforward and I can create a PR. Nonetheless, it would be my first contribution so I might need some help along the way.

Comments:
Thanks for explaining your issue in depth, and happy to review a PR!