Trainer training batches estimation for map-style datasets seems wrong

[robertomest]

Hello,

It seems that setting max_steps and not setting max_epochs for map-style datasets will lead to estimated_stepping_batches returning the batches for a single epoch. From the docstring, it indicates that this would be the estimated stepping batches for the whole training. I believe the line below (with max(self.max_epochs, 1)) is not appropriate in this scenario. Shouldn't it always return max_steps in case max_epochs == -1?
https://github.com/Lightning-AI/lightning/blob/deaadc157b87b653f1c3aa7926a5ff092ab81863/src/pytorch_lightning/trainer/trainer.py#L2783

Thanks!

cc @carmocca @justusschock @ananthsub @ninginthecloud @rohitgr7

[awaelchli]

@rohitgr7 Do you have an idea what @robertomest means here?

[robertomest]

Hi @awaelchli,
Sorry if my explanation was confusing, let me try to make it clearer. The way I understand estimate_stepping_batches should work is that it should return the number of stepping batches for the whole training.

In cases where you have a dataset with specific length, and you define a maximum number of epochs, the number of steps would be computed based on batch size, accumulation steps, etc (this seems all good and you take care of a lot of stuff). For cases where you supply a maximum number of steps, but not number of epochs, then you should use max_steps instead of assuming one epoch (which is what max(self.max_epochs, 1) does), regardless of the dataset length, since max_steps will necessarily be the defining factor.

From what I understand, the current implementation will use max_steps if the dataset has infinite length (iterable dataset), but it will default to one epoch for map-style datasets.

Does that make more sense?

[rohitgr7]

https://github.com/Lightning-AI/lightning/blob/deaadc157b87b653f1c3aa7926a5ff092ab81863/src/pytorch_lightning/trainer/trainer.py#L2785

but it will consider max_steps in the very next line to make sure in case your total batches are less than max_stepsit computes this value correctly since it's not going to run for more than that during training.

[robertomest]

This next line will make sure to pick the lower of both conditions, which makes sense to me. The problem is that if max_epochs == -1, line 2783 will evaluate to a single epoch, and max_estimated_steps will be limited by that value. Below is a simple repro of the issue: in this case I'd expect self.trainer.estimated_stepping_batches to evaluate to 1000, since max_epochs is not used, but it's computed as a single epoch of steps.

import pytorch_lightning as pl
import torch
from torch.utils.data import Dataset, DataLoader

class ExampleDataset(Dataset):
    def __init__(self, size: int):
        self.data = torch.randn(size, 10)
    
    def __getitem__(self, idx: int):
        return self.data[idx]
    
    def __len__(self):
        return len(self.data)

class ExampleModule(pl.LightningModule):
    def configure_optimizers(self):
        print(f"{self.trainer.estimated_stepping_batches=}")
        raise RuntimeError()
    
    def training_step(self, *args):
        raise RuntimeError()

dataset = ExampleDataset(100)
loader = DataLoader(dataset, batch_size=5)
pl_module = ExampleModule()
trainer = pl.Trainer(logger=False, max_epochs=-1, max_steps=1000)

trainer.fit(pl_module, loader)
# prints self.trainer.estimated_stepping_batches=20

[rohitgr7]

ah! good catch! are you interested in sending a fix?

[robertomest]

Sure, I can do that 🙂