Issue with gradient accumulation in CodeParrot example
Author: lvwerra
Pull request: no

There is a bug in the gradient accumulation that causes the [training script](https://github.com/huggingface/transformers/blob/main/examples/research_projects/codeparrot/scripts/codeparrot_training.py) to run slower than necessary. Currently we have the following:

```python
for step, batch in enumerate(train_dataloader, start=1):
    if args.resume_from_checkpoint and step < resume_step:
        continue  # we need to skip steps until we reach the resumed step
    loss = model(batch, labels=batch, use_cache=False).loss
    avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
    loss_tracking += avg_loss.item() / args.gradient_accumulation_steps
    log_metrics(step, {"samples": step * samples_per_step,
                       "loss_per_step/train": loss.item()})
    loss = loss / args.gradient_accumulation_steps
    if step % args.gradient_accumulation_steps != 0:
        # Prevent backward from doing gradient all_reduce in every step
        if accelerator.distributed_type == DistributedType.MULTI_GPU:
            with model.no_sync():
                accelerator.backward(loss)
        else:
            accelerator.backward(loss)
    else:
        lr = get_lr()
        accelerator.backward(loss)
        accelerator.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        elapsed_time = time.time() - t_start
        tflops = compute_tflops(elapsed_time, accelerator, args)
        log_metrics(
            step,
            {
                "steps": completed_steps,
                "loss/train": loss_tracking,
                "lr": lr,
                "tflops": tflops,
                "time_per_iteration": elapsed_time,
            },
        )
        t_start = time.time()
        loss_tracking = 0
        completed_steps += 1
```

When it should be something along the lines of this:

```python
for step, batch in enumerate(train_dataloader, start=1):
    with accelerator.accumulate(model):
        if args.resume_from_checkpoint and step < resume_step:
            continue  # we need to skip steps until we reach the resumed step
        lr = get_lr()   
        loss = model(batch, labels=batch, use_cache=False).loss
        avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
        loss_tracking += avg_loss.item() / args.gradient_accumulation_steps
        log_metrics(step, {"samples": step * samples_per_step,
                        "loss_per_step/train": loss.item()})
        accelerator.clip_grad_norm_(model.parameters(), 1.0)
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        if accelerator.gradient_state.sync_gradients:
            elapsed_time = time.time() - t_start
            tflops = compute_tflops(elapsed_time, accelerator, args)
            log_metrics(
                step,
                {
                    "steps": completed_steps,
                    "loss/train": loss_tracking,
                    "lr": lr,
                    "tflops": tflops,
                    "time_per_iteration": elapsed_time,
                },
            )
            t_start = time.time()
            loss_tracking = 0
            completed_steps += 1
```

We're not actually pausing the gradient accumulation. Here's an example:
https://github.com/huggingface/accelerate/blob/92d072043eb24eddf714edd578bceff07a2d9470/examples/by_feature/gradient_accumulation.py#L171-L183


And here some explanation:
https://huggingface.co/docs/accelerate/concept_guides/gradient_synchronization.

This could speed-up training up to 2x as [reported](https://github.com/muellerzr/timing_experiments) by @muellerzr! Thanks for reporting!
![Screenshot 2023-04-03 at 19 08 31](https://user-images.githubusercontent.com/8264887/229579509-5780b821-3492-4380-98d5-1b0520ea3db0.png)
cc @ArmelRandy @loubnabnl 


Comments:
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.