Gradient accumulation calcluation may be incorrect

[tyler-rt]

Bug description

See https://unsloth.ai/blog/gradient for in depth explanation, and this PR for how huggingface fixed it.

I verified that I see worse performance with gradient accumulation than multiple devices, so I suspect this bug also applies to Lightning

What version are you seeing the problem on?

v2.4

How to reproduce the bug

No response

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.4.0):
#- PyTorch Version (e.g., 2.4):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):

More info

No response

[andyschuster]

I am experiencing the same issue with the latest lightning release. Gradient accumulation delivers worse results compared to increasing the batch size to the same value, which leads to a very noticabel degradation when using uneven length (and padded) samples when training a transformer model.

HuggingFace has a useful writeup of the fix triggered by the same unsloth blog as mentioned above here: https://huggingface.co/blog/gradient_accumulation

[bhimrazy]

Hi @tyler-rt @andyschuster , thanks for raising this! 🙌

Would you be interested in adding a minimal reproducible example to help the community reproduce and debug the issue more effectively?
You could also consider sharing a Lightning Studio link if that’s easier.

Thanks again!

[andyschuster]

OK, here's an example of a minimalist transformer-style language model implementation that showcases the problem.
The underlying reason for the issue is further explained as something other libraries ran into and fixed in the blogs linked above: 1, 2

import lightning as L
import numpy as np
import torch
import torch.nn.functional as F
from torch import Tensor, nn


class MyModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(5, 16)
        self.transformer = nn.TransformerEncoder(
            encoder_layer=nn.TransformerEncoderLayer(
                d_model=16,
                nhead=1,
                dim_feedforward=32,
                batch_first=True,
            ),
            num_layers=2,
        )
        self.head = nn.Linear(16, 5, bias=False)

    def training_step(self, batch: tuple[Tensor, Tensor]):
        tokens, mask = batch
        predicted = self.head(
            self.transformer(
                self.embedding(tokens[:, :-1])
            )
        )
        expected = tokens[:, 1:]
        mask = mask[:, 1:] != 0
        loss = F.cross_entropy(
            predicted[mask].view(-1, predicted.shape[2]),
            expected[mask].view(-1),
        )
        self.log("train_loss", loss, on_step=False, on_epoch=True, prog_bar=True)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=3e-3)


class MyDataset(torch.utils.data.Dataset):
    def __init__(self) -> None:
        super().__init__()
        self.token_strings = [
            np.array([1, 2, 3, 4]),
            np.array([1, 2, 3, 4, 1, 2, 3, 4]),
            np.array([1, 2, 3, 4, 1, 2, 3, 4, 4, 3, 2, 1]),
            np.array([1, 2, 3, 4, 1, 2, 3, 4, 4, 3, 2, 1, 4, 3, 2, 1]),
        ]

    def __len__(self):
        return len(self.token_strings)

    def __getitem__(self, idx) -> tuple[Tensor, Tensor]:
        tokens = self.token_strings[idx]
        mask = np.full_like(tokens, 1)
        return (
            torch.from_numpy(tokens).type(torch.long),
            torch.from_numpy(mask).type(torch.uint8),
        )

    @staticmethod
    def collate_padded(batch: list[tuple[Tensor, Tensor]]) -> tuple[Tensor, Tensor]:
        tokens, mask = zip(*batch)
        tokens_padded: Tensor = torch.nn.utils.rnn.pad_sequence(tokens, batch_first=True, padding_value=0)
        mask_padded: Tensor = torch.nn.utils.rnn.pad_sequence(mask, batch_first=True, padding_value=0)
        return tokens_padded, mask_padded

I would expect the training loss progression to be similar when
a) training with batch_size of 4 and no gradient accumulation, or
b) with batch_size of 1 and gradient accumulation across 4
However the examples below lead to different training loss progressions.

L.pytorch.seed_everything(42)
model = MyModel()

dataloader = torch.utils.data.DataLoader(
    dataset=MyDataset(),
    batch_size=4,  # No gradient accumulation: train on full batches of 4
    collate_fn=MyDataset.collate_padded,
)

