on_validation_batch_end() return empty list in 'dp' mode

[allanchan339]

🐛 Bug

To Reproduce

The following tutorial is used and cloned to my machine to test the functionality of Distributed Data-Parallel and Data Parallel.

https://colab.research.google.com/github/wandb/examples/blob/master/colabs/pytorch-lightning/Optimize_Pytorch_Lightning_models_with_Weights_%26_Biases.ipynb

The validation step will return preds for a customized callback LogPredictionsCallback() that saves images.
The results between DDP and DP are completely different. In DDP mode, the sanity checking pass.
However, in DP mode, the "outputs" in LogPredictionsCallback() are just an empty list and fail the sanity checking.
Where the error codes like :

File "/home/user/Desktop/Code/ALEN/lightning_log.py", line 151, in on_validation_batch_end
    for y_i, y_pred in list(zip(y[:n], outputs[:n]))]
IndexError: dimension specified as 0 but tensor has no dimensions

As 2 GPU is required, the bug cannot reproduce in the Colab environment as Colab will not provide multiple free GPUs.

Expected behavior

The preds from validation_step() should be accumulated and can be called "outputs" in Callback, just like DDP mode.

Environment

• CUDA:
  - GPU:
  - NVIDIA GeForce RTX 3090
  - NVIDIA GeForce RTX 3090
  - available: True
  - version: 11.2
• Packages:
  - numpy: 1.21.5
  - pyTorch_debug: False
  - pyTorch_version: 1.10.0
  - pytorch-lightning: 1.5.8
  - tqdm: 4.62.3
• System:
  - OS: Linux
  - architecture:
  - 64bit
  - ELF
  - processor: x86_64
  - python: 3.9.7
  - version: #202201071026-Ubuntu SMP Fri Jan 7 16:52:09 UTC 2022

Additional context

cc @justusschock @awaelchli @akihironitta @rohitgr7

[morganmcg1]

Hey @allanchan339 I'm from Weights & Biases, we'll look into this, thanks for raising it :)

[manangoel99]

Hello! recreated this with DataParallel. For 2 GPUs, DataParallel splits the batch (size 256) into two parts (size 128 each) and runs computation on them seperately. This happens in the validation_step function as well but when the on_validation_batch_end function is called, the batch variable has 256 values while outputs turns out to be a scalar.

The underlying reason could be that after the validation_step function, the outputs from the 2 GPUs aren't put back together before sending them to on_validation_batch_end. It's most likely an internal bug in pytorch_lightning.

The code below can be used to recreate this.

Expected Output

torch.Size([256, 1, 28, 28]) torch.Size([256])

Actual Output

torch.Size([256, 1, 28, 28]) torch.Size([ ])

from pytorch_lightning import LightningDataModule
from torch.utils.data import DataLoader, random_split
from torchvision.datasets import MNIST
from torchvision import transforms
from pytorch_lightning.callbacks import Callback
from pytorch_lightning import Trainer
import torch
from torch.nn import Linear, CrossEntropyLoss, functional as F
from torch.optim import Adam
from torchmetrics.functional import accuracy
from pytorch_lightning import LightningModule
from pytorch_lightning import seed_everything

seed_everything(42)

# Data Module

class MNIST_DataModule(LightningDataModule):

    def __init__(self, data_dir='./', batch_size=256):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))])

    def prepare_data(self):
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage=None):
        if stage in [None, 'fit', 'validate']:
            mnist_train = MNIST(self.data_dir, train=True, transform=self.transform)
            self.mnist_train, self.mnist_val = random_split(mnist_train, [55000, 5000])
        if stage in [None, 'test']:
            self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)

    def train_dataloader(self):
        mnist_train = DataLoader(self.mnist_train, batch_size=self.batch_size)
        return mnist_train

    def val_dataloader(self):
        mnist_val = DataLoader(self.mnist_val, batch_size=self.batch_size)
        return mnist_val

    def test_dataloader(self):
        mnist_test = DataLoader(self.mnist_test, batch_size=self.batch_size)
        return mnist_test

# Training Module

class MNIST_LitModule(LightningModule):

    def __init__(self, n_classes=10, n_layer_1=128, lr=1e-3):
        super().__init__()
        self.layer_1 = Linear(28 * 28, n_layer_1)
        self.layer_2 = Linear(n_layer_1, n_classes)
        self.loss = CrossEntropyLoss()
        self.lr = lr
        self.save_hyperparameters()

    def forward(self, x):
        batch_size, channels, width, height = x.size()
        x = x.view(batch_size, -1)
        x = self.layer_1(x)
        x = F.relu(x)
        x = self.layer_2(x)

        return x

    def training_step(self, batch, batch_idx):
        _, loss, acc = self._get_preds_loss_accuracy(batch)
        return loss

    def validation_step(self, batch, batch_idx):
        preds, loss, acc = self._get_preds_loss_accuracy(batch)
        return preds

    def test_step(self, batch, batch_idx):
        _, loss, acc = self._get_preds_loss_accuracy(batch)
    
    def configure_optimizers(self):
        return Adam(self.parameters(), lr=self.lr)
    
    def _get_preds_loss_accuracy(self, batch):
        x, y = batch
        logits = self(x)
        preds = torch.argmax(logits, dim=1)
        loss = self.loss(logits, y)
        acc = accuracy(preds, y)
        return preds, loss, acc

mnist = MNIST_DataModule()
model = MNIST_LitModule(n_layer_1=128)

# Callback Function Which Has The Error

class LogPredictionsCallback(Callback):   
    def on_validation_batch_end(
        self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
        print(batch[0].shape, outputs.shape)

trainer = Trainer(
    callbacks=[LogPredictionsCallback()],   # logging of sample predictions
    gpus=-1,                                # use all available GPU's
    max_epochs=5,                           # number of epochs
    precision=16,                           # train in half precision
    deterministic=True, strategy='dp')                     # for reproducibility

trainer.fit(model, datamodule=mnist)

[rohitgr7]

must be happening due to this:
https://github.com/PyTorchLightning/pytorch-lightning/blob/1a2536361749d13e3f30cec11c7da11747f425ab/pytorch_lightning/strategies/dp.py#L139-L152

it reduces (basically takes the mean by default) for everything.
two possible solutions I can think of:

1. Either override specific_step_end hook to avoid the reductions.
2. Lightning shouldn't reduce everything by default. Just loss should be reduced.

[justusschock]

@rohitgr7 You are right with the reason for this. However, maybe we should also check the callbacks whether they implement it? Because I agree that the current behaviour is weird.

[rohitgr7]

However, maybe we should also check the callbacks whether they implement it?

which hooks in callbacks? @justusschock

[justusschock]

Every hook getting the outputs (on_validation_batch_end etc. )

[rohitgr7]

Every hook getting the outputs (on_validation_batch_end etc. )

hmm.. but they cannot change the outputs and return them but LM.validation_step_end can and return them.
The reason, I think, we check this is because we expect users to reduce the loss too if they are overriding LM.validation_step_end.

Anyways if we move forward to remove the reduction, I think it should not reduce anything else except the loss.

[justusschock]

Well they could still change it without returning (python uses call by reference most of the time).

However, we may also keep it the current way