How to set num_training_steps in lr_scheduler properly

Usually I call something like this to set the scheduler

from transformers import get_linear_schedule_with_warmup
scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=warmup_step, num_training_steps=num_training_steps
    )

And the num_training_steps usually equals to

t_total = int(len(train_dataloader) * num_epochs) ## len(train_dataloader) = num_batches

If I use accelerator, should I change the num_training_steps to something like this? And how to understand this:

I believe the len(train_dataloader) = batch size at each device

t_total = int(len(train_dataloader) * num_epochs // accelerator.num_processes)

All the above operations are before accelerator.prepare

[muellerzr]

@allanj everything and anything that has to do with gradient accumulation accelerate will now handle for you. Just pass in the gradient_accumulation_steps arg and make no changes to your code, as if you weren’t using gradient accumulation at all 😄

[takiholadi]

Hi @muellerzr , I don’t quite understand, I thought the sample sent by @takiholadi has no accelerator.num_processed. Why we should revise that

Since: https://github.com/huggingface/accelerate/blob/b0f8189d34fa42821ca041e2cba161db864c76b5/src/accelerate/scheduler.py#L31-L32