Scheduler num_warmup_steps and num_training_steps
Author: KwangryeolPark
Pull request: no

### System Info

transformers: 4.34.0
Ubuntu
torch:2.1
2 * RTX3090 GPUs

t5-small, translation on wmt14 en -> de task
metric: bleu every 1,000 training iteration.

### Who can help?

@ArthurZucker
I'm runing run_translation_no_trainer.py code and using 2 GPUs with accelerate package provided in the code.
I'm using linear scheduler. 

I find that the learning rate is 0 after half of args.max_train_steps. Also the warmup schedule looks work until the half of args.num_warmup_steps. The related code is here:
`
lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps
        num_training_steps=args.max_train_steps
    )
`

I think the problem is caused from multi GPU settings and the scheduler does not consider it.
So I think the code should be modified like below
`
lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps * accelerator.num_processes,
        num_training_steps=args.max_train_steps * accelerator.num_processes,
    )
`



Here is my learning graph. As you can see, the learning rate becomes zero after half max iteration.(And also the at warmup steps) I logged 
![image](https://github.com/huggingface/transformers/assets/48284967/66bce282-bacb-4268-8831-e20216fbcf9f)

And the below is bleu scores every 1000 iteration.
![image](https://github.com/huggingface/transformers/assets/48284967/600918a4-2c12-49c2-b209-bc83d2ae1df4)




### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Change the example code like below for logging every 1000 iteration. (There is nothing much changed. I just modified the logging interval from epoch -> 1000 iteration)

`

    for epoch in range(starting_epoch, args.num_train_epochs):
        model.train()
            
        if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
            # We skip the first `n` batches in the dataloader when resuming from a checkpoint
            active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
        else:
            active_dataloader = train_dataloader
        for step, batch in enumerate(active_dataloader):
            outputs = model(**batch)
            loss = outputs.loss
            # We keep track of the loss at each epoch
            if args.with_tracking:
                total_loss += loss.detach().float()
            loss = loss / args.gradient_accumulation_steps
            accelerator.backward(loss)
            if step % args.gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()
                progress_bar.update(1)
                completed_steps += 1

            if isinstance(checkpointing_steps, int):
                if completed_steps % checkpointing_steps == 0:
                    output_dir = f"step_{completed_steps}"
                    if args.output_dir is not None:
                        output_dir = os.path.join(args.output_dir, output_dir)
                    accelerator.save_state(output_dir)
                    
            if args.with_tracking:
                total_memory += torch.cuda.max_memory_allocated() / 1024.0 / 1024.0
                
            if completed_steps >= args.max_train_steps:
                break
            
            # Log metrics
            if completed_steps % 1000 == 0:
                model.eval()

                if args.val_max_target_length is None:
                    args.val_max_target_length = args.max_target_length

                gen_kwargs = {
                    "max_length": args.val_max_target_length if args is not None else config.max_length,
                    "num_beams": args.num_beams,
                }
                samples_seen = 0
                for step, batch in enumerate(eval_dataloader):
                    with torch.no_grad():
                        generated_tokens = accelerator.unwrap_model(model).generate(
                            batch["input_ids"],
                            attention_mask=batch["attention_mask"],
                            **gen_kwargs,
                        )

                        generated_tokens = accelerator.pad_across_processes(
                            generated_tokens, dim=1, pad_index=tokenizer.pad_token_id
                        )
                        labels = batch["labels"]
                        if not args.pad_to_max_length:
                            # If we did not pad to max length, we need to pad the labels too
                            labels = accelerator.pad_across_processes(batch["labels"], dim=1, pad_index=tokenizer.pad_token_id)

                        generated_tokens = accelerator.gather(generated_tokens).cpu().numpy()
                        labels = accelerator.gather(labels).cpu().numpy()

                        if args.ignore_pad_token_for_loss:
                            # Replace -100 in the labels as we can't decode them.
                            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)

                        decoded_preds = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
                        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

                        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

                        # If we are in a multiprocess environment, the last batch has duplicates
                        if accelerator.num_processes > 1:
                            if step == len(eval_dataloader) - 1:
                                decoded_preds = decoded_preds[: len(eval_dataloader.dataset) - samples_seen]
                                decoded_labels = decoded_labels[: len(eval_dataloader.dataset) - samples_seen]
                            else:
                                samples_seen += len(decoded_labels)

                        metric.add_batch(predictions=decoded_preds, references=decoded_labels)
                eval_metric = metric.compute()
                logger.info({"bleu": eval_metric["score"]})

                if args.with_tracking:
                    accelerator.log(
                        {
                            "bleu": eval_metric["score"],
                            "train_loss": total_loss.item() / args.log_interval,
                            "epoch": epoch,
                            "step": completed_steps,
                            "train_mem": total_memory / args.log_interval,
                            "lr": optimizer.param_groups[0]["lr"],
                        },
                        step=completed_steps,
                    )
                    logger.info("Upload complete")
                    
                total_loss = 0
                total_memory = 0
                    
                if args.push_to_hub and epoch < args.num_train_epochs - 1:
                    accelerator.wait_for_everyone()
                    unwrapped_model = accelerator.unwrap_model(model)
                    unwrapped_model.save_pretrained(
                        args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
                    )
                    if accelerator.is_main_process:
                        tokenizer.save_pretrained(args.output_dir)
                        repo.push_to_hub(
                            commit_message=f"Training in progress epoch {epoch}", blocking=False, auto_lfs_prune=True
                        )

                if args.checkpointing_steps == "epoch":
                    output_dir = f"epoch_{epoch}"
                    if args.output_dir is not None:
                        output_dir = os.path.join(args.output_dir, output_dir)
                    accelerator.save_state(output_dir)
                model.train()

`


RTX 3090 x 2 without quantization and without any kinds of additional features like DeepSpeed.

Args:
`

accelerate launch run_translation.py \
    --model_name_or_path  t5-small\
    --dataset_name wmt14 \
    --source_lang en \
    --target_lang de \
    --source_prefix "translate English to German: " \
    --preprocessing_num_workers 16 \
    --per_device_train_batch_size 32 \
    --per_device_eval_batch_size 32 \
    --max_train_steps 262144 \
    --dataset_config_name de-en \
    --with_tracking \
    --log_interval 1000 \
    --max_length 512 \
    --num_beams 7 \
    --report_to wandb \
    --learning_rate 0.001 \
    --lr_scheduler_type linear \
    --num_warmup_steps 10000 \
    --wandb_id fasldkfh \
    --output_dir ./
`

### Expected behavior

I hope that the scheduler works as I intend

Comments:
Hey! Thanks for reporting, it would be great if you can isolate the bug by creating a very small reproducer with a dummy model for example! This way we can be sure that indeed the problem is from the lr scheduler on multi GPU and not the custom training loop / dataset / model 🤗 

> Hey! Thanks for reporting, it would be great if you can isolate the bug by creating a very small reproducer with a dummy model for example! This way we can be sure that indeed the problem is from the lr scheduler on multi GPU and not the custom training loop / dataset / model 🤗

Thanks for reply. What is reproducer and dummy model? Do you mean the whole code and corresponding args?

Sorry, actually it's a bit expected that the example codes should be updated to each's custom usage. By isolating the bug I mean testing with your proposed changes and see if that fixes the problem or if the issue is from the scheduler itself, in which case creating a small reproducer is usually easier for us to debug! 
If your proposed fix works as expected you can open a PR 😉 


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/main/CONTRIBUTING.md) are likely to be ignored.