Fix correct gradient accumulation

[younesbelkada]

What does this PR do?

This PR correctly fixes the gradient accumulation, firstly introduced in #220

Fixes #321

cc @lvwerra

[HuggingFaceDocBuilderDev]

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

[vwxyzjn]

The PR makes sense! Would it be possible to verify that it in fact can reproduce the exact same gradient with different number of gradient accumulation steps?

Here is a snippet:

import torch
from torch.utils.data import TensorDataset, DataLoader
import copy

# seed
torch.manual_seed(0)

# define toy inputs and labels
x = torch.tensor([1., 2., 3., 4.])
y = torch.tensor([2., 4., 6., 8.])

# define dataset and dataloader
dataset = TensorDataset(x, y)
dataloader = DataLoader(dataset, batch_size=2)

# define model, optimizer and loss function
model = torch.nn.Linear(1, 1)

# clone the model
model_clone = copy.deepcopy(model)
criterion = torch.nn.MSELoss()
accumulation_steps = 2

# loop over batches
for i, (inputs, labels) in enumerate(dataloader):
    # reshape inputs and labels
    inputs = inputs.view(-1, 1)
    labels = labels.view(-1, 1)
    # forward pass
    outputs = model(inputs)
    loss = criterion(outputs, labels) / accumulation_steps
    # backward pass
    loss.backward()
    # check if accumulation is done
    if (i + 1) % accumulation_steps == 0:
        print("w/ accumulation, the final model grad is", model.weight.grad) 
        break

loss = criterion(model_clone(x.view(-1, 1)), y.view(-1, 1))
loss.backward()
print("w/o accumulation, the final model grad is", model_clone.weight.grad)

w/ accumulation, the final model grad is tensor([[-27.4301]])
w/o accumulation, the final model grad is tensor([[-27.4301]])

[younesbelkada]

Yes definitelty! Will add a test for that, thanks a lot for the snippet and the pointer!

[vwxyzjn]

LGTM! Thanks so much @younesbelkada