Lightning-AI / Lightning-AI/pytorch-lightning

Resuming not correct when `max_steps` corresponds to the end of an epoch

Open
#19,617 0 comments 0 reactions 1 assignee View on GitHub

@awaelchli is already working on this.

Since Mar 12, 2024.

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

Description

Bug description

The resuming logic for the training loop is not correct in an edge case where max_steps corresponds to the end of an epoch.

What version are you seeing the problem on?

v2.2, master

How to reproduce the bug
import shutil
import time
from pathlib import Path

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, dataloader_iter):
        batch, _, _ = next(dataloader_iter)
        batch, _, _ = next(dataloader_iter)
        loss = self(batch).sum()
        time.sleep(0.5)
        return {"loss": loss}

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


train_data = DataLoader(RandomDataset(32, 40), batch_size=2)  # length = 20

checkpoint_dir = Path.cwd() / "checkpoints"
if checkpoint_dir.is_dir():
    shutil.rmtree(checkpoint_dir)


model = BoringModel()
trainer = Trainer(
    default_root_dir=checkpoint_dir,
    max_steps=10,
    # max_epochs=1,
    enable_model_summary=False,
    accelerator="cpu",
)
trainer.fit(model, train_data)

# Resume
trainer = Trainer(
    default_root_dir=checkpoint_dir,
    max_steps=20,
    # max_epochs=2,
    enable_model_summary=False,
    accelerator="cpu",
)
trainer.fit(model, train_data, ckpt_path=checkpoint_dir / "lightning_logs/version_0/checkpoints/epoch=0-step=10.ckpt")
Error messages and logs

The expected behavior is that the loop/progress bar behaves the same when setting max_epochs (commented out above). Instead, when reaching max_steps, the epoch doesn't get bumped and when resuming the checkpoint, we resume in a mid-epoch state.

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 @carmocca @justusschock

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.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.