Epochs terminating early incorrectly

[jasonxzhou]

🐛 Bug

My understanding is that custom dataloaders are expected to reset themselves and throw StopIteration when __next__ is called and there is nothing more to yield (e.g. at end of an epoch).
However, starting from pytorch-lightning == 1.6.0, Lightning appears to recognize that there are no more batches to yield, and this final __next__ call at the end of the epoch is no longer made. Subsequently, on the next epoch, when the first such __next__ call is made, the reset and StopIteration steps are triggered and the epoch ends immediately.
This results in every other epoch being skipped/terminating early:

Time for epoch 0: 0.14461779594421387
Time for epoch 1: 0.000400543212890625   <- skipped 
Time for epoch 2: 0.11101579666137695
Time for epoch 3: 0.00034236907958984375 <- skipped 
Time for epoch 4: 0.11014986038208008
Time for epoch 5: 0.0003437995910644531  <- skipped 
Time for epoch 6: 0.1101064682006836
Time for epoch 7: 0.00034737586975097656 <- skipped 
Time for epoch 8: 0.1124570369720459
Time for epoch 9: 0.0006518363952636719  <- skipped 

To Reproduce

from nvidia.dali import pipeline_def
import nvidia.dali.fn as fn
from nvidia.dali.plugin.pytorch import DALIGenericIterator, LastBatchPolicy

from pytorch_lightning import Trainer
from pytorch_lightning.core.lightning import LightningModule

import torch
import os
import time 

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

    def forward(self, x):
        x = x[0]['random'].float() # unpack random data
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)
    
    def on_train_epoch_start(self):
        self.start_time = time.time()
    
    def on_train_epoch_end(self):
        print(f'Time for epoch {self.current_epoch}: {time.time() - self.start_time}')

    def setup(self, stage=None):
        device_id = self.local_rank
        shard_id = self.global_rank
        num_shards = self.trainer.world_size
        mnist_pipeline = BoringPipeline(batch_size=2, device='gpu', device_id=device_id, shard_id=shard_id, num_shards=num_shards, num_threads=8)
        self.train_loader = DALIGenericIterator(mnist_pipeline, output_map=['random'], size=100, last_batch_policy=LastBatchPolicy.PARTIAL, auto_reset=True)

    def train_dataloader(self):
        return self.train_loader

@pipeline_def
def BoringPipeline(device, shard_id, num_shards):
    return fn.random.coin_flip(shape=32)

def run():
    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_val_batches=0,
        limit_test_batches=1,
        num_sanity_val_steps=0,
        max_epochs=10,
        enable_model_summary=False,
    )
    trainer.fit(model)

if __name__ == "__main__":
    run()

Expected behavior

On pytorch-lightning==1.5.10, there is no such issue, and we get the following times:

Time for epoch 0: 0.12346410751342773
Time for epoch 1: 0.09711241722106934
Time for epoch 2: 0.09668135643005371
Time for epoch 3: 0.09659314155578613
Time for epoch 4: 0.09643316268920898
Time for epoch 5: 0.0983741283416748
Time for epoch 6: 0.09633684158325195
Time for epoch 7: 0.09726572036743164
Time for epoch 8: 0.09633612632751465
Time for epoch 9: 0.10026073455810547

Environment

* CUDA:
        - GPU:
                - NVIDIA A40
                - NVIDIA A40
                - NVIDIA A40
                - NVIDIA A40
                - NVIDIA A40
                - NVIDIA A40
                - NVIDIA A40
                - NVIDIA A40
        - available:         True
        - version:           11.5
* Packages:
        - numpy:             1.22.3
        - pyTorch_debug:     False
        - pyTorch_version:   1.11.0+cu115
        - pytorch-lightning: 1.6.2
        - tqdm:              4.64.0
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                - ELF
        - processor:         x86_64
        - python:            3.8.10
        - version:           #1 SMP Debian 5.15.15-2~bpo11+1 (2022-02-03)

Additional context

cc @justusschock @awaelchli @ninginthecloud @rohitgr7

[awaelchli]

@jasonxzhou
Thanks for reporting

I recognize the behavior you describe because I solved that a long time ago (pre 1.0.0) by prefetching a batch in advance. This way, if we know in advance that the stop iteration gets triggered and can stop the epoch on the last batch instead of going over into an empty next epoch.

