Bug in ThroughputMonitor.on_validation_end: using sum() instead of last value can corrupt

[itzhakstern]

Bug description

When using the built-in ThroughputMonitor callback, training can crash with:
ValueError: Expected the value to increase, last: X, current: X

This happens because in on_validation_end Lightning currently computes training_finished as:
training_finished = self._t0s[RunningStage.TRAINING] + sum(self._throughputs[RunningStage.TRAINING]._time)

But _time already stores cumulative elapsed times per batch (monotonic, increasing).
Summing them again produces a huge, incorrect number, which eventually makes the internal _time.append assertion fail (Expected the value to increase).

What version are you seeing the problem on?

v2.5

Reproduced in studio

No response

How to reproduce the bug

import os
from torch import optim, nn, utils
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
import lightning as L
from torch.utils.data import random_split
from lightning.pytorch.callbacks import ThroughputMonitor
from lightning.pytorch.trainer.trainer import Trainer

encoder = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))

class CustomThroughputMonitor(ThroughputMonitor):
    def __init__(self, batch_size_fn=None, *args, **kwargs) -> None:
        super().__init__(batch_size_fn=batch_size_fn, *args, **kwargs)

    def setup(self, trainer: Trainer, pl_module, stage):
        self.batch_size_fn = lambda batch: 1
        return super().setup(trainer, pl_module, stage)

class LitAutoEncoder(L.LightningModule):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def training_step(self, batch, batch_idx):
        x, _ = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = nn.functional.mse_loss(x_hat, x)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, _ = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        val_loss = nn.functional.mse_loss(x_hat, x)
        self.log("val_loss", val_loss)

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=1e-3)

autoencoder = LitAutoEncoder(encoder, decoder)
dataset = MNIST(os.getcwd(), download=True, transform=ToTensor())
train_set, val_set = random_split(dataset, [55000, 5000])
train_loader = utils.data.DataLoader(train_set, batch_size=64)
val_loader = utils.data.DataLoader(val_set, batch_size=64)

trainer = L.Trainer(
    limit_train_batches=500,
    max_epochs=20,
    callbacks=[CustomThroughputMonitor()],
    log_every_n_steps=100,
)

trainer.fit(model=autoencoder, train_dataloaders=train_loader, val_dataloaders=val_loader)

Error messages and logs

raise ValueError(f"Expected the value to increase, last: {last}, current: {x}")
ValueError: Expected the value to increase, last: 358569066962426.94, current: 358569066962426.94

Environment

No response

More info

Proposed fix:
Replace the sum() with just the last elapsed value

    @rank_zero_only
    def on_validation_end(self, trainer, *_):
        if trainer.sanity_checking or trainer.state.fn != TrainerFn.FITTING:
            return

        training_finished = self._t0s[RunningStage.TRAINING] + self._throughputs[RunningStage.TRAINING]._time[-1]
        time_between_train_and_val = self._t0s[RunningStage.VALIDATING] - training_finished
        val_time = self._throughputs[RunningStage.VALIDATING]._time[-1]
        self._t0s[RunningStage.TRAINING] += time_between_train_and_val + val_time

[dfhssilva]

I was also investigating this bug independently. I am on V2.3.3. I initially found it since I was noticing that the train/device/mfu curves I was getting with the ThroughputMonitor had this sinking pattern for Throughput.window_size steps at the beginning of each subsequent training epoch:

I tried @itzhakstern proposed solution and the sinking is still there but it's much less steep (instead of dropping from ~0.16 to ~0 it goes to ~0.13):

My understanding is that this is happening because how MFU is calculated in Throughput.compute():

elapsed_flops = sum(self._flops) - self._flops[0]
elapsed_time = self._time[-1] - self._time[0]
flops_per_sec = elapsed_flops / elapsed_time
dev_flops_per_sec = flops_per_sec / self.world_size
metrics[f"device{self.separator}mfu"] = (
    dev_flops_per_sec / self.available_flops
            )

It's possible to see that MFU is affected by the value of elapsed_time which in the first window_size batches of epochs > 0, will have an inflated value (extremely with the existing bug and slightly with the proposed fix). After window_size steps, the self._time queue resets and the elapsed_time values will return to the original scale which is why the MFU returns back to the original scale too.

