Make the Scheduler adjust the steps taken relative to the gradient accumulation steps

[muellerzr]

Let the AcceleratedScheduler handle gradient accumulation steppage

What does this add?

This PR adjusts the logic in the AcceleratedScheduler to take into account gradient accumulation steps.

Who is it for?

Closes #1170
Closes #1160

Why is it needed?

Currently a behavior does not exist that will automatically "cut" the LR scheduler of a user if they pass in gradient_accumulation_steps, so unless they are careful and adjust their LR scheduler beforehand, they're not actually stepping properly with the lr scheduler.

What parts of the API does this impact?

User-facing:

A new GradientAccumulationPlugin is being added which will handle gradient_accumulation_steps and optionally disabling the extra steppage involved with the scheduler when performing gradient accumulation

Internal structure:

AcceleratedScheduler's step function now will run n*num_processes where n==gradient_accumulation_steps to account for the difference.

To test performance, I ran the equivalent training of gradient_accumulation_steps==2 on a bs of 16, and a regular batch size of 32, with negligible performance differences (accuracy += 0.5%, F1 of 0.3). Likely due to batch norm layers

Usage Example:

When building the Accelerator, pass in adjust_scheduler_to_accumulation (default False) to enable this behavior:

    accelerator = Accelerator(
        cpu=args.cpu, 
        mixed_precision=args.mixed_precision, 
        gradient_accumulation_steps=2,
    )

    accelerator = Accelerator(
        cpu=args.cpu, 
        mixed_precision=args.mixed_precision, 
        gradient_accumulation_plugin=GradientAccumulationPlugin(adjust_scheduler=False, num_steps=2)
    )

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sgugger]

Thanks for working on this. Just not sure about the API (adding yet another argument to the Accelerator) for this.

[muellerzr]

@sgugger good for review, a few notes:

• You'll notice that in the test for grad_sync changes, for it to pass (and for everywhere to pass in general), having accelerator use the new stepped scheduler to True is used, and on every script we have this behavior of what should exist is what we've advertised using. As such, I think it would be good to change the default to True not False here, as users are expecting one thing rather than the other
• Should we keep gradient_accumulation_steps as an arg in Accelerator? I'm leaning towards yes if we change the default to True for reasons above, though we can also just keep it there for usability. (In my eyes, users would likely not need to adjust the other param, hence we only need to expose gradient_accumulation_steps there)

[sgugger]

Thanks for iterating! I agree that the default should be True for adjust_scheduler. The gradient accumulation API is quite new anyway, so it's okay if it's adjusted a bit like this.

[sgugger]

This should be True by default I think.