update dataset at "on_train_epoch_start", but "training_step" still get old data

[Yak1m4Sg]

Bug description

I use trainer.fit(model, datamodule=dm) to start training.
"dm" is an object whose class inherited from pl.LightningDataModule, and in the class, I override the function:

def train_dataloader(self):
    train_dataset = MixedBatchMultiviewDataset(self.args, self.tokenizer, 
         known_exs=self.known_train, 
         unknown_exs=self.unknown_train, 
         feature=self.args.feature) 
        
    train_dataloader = DataLoader(train_dataset, 
         batch_size = self.args.train_batch_size, 
         shuffle=True, num_workers=self.args.num_workers, 
         pin_memory=True, collate_fn=self.collate_batch_feat)
        
    return train_dataloader

at the model's hook on_train_epoch_start, I update the dataset:

train_dl = self.trainer.train_dataloader
train_dl.dataset.update_pseudo_labels(uid2pl)

loop = self.trainer.fit_loop
loop._combined_loader = None
loop.setup_data()

in the training_step, the batch data is still old data, but trainer.train_dataloader.dataset is new:

def training_step(self, batch: List[Dict[str, torch.Tensor]], batch_idx: int):

    self.mv_model._on_train_batch_start()

    logger.info(self.trainer.train_dataloader.dataset.unknown_feats) # new
    logger.info(batch) # old

What version are you seeing the problem on?

v2.3

How to reproduce the bug

No response

Error messages and logs

# Error messages and logs here please

Environment

Current environment

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

cc @justusschock

[lantiga]

hey, thanks for filing the issue. Can you provide a minimal runnable reproduction I can just copy-paste? that'll speed up the investigation, thanks!

[SkafteNicki]

A tried to reproduce the bug but was unable to do it on master. Maybe bug have been fixed at some point this issue was posted. Posting script here:

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


class DynamicDataset(Dataset):
    def __init__(self, initial_data):
        self.data = initial_data.copy()
        self.version = 1
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return {
            "value": torch.tensor(self.data[idx], dtype=torch.float32),
            "version": self.version,
            "idx": idx
        }
    
    def update_pseudo_labels(self, new_data):
        """Update pseudo labels - reproduces the exact pattern from the issue."""
        print(f"📝 Dataset update: {self.data} → {new_data} (v{self.version} → v{self.version+1})")
        self.data = new_data.copy()
        self.version += 1


class TestDataModule(pl.LightningDataModule):
    def __init__(self):
        super().__init__()
        
    def setup(self, stage=None):
        if stage == "fit" or stage is None:
            self.train_dataset = DynamicDataset([100, 200, 300, 400])
        
    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=2, shuffle=False, num_workers=0)


class TestModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)
        self.epoch = 0
        
    def on_train_epoch_start(self):
        self.epoch += 1
        print(f"\n🔄 EPOCH {self.epoch}: on_train_epoch_start()")
        
        # Reproduce the exact pattern from the issue report
        train_dl = self.trainer.train_dataloader
        print(f"   Before: train_dl.dataset data = {train_dl.dataset.data}")
        
        # Update with new pseudo labels
        uid2pl = [1000 + self.epoch, 2000 + self.epoch, 3000 + self.epoch, 4000 + self.epoch]
        train_dl.dataset.update_pseudo_labels(uid2pl)
        
        print(f"   After:  train_dl.dataset data = {train_dl.dataset.data}")
        
    def training_step(self, batch, batch_idx):
        print(f"   training_step(batch_idx={batch_idx}):")
        
        # Check the exact scenario described in the issue
        dataset_data = self.trainer.train_dataloader.dataset.data  # Should be NEW
        dataset_version = self.trainer.train_dataloader.dataset.version
        
        batch_values = batch["value"].cpu().tolist()  # Was this OLD in the issue?
        batch_versions = batch["version"].cpu().tolist()
        batch_indices = batch["idx"].cpu().tolist()
        
        # What should the batch contain based on current dataset?
        expected_values = [dataset_data[i] for i in batch_indices]
        
        print(f"     self.trainer.train_dataloader.dataset data: {expected_values} (NEW)")
        print(f"     batch data: {batch_values} (should match above)")
        
        # Detect the issue described in #20407
        if batch_values != expected_values or any(v != dataset_version for v in batch_versions):
            print("     ❌ BUG REPRODUCED: training_step gets OLD data, dataset shows NEW data!")
            return torch.tensor(0.0, requires_grad=True)
        else:
            print("     ✅ WORKING CORRECTLY: Batch data matches updated dataset")
        
        # Simple training step
        x = batch["value"].view(-1, 1)
        y = self.linear(x)
        loss = nn.functional.mse_loss(y, torch.zeros_like(y))
        return loss
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.01)


def main():
    print("🔍 Verifying PyTorch Lightning issue #20407")
    print(f"📦 PyTorch Lightning version: {pl.__version__}")
    print("📄 Issue: Dataset updated in on_train_epoch_start but training_step gets old data\n")
    
    # Setup
    model = TestModel()
    datamodule = TestDataModule()
    
    # Train with minimal configuration
    trainer = pl.Trainer(
        max_epochs=2,
        limit_train_batches=2,
        enable_checkpointing=False,
        logger=False,
        enable_progress_bar=False,
        accelerator="cpu"
    )
    
    trainer.fit(model, datamodule=datamodule)

if __name__ == "__main__":
    main()

[SkafteNicki]

Closing issue as it seems fixed. Feel free to ping me and reopen if necessary and reproducible example can be provided.