on_validation_batch_end is not called when Loss is NaN

[PeiyuanQi]

Bug description

Problem: on_validation_batch_end is not called when Loss is NaN.

How to reproduce: create a model with loss as NaN invalidation, and then on_validation_batch_end won't be called.

What version are you seeing the problem on?

v2.5

Reproduced in studio

No response

How to reproduce the bug

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- 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

[deependujha]

Hi @PeiyuanQi ,
Thank you for opening the issue, I'm currently looking into it.

As a suggestion, it would be very helpful if you could share a lightning studio that reproduces the issue. This would greatly assist us in triaging and resolving it more efficiently.

Thanks again!
Deependu

[deependujha]

Hi @PeiyuanQi,

I ran a test with a model that produces NaN during validation, and I’m seeing that on_validation_batch_end is being called as expected (screenshot attached).

Here's my code

import math
import argparse

import lightning as L
import torch
import torch.nn as nn
from lightning.pytorch.callbacks import ModelCheckpoint, RichProgressBar
from torch.utils.data import DataLoader, TensorDataset


callbacks = [
    RichProgressBar(),
]


class SimpleModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 1)
        self.loss_fn = nn.MSELoss()

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

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.loss_fn(y_hat, y)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.loss_fn(y_hat, y)
        self.log("val_loss", loss)
        return torch.tensor(math.nan)

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


def create_random_data(num_samples=1000):
    x = torch.randn(num_samples, 10)
    y = torch.randn(num_samples, 1)
    return x, y


if __name__ == "__main__":
    x_train, y_train = create_random_data(1000)
    x_val, y_val = create_random_data(200)

    train_dataset = TensorDataset(x_train, y_train)
    val_dataset = TensorDataset(x_val, y_val)

    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=1, persistent_workers=True)
    val_loader = DataLoader(val_dataset, batch_size=32, num_workers=1, persistent_workers=True)

    model = SimpleModel()

    trainer = L.Trainer(
        max_epochs=3,
        enable_progress_bar=True,
        log_every_n_steps=1,
        val_check_interval=0.5,
        callbacks=callbacks,
        devices=[0],
    )

    trainer.fit(model, train_loader, val_loader)

It’s possible there might be a difference in how the NaN is being introduced or handled in your setup. If you could share a minimal Lightning Studio that reproduces the issue, I’d be happy to investigate further.

Best,
Deependu

[deependujha]

Hi @PeiyuanQi , just checking in to see if you had a chance to create a minimal reproducible example for this issue. That would help us investigate further.

[PeiyuanQi]

ah sry, just realize we didn't reply to this thread, based on further investigation we think the problem is on the app implication side, we could close this thread.