You reporting this again hints that we might have broken this behavior.

[carmocca]

If that's the problem, you can verify it (and work around it) by doing this:

from pytorch_lightning import Trainer
from pytorch_lightning.demos.boring_classes import BoringModel
from pytorch_lightning.loops import FitLoop

model = BoringModel()
trainer = Trainer(max_epochs=10)

class WorkaroundFitLoop(FitLoop):
    @property
    def prefetch_batches(self) -> int:
        return 1

trainer.fit_loop = WorkaroundFitLoop(trainer.fit_loop.min_epochs, trainer.fit_loop.max_epochs)

trainer.fit(model)

[jasonxzhou]

Verified - that is indeed the problem. Overriding trainer.fit_loop with WorkaroundFitLoop resolves the issue.

I had worked around the issue by modifying the custom dataloader provided by NVIDIA DALI (DALIGenericIterator), since it was not clear to me if this was an intended change in behavior on the PL side.

prefetch_batches also first appeared in this PR (#12101), which landed between PL 1.5.10 and 1.6.0 as expected.

[carmocca]

Here is a reduced reproduction code:

from nvidia.dali import pipeline_def
import nvidia.dali.fn as fn
from nvidia.dali.plugin.pytorch import DALIGenericIterator

from pytorch_lightning.utilities.fetching import DataFetcher


@pipeline_def
def BoringPipeline(device, shard_id, num_shards):
    return fn.random.coin_flip(shape=32)


pipeline = BoringPipeline(batch_size=2, device="cpu", device_id=0, shard_id=0, num_shards=0, num_threads=1)

# No fetcher

loader = DALIGenericIterator(pipeline, output_map=["random"], size=6, auto_reset=True)
for epoch in range(3):
    for i, _ in enumerate(loader):
        print(epoch, i)

print("-" * 20)

# Fetcher no prefetch

loader = DALIGenericIterator(pipeline, output_map=["random"], size=6, auto_reset=True)
fetcher = DataFetcher(prefetch_batches=0)
fetcher.setup(loader)
for epoch in range(3):
    for i, _ in enumerate(fetcher):
        print(epoch, i)


print("-" * 20)

# Fetcher with prefetch

loader = DALIGenericIterator(pipeline, output_map=["random"], size=6, auto_reset=True)
fetcher = DataFetcher(prefetch_batches=1)
fetcher.setup(loader)
for epoch in range(3):
    for i, _ in enumerate(fetcher):
        print(epoch, i)

Since the DALI iterator is sized, when we detect that it is exhausted, we signal to stop fetching by setting self.done: https://github.com/PyTorchLightning/pytorch-lightning/blob/eb21135b2aad20aaad45aae44858c090f1e780e5/pytorch_lightning/utilities/fetching.py#L285-L289. This means on the next __next__ call, we will raise the StopIteration ourselves: https://github.com/PyTorchLightning/pytorch-lightning/blob/eb21135b2aad20aaad45aae44858c090f1e780e5/pytorch_lightning/utilities/fetching.py#L275-L277
but the DALI iterator does not know this, and needs this extra __next__ call to reset its state: https://github.com/NVIDIA/DALI/blob/9d790c7cca5b383c99c79b430117e13f14764d3d/dali/python/nvidia/dali/plugin/base_iterator.py#L250-L251. This is why the first fetch of the next epoch raises the StopIteration.

IMO the best solution would be that (1) DALI does the same state reset after returning the last batch by checking the current counter vs total size or/and (2) we manually call a reset method on the DALI iterator when this happens.

I would appreciate input here from a DALI dev: cc @JanuszL as its top contributor

[JanuszL]

Hi @carmocca,

Thank you for your thorough analysis.
Regarding the options:

• (1) DALI does the same state reset after returning the last batch by checking the current counter vs total size - I'm thinking about extending auto-reset to have 3 values: True, False, Silent. The silent mode would not rise StopIteration and this is probably what we are looking for now. Let me check with the team if this is the way to go.
• (2) we manually call a reset method on the DALI iterator when this happens. - I guess this requires a dedicated code path for DALI (unless the user can provide a custom function for fetching)