Regarding the proposed fix, it does improve the situation but it still creates the sinking in the MFU for the initial steps of subsequent epochs, which can be misleading and confusing for someone that's trying to assess the MFU values. We should be adding the exact time between the 2 epochs to the self._t0s[RunningStage.TRAINING], which doesn't correspond only to time_between_train_and_val + val_time.

Instead, I would propose the following solution - adding a on_train_epoch_start hook for incrementing the t0 of training with a more accurate time between the new and the last epoch (this requires the deletion of the on_validation_end hook):

@override
@rank_zero_only
def on_train_epoch_start(self, trainer, pl_module):
    # increment t0s at each epoch to avoid sinking the throughput due to validation or other stuff happening
    stage = trainer.state.stage
    assert stage is not None
    if len(self._throughputs[stage]._time) > 0:
        last_epoch_training_finished = self._t0s[stage] + self._throughputs[stage]._time[-1]
        time_between_epochs = time.perf_counter() - last_epoch_training_finished
        self._t0s[stage] += time_between_epochs

The corresponding MFU curve below still shows some sinking patterns, however the difference is now much less abrupt now (from 0.1654 to 0.1648 from the 1st to the 2nd epoch):

[SkafteNicki]

@itzhakstern thanks for raising this issue. It does indeed seem like there is a bug in the code.
@dfhssilva just to confirm does your "fix" actually solve the orignal valueerror mentioned at the top of the issue or do think think what you are observing is a separate issue?

[radulescupetru]

@SkafteNicki Hijacking this issue cause it's related, but I think the way that Throughput is defined might not be correct for cases when using advanced parallelism methods like TP + FSDP.
The current approach uses trainer.world_size to scale the values, but it should instead compute the value depending if pl.module.device_mesh is None or not.

Could look something like:

def world_size(trainer: Trainer) -> int:
    """Get distributed training world size information.

    Retrieves rank and world size either from device mesh if available,
    or falls back to trainer's global values.

    Args:
        trainer: Lightning Trainer instance

    Returns:
        int: world_size
    """
    assert trainer.lightning_module is not None, "Expected trainer.lightning_module to be setup."
    pl_module = trainer.lightning_module
    if (
        pl_module.device_mesh is not None
        and pl_module.device_mesh.mesh_dim_names is not None
        and "data_parallel" in pl_module.device_mesh.mesh_dim_names
    ):
        dp_mesh = pl_module.device_mesh["data_parallel"]
        return dp_mesh.size()
    return trainer.world_size

[dfhssilva]

@dfhssilva just to confirm does your "fix" actually solve the orignal valueerror mentioned at the top of the issue or do think think what you are observing is a separate issue?

I don't know since I didn't experience the same valueerror above. When I was digging in the code I spotted this same bug with the sum of the elapsed times (the calculation is wrong because we are summing cumulative values) and that's where I noticed they were causing the drop in the MFU at the beginning of each epoch.
It seems that the same bug is causing these two separate issues (the valueerror and the drop of MFU after each epoch).
However, I believe the solution I am proposing it's just a more accurate way to append _t0s because we are doing it at the epoch start instead of at validation end.

[itzhakstern]

@SkafteNicki The solution that @dfhssilva proposed removes the unnecessary calculation performed in the on_validation_end function, where the validation time is subtracted from the throughput calculation.
This change should, in any case, resolve the overflow issue caused by the incorrect use of sum() in this function.

Currently, there are actually two separate bugs in the on_validation_end function:

1. The use of sum() to compute the training time for the current epoch, while it is sufficient (and correct) to use only the last value of the _time array.
2. The _time array is not properly reset between epochs, leading to cumulative errors over time.

My focus was on addressing the first bug, which caused the overflow error I encountered.
However, the solution and update in the on_train_epoch_start function is intended to addresses both issues — the incorrect summation and the accumulation — and should therefore fully resolve both the overflow and the throughput inaccuracy problems - but of course this needs to be tested in a more qualitative way.