Lightning-AI / Lightning-AI/pytorch-lightning

Handle gradient accumulations at the end of epoch differently

Open
#19,987 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug optimization ver: 2.0.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

At the end of an epoch with accumulate_grad_batches>1 the dataloader may run out of data before the required number of accumulations. The lightning docs do not say what happens. It could

1. not update the gradients
2. update gradients correctly but with an effectively smaller batch size
3. update gradients incorrectly because the gradients are scaled by accumulate_grad_batches instead of the actual number of accumulations

My experiments suggest its option 3 but happy to be wrong.

image

### What version are you seeing the problem on?

v2.0, v2.2

### How to reproduce the bug

```python
import torch
from torch.optim.optimizer import Optimizer
from torch.utils.data import DataLoader, TensorDataset
import pytorch_lightning as pl

# Generate some dummy data
X = torch.randn(1000, 28*28) # 1000 samples of 28*28 features
y = torch.randint(0, 10, (1000,)) # 1000 labels for 10 classes

# Create a TensorDataset
dataset = TensorDataset(X, y)
# Create DataLoaders
train_loader = DataLoader(dataset, batch_size=1, shuffle=True)

class SimpleModel(pl.LightningModule):
def __init__(self):
super(SimpleModel, self).__init__()
self.layer_1 = torch.nn.Linear(28 * 28, 128)
self.layer_2 = torch.nn.Linear(128, 10)
self.loss = 0
self.steps = 0

def forward(self, x):
x = x.view(x.size(0), -1)
x = torch.relu(self.layer_1(x))
return torch.log_softmax(self.layer_2(x), dim=1)

def training_step(self, batch, batch_idx):
x, y = batch
y_hat = self(x)
loss = torch.nn.functional.nll_loss(y_hat, y)
self.loss += loss.detach()
self.steps += 1
self.log('train_loss_batch', loss.detach(), on_step=True, on_epoch=False, logger=True)
return loss

def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=0.001)

def on_before_optimizer_step(self, optimizer):
self.log('train_loss', self.loss/self.steps, prog_bar=True, logger=True)
self.log('accumulations', self.steps, prog_bar=False, logger=True)
self.loss = 0
self.steps = 0

grad_norm = 0
for param in self.parameters():
if param.grad is not None:
grad_norm += param.grad.norm(2).item()
self.log('grad norm', grad_norm, prog_bar=False, logger=True)

return super().on_before_optimizer_step(optimizer)

model = SimpleModel()
trainer = pl.Trainer(max_epochs=5, accumulate_grad_batches=32, log_every_n_steps=1)
trainer.fit(model, train_loader)
```

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):
```

### More info

_No response_

cc @borda

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by running the provided Trainer reproduction with accumulate_grad_batches=32 and inspect the training_step and optimizer-step flow at the end of the dataloader. Determine which of the three described outcomes occurs, then trace the gradient scaling and define the expected behavior for a partial accumulation before adding coverage for that case.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.