Add Tri-Stage Scheduler, proposed in SpecAugment
Author: jp1924
Pull request: no

### Feature request

paper: [SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition](https://arxiv.org/abs/1904.08779)
code: [Fairseq Tri-Stage Scheudler](https://github.com/facebookresearch/fairseq/blob/main/fairseq/optim/lr_scheduler/tri_stage_lr_scheduler.py)

i want to add Tri-Stage Scheduler in huggingface


### Motivation

i have two motivation
- first, many ASR model using tri-stage scheduler on training, typically wav2vec2 in this case
- second, when i making the model, use tri-stage scheduler, so i thought it'd be better to post it while I'm making it.


### Your contribution

maybe it's need to modify optimization.py, trainer.py, trainingargument.py codes
```python
import math
from typing import Tuple, Union

from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR


# [NOTE]: copied from https://github.com/facebookresearch/fairseq/blob/main/fairseq/optim/lr_scheduler/tri_stage_lr_scheduler.py
def get_tri_stage_scheduler_with_warmup(
    optimizer: Optimizer,
    num_training_steps: int,
    final_lr: float,
    num_warmup_steps: Union[int, float],
    num_hold_steps: Union[int, float],
    num_decay_steps: Union[int, float],
    last_epoch: int = -1,
) -> LambdaLR:
    default_lr = optimizer.defaults["lr"]

    check_warmup_type = isinstance(num_warmup_steps, int)
    warmup_steps = num_warmup_steps if check_warmup_type else num_warmup_steps * num_training_steps

    check_hold_type = isinstance(num_hold_steps, int)
    hold_steps = num_hold_steps if check_hold_type else num_hold_steps * num_training_steps

    check_decay_type = isinstance(num_decay_steps, int)
    decay_steps = num_decay_steps if check_decay_type else num_decay_steps * num_training_steps

    if not (warmup_steps + hold_steps + decay_steps) <= num_training_steps:
        raise ValueError("must don't exceed max_steps. but lr steps exceed max_step, please setting again")

    warmup_factpr = default_lr / warmup_steps
    decay_factor = -math.log(final_lr) / decay_steps

    def _decide_stage(step: int) -> Tuple[int, int]:
        # [NOTE]: warmup(rampup) stage
        if step < warmup_steps:
            return ("warm", step)
        offset = warmup_steps

        # [NOTE]: hold stage
        if step < offset + hold_steps:
            return ("hold", step - offset)
        offset += hold_steps

        # [NOTE]: decay stage
        if step <= offset + decay_steps:
            return ("decay", step - offset)

        # [NOTE]: over stage
        return "over", step - offset

    def lr_lambda(current_step: int) -> float:
        stage, step = _decide_stage(current_step)
        if "warm" == stage:
            compensator = (current_step if current_step else 1) * default_lr
            learning_rate = (warmup_factpr * step) + compensator
        elif "hold" == stage:
            compensator = default_lr
            learning_rate = default_lr**compensator
        elif "decay" == stage:
            compensator = default_lr
            learning_rate = (default_lr**compensator) * math.exp(-decay_factor * step)
        elif "over" == stage:
            learning_rate = final_lr

        return learning_rate

    return LambdaLR(optimizer, lr_lambda, last_epoch)
```

## testing & result
I compared the fairseq tri-stage on the link with the tri-stage I made using matplotlib.  
And I compared it with linear scheduler's warmup.


### maked tri-stage
![maked_tri_stage](https://user-images.githubusercontent.com/93233241/211130497-27a34222-ec4c-4f1e-9bd3-65067afb6493.png)


### fairseq tri-stage
![fairseq_tri_stage](https://user-images.githubusercontent.com/93233241/211130502-978d43d7-35d5-420e-9223-e08d0d84dfd0.png)

### linear tri stage
![linear](https://user-images.githubusercontent.com/93233241/211130508-0521697b-dd2d-4570-ae8b-c958a34d54d3.png)


### all gather
![all_gather](https://user-images.githubusercontent.com/93233241/211130509-80c5107a-2827-4471-8caa-3331cd5a2e5b.png)


it's worked well!
If you zoom in on 0.00025 hold steps, you can also see green!

You don't have to read it after this.

## why use compensator?
actually this tri-stage have float mismatch, if you compare output lr of this tri-stage and fairseq tri-stage, values is different when you test.

it's because LamdbaLR, LamdbaLR's a part that multiplies the value from the scheduler lr by default_lr. as a result output lr differentt  fairseq tri-stage.  so i used compensator for solved this issue

Since it's mathematically corrected, there could be some errors.
- average error of 9.361272723949663e-08 at warmup stage.
- average error of -4.569185665208767e-08 at hold stage.
- average error of -4.164950459092579e-08 at decay stage.


Comments:
Note that we won't accept new optimizer/scheduler in the Transformers library as the main goal of Transformers is models :-)
You can add the scheduler directly to an example however!

ok! thank you for your reply!