Fix LR scheduler based on bs from auto bs finder
Author: muellerzr
Pull request: yes

# What does this PR do?

This PR presents an alternative to https://github.com/huggingface/transformers/pull/23038, since we can't actually modify the schedulers realistically in Accelerate, this sets the correct "total batch size" based on the new bs being used, which gets trickled down to the creation of the scheduler via max steps if applicable. We can't go further than this however, as modifying the scheduler further would allude to auto-gradient accumulation, which while good, should be done after this :) 

(and might be as simple as just modifying `self.accelerator.gradient_accumulation_steps`



Comments:
_The documentation is not available anymore as the PR was closed or merged._

cc @mzamini92

> cc @mzamini92

yup that works. 👍 

@muellerzr I think there might be a problem with this PR, after building from source, the following snippet gives me a learning rate of 0 before the end of training:

```python
import evaluate
import numpy as np
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from transformers import TrainingArguments, Trainer

dataset = load_dataset("yelp_review_full", split="train[99%:]")
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")


def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)


def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)


tokenized_datasets = dataset.map(tokenize_function, batched=True).shuffle(seed=42)
tokenized_datasets = tokenized_datasets.train_test_split(test_size=0.1)

small_train_dataset = tokenized_datasets["train"]
small_eval_dataset = tokenized_datasets["test"]

metric = evaluate.load("accuracy")

training_args = TrainingArguments(output_dir="test_trainer", evaluation_strategy="epoch", report_to=[], num_train_epochs=3, auto_find_batch_size=True,
                                  lr_scheduler_type="linear", per_device_train_batch_size=1024, logging_strategy="steps", logging_steps=5, fp16=True)
model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=5)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=small_train_dataset,
    eval_dataset=small_eval_dataset,
    compute_metrics=compute_metrics,
)

trainer.train()
```

@thomas-schillaci what version of Accelerate are you using? 

@muellerzr `0.20.3` using `accelerate launch` and the following config:

```
compute_environment: LOCAL_MACHINE
distributed_type: MULTI_GPU
downcast_bf16: 'no'
gpu_ids: all
machine_rank: 0
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 3
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```

Thank you so much @thomas-schillaci! https://github.com/huggingface/transformers/pull/24758 should fix it right up, I verified it was being decayed properly and in-turn, a new lr scheduler actually being made. You can toy with `pip install git+https://github.com/huggingface/transformers@fix-train-bs` to try :)