sync distributed metric error

[rogertrullo]

Hi all,
I am using lightning with multiple GPUs (DDP).
I am using self.log(sync_dist=True, on_epoch=True, prog_bar=True, logger=True) in the validation_step function.
I was trying to replicate the logged value, and I have found some issues.
To replicate the values I am doing this in the validation_epoch_end:

allgather=self.all_gather(outputs)
if self.trainer.global_rank==0:
    loss=[v['valid_loss'] for v in allgather]
    loss=torch.cat(loss).mean().item()
    print('expected loss', loss)

It turns out this only works when the dataloader has all batches of the same size , escentially it works fine if I put in the dataloader drop_last=True

Is this something expected or is it a bug?
For completeness here is the boring model modified a little bit:

import os
import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


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

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

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = 0*self(batch).sum()+self.trainer.global_rank+batch_idx
        self.log("valid_loss", loss, sync_dist=True, on_epoch=True, prog_bar=True, logger=True)
        return {"valid_loss": loss}
    def validation_epoch_end(self, outputs):
        allgather=self.all_gather(outputs)
        if self.trainer.global_rank==0:
            loss=[v['valid_loss'] for v in allgather]
            loss=torch.cat(loss).mean().item()
            print('expected 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 run():
    from pytorch_lightning.utilities.seed import  seed_everything
    from pytorch_lightning.callbacks import ModelCheckpoint

    seed_everything(42, workers=True)
    train_data = DataLoader(RandomDataset(32, 128), batch_size=4, num_workers=8, shuffle=True,pin_memory=True)
    val_data = DataLoader(RandomDataset(32, 82), batch_size=4, num_workers=8, shuffle=False,pin_memory=True, drop_last=False)
    
    checkpoint_callback = ModelCheckpoint(
        monitor='valid_loss',
        dirpath='checkpoints/',
        filename='system-{epoch:02d}-{valid_loss:.4f}'
        )
    model = BoringModel()
    trainer = Trainer(
        gpus=-1,
        strategy="ddp",
        callbacks=[checkpoint_callback],
        default_root_dir=os.getcwd(),
        num_sanity_val_steps=0,
        max_epochs=10,
        enable_model_summary=False
    )
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)

if __name__ == "__main__":
    run()