trainer = L.Trainer(
    max_epochs=50,
)
trainer.fit(model, train_dataloaders=dataloader)

L.pytorch.seed_everything(42)
model = MyModel()

dataloader = torch.utils.data.DataLoader(
    dataset=MyDataset(),
    batch_size=1,  # Batch size reduced to 1
    collate_fn=MyDataset.collate_padded,
)

trainer = L.Trainer(
    accumulate_grad_batches=4,  # Gradient accumulation across 4 batches
    max_epochs=50,
)

trainer.fit(model, train_dataloaders=dataloader)

[andyschuster]

@bhimrazy, are you able to replicate with this code snippet?

[bhimrazy]

Thanks @andyschuster for the clear example! 🙌
I’ll test it out shortly and get back with my observations.

[Sohaib-Ahmed21]

@SkafteNicki The gradient accumulation discrepancy arises from how CrossEntropyLoss normalizes losses.

In Lightning, each micro-batch computes a mean loss:

loss_i = CrossEntropyLoss(y_pred_i, y_true_i,reduction="mean")

These are later accumulated for N micro-batches as:

$$L_{accumulated} = \frac{1}{N} \sum_{i=1}^{N} loss_i$$

instead of the correct large-batch equivalent:

$$L_{true} = \frac{\sum_{i=1}^{N} \text{sum}_i}{\sum_{i=1}^{N} \text{count}_i}$$

where
• sum_i is the CrossEntropyLoss with reduction="sum" for micro-batch i,
• count_i is the number of valid elements (e.g., tokens or samples not equal to ignore_index) contributing to that loss for micro-batch i.

This effectively re-averages already-averaged losses, producing gradients that differ from true large-batch training(as detailed in the Unsloth article.

Since Lightning only receives the final reduced loss tensor from the user’s training_step, it has no visibility into which loss function (e.g. CrossEntropyLoss) or reduction mode (mean, sum) was applied.

This makes it difficult to apply a targeted fix only for CrossEntropy-based cases.
I have been following this issue in recent days and would appreciate your thoughts on the best way to handle this within Lightning’s internal codebase, thanks!

[SkafteNicki]

@Sohaib-Ahmed21 thanks for reaching out. I think you are correctly hitting on a hard limitation of lightning: we are only receiving the final/reduced loss and we have no idea how it was calculated. This is core part of lightning and not really anything that can be touched. So changing the internals of lightning backward pass is probably not something we can do.

However, I do wonder if we in some way can update the GradientAccumulation callback. Especially the on_before_backward is interesting because we can intercept the loss before backward and on_before_optimizer_step hook is called after accumulation but before stepping.

[Sohaib-Ahmed21]

@Sohaib-Ahmed21 thanks for reaching out. I think you are correctly hitting on a hard limitation of lightning: we are only receiving the final/reduced loss and we have no idea how it was calculated. This is core part of lightning and not really anything that can be touched. So changing the internals of lightning backward pass is probably not something we can do.

However, I do wonder if we in some way can update the GradientAccumulation callback. Especially the on_before_backward is interesting because we can intercept the loss before backward and on_before_optimizer_step hook is called after accumulation but before stepping.

Thanks for clarifying! so does that mean this is a blocker for this issue? As in callbacks also, only final/reduced loss is accessible.

[SkafteNicki]

So here is my quick thoughts: we cannot natively implement this into lightning in the sense that lightning can automatically take care of everything due to the above limitations. Instead, what is the minimum extra information that we should ask the user to provide to implement this feature and what is a easy interface for the user to provide this information? Essentially, if we can do:

def training_step(self, batch, batch_idx):
    x, y = batch
    pred = self.forward(x)
    loss = self.loss_fn(pred, y)  # this is normalized cross entropy
    num_ignore_val = sum(y == self.ignore_val)
    return {"loss": loss, "num_ignore_val": num_ignore_val}

and then GradientAccumulation should be able to calculate loss * num_ignore_val which is the unnormalized value, backprop and accumulate that and then normalize by sum(num_ignore_val for all batches in the end.

[Sohaib-Ahmed21]

Okay, thanks! I'll look into what and how we can take from the user, and update accordingly.