Validation has an impact on training

Bug description

I have an inconsistency when training models. The training changes depending on the validation. I don't mean it changes because of the LR scheduler being fed with validation results, or early stopping, but simply doing (or not) the validation at the end of each epoch changes the loss of the model afterwards.

What I know:

• It only happens when training in multi-gpu mode with DDP
• If I change the validation batch size the training is different
• If I don't go through validation at the end of each epoch, the training loss is different too
• A larger batch size seems to provide results closer to the ones I get without validation
• However if I don't change the validation process (bsz, or no validation), the results are consistent between each other
• The model itself (state_dict) is exactly the same before and after validation
• Before the first validation, the training losses are always exactly the same between the models. It's always after that first validation that the training losses start to diverge
• I tried the same experiments without a HuggingFace transformer model (just a LinearLayer instead, or a torchaudio.models.Wav2Vec2Model), and I didn't noticed any difference. At first, this led me thinking the issue was from HuggingFace, but I'm not able to reproduce the bug without Pytorch Lightning.

On the screen below you can find 6 different runs using a MWE.

• models ending by debug128 are validated using a valid batch size of 128
• models ending by debug8 are validated using a valid batch size of 8
• models ending with debugX-novalid are not validated, but the validation dataloader has been created with a batch size of X
  

How to reproduce the bug

Here is a example code that reproduce this bug. I used a first script to generate dummy random datasets. I used Wav2Vec2ForSequenceClassification as the HuggingFace model, but I had the issue with other HF models.

import torch

NUM_LABELS = 100

def generate_dummy_set(num_samples):
    input_lengths = torch.randint(low=16_000, high=4 * 16_000, size=(num_samples,)).tolist()
    return [
        {
            'input': torch.rand((input_lengths[sample],)),
            'label': torch.randint(low=0, high=NUM_LABELS, size=(1,))
        } for sample in range(num_samples)
    ],

torch.save(generate_dummy_set(1000), "random_train_1000.pt")
torch.save(generate_dummy_set(600), "random_valid_600.pt")

#!/usr/bin/env python3

import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--valid_bsz', required=True, type=int)
parser.add_argument('--do_valid', action='store_true', default=False)
parser.add_argument('--model_version', required=True, type=str)
parser.add_argument('--wav2vec2_pretrained', default='facebook/wav2vec2-base', type=str)

args = parser.parse_args()

NUM_LABELS = 100

BSZ_TRAIN = 64
BSZ_VALID = args.valid_bsz

################################################################################

import os
import sys
import logging
import tempfile

import torch
import pytorch_lightning as pl
import torch.distributed as dist

logging.basicConfig(
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=os.environ.get("LOGLEVEL", "INFO").upper(),
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger()

################################################################################

# bypass slurm related mechanics of PL
for var in list(os.environ.keys()):
    if var.startswith("SLURM_"):
        del os.environ[var]

def get_trainer():
    return pl.Trainer(
        callbacks=[],
        logger=pl.loggers.TensorBoardLogger(
            save_dir='tensorboard/',
            version=args.model_version,
            name='debug',
        ),
        log_every_n_steps=1,
        strategy='ddp',
        accelerator='gpu',
        devices=-1,
        accumulate_grad_batches=4, # same issue with no accumulate grad
    )

class DataModule(pl.LightningDataModule):

    def __init__(self):
        super().__init__()

    def setup(self, stage):
        self.num_workers = 10 # 0 produces the same issue
        self.dummy_train = torch.load("random_train_1000.pt")
        self.dummy_val = torch.load("random_valid_600.pt")

    def train_dataloader(self):
        return torch.utils.data.DataLoader(
            self.dummy_train,
            shuffle=True,
            batch_size=BSZ_TRAIN,
            num_workers=self.num_workers,
            collate_fn=self.collate_fn
        )

    def val_dataloader(self):
        return torch.utils.data.DataLoader(
            self.dummy_val,
            shuffle=False,
            batch_size=BSZ_VALID,
            num_workers=self.num_workers,
            collate_fn=self.collate_fn
        )

    def collate_fn(self, samples):
        return {
            'input': torch.nn.utils.rnn.pad_sequence(
                [samples[i]['input'] for i in range(len(samples))],
                batch_first=True, padding_value=0
            ),
            'label': torch.nn.utils.rnn.pad_sequence(
                [samples[i]['label'] for i in range(len(samples))],
                batch_first=True, padding_value=0
            ),
            'bsz': len(samples)
        }

class HFModel(pl.LightningModule):

    def __init__(self):
        super().__init__()
        from transformers import Wav2Vec2ForSequenceClassification
        self.model = Wav2Vec2ForSequenceClassification.from_pretrained(
            args.wav2vec2_pretrained,
            num_labels=NUM_LABELS,
            gradient_checkpointing=False
        )

    def forward(self, batch, labels=None):
        return self.model(input_values=batch['input'], labels=labels)

    def training_step(self, batch, batch_idx):
        outputs = self(batch, labels=batch['label'][:,0])
        loss = outputs.loss
        self.log('train__loss', loss, batch_size=batch['bsz'])
        return loss

    def validation_step(self, batch, batch_idx):
        if args.do_valid:
            outputs = self(batch)

    def configure_optimizers(self):
        return {
            'optimizer': torch.optim.AdamW(self.parameters(), lr=0.02, weight_decay=0.0)
        }

def main():
    pl.seed_everything(42, workers=True)
    trainer = get_trainer()

    data = DataModule()
    model = HFModel()

    trainer.fit(
        model=model,
        datamodule=data,
    )

if __name__ ==  '__main__':
    main()

Error messages and logs

Tensorboard results can be found here: https://tensorboard.dev/experiment/gKs1b3ZuQTWnEpt4wlGdUg/#scalars&_smoothingWeight=0
Not all models were stopped at the same step, I manually stopped them when I noticed they diverged.

Environment

* CUDA:
        - GPU:               None
        - available:         False
        - version:           11.6
* Lightning:
        - lightning-utilities: 0.7.0
        - pytorch-lightning: 1.9.2
        - torch:             1.13.1+cu116
        - torchaudio:        0.13.1+cu116
        - torchmetrics:      0.11.1
        - torchvision:       0.14.1+cu116
* Packages:
        - aiofiles:          22.1.0
        - aiohttp:           3.8.4
        - aiosignal:         1.3.1
        - aiosqlite:         0.18.0
        - alembic:           1.9.4
        - anyio:             3.6.2
        - argon2-cffi:       21.3.0
        - argon2-cffi-bindings: 21.2.0
        - arrow:             1.2.3
        - asttokens:         2.2.1
        - async-generator:   1.10
        - async-timeout:     4.0.2
        - attrs:             22.2.0
        - babel:             2.11.0
        - backcall:          0.2.0
        - beautifulsoup4:    4.11.2
        - bleach:            6.0.0
        - certifi:           2022.12.7
        - certipy:           0.1.3
        - cffi:              1.15.1
        - charset-normalizer: 3.0.1
        - comm:              0.1.2
        - cryptography:      39.0.1
        - datasets:          2.9.0
        - debugpy:           1.6.6
        - decorator:         5.1.1
        - defusedxml:        0.7.1
        - dill:              0.3.6
        - executing:         1.2.0
        - fastjsonschema:    2.16.2
        - filelock:          3.9.0
        - fqdn:              1.5.1
        - frozenlist:        1.3.3
        - fsspec:            2023.1.0
        - greenlet:          2.0.2
        - huggingface-hub:   0.12.1
        - idna:              3.4
        - ipdb:              0.13.11
        - ipykernel:         6.21.2
        - ipython:           8.10.0
        - ipython-genutils:  0.2.0
        - isoduration:       20.11.0
        - jedi:              0.18.2
        - jinja2:            3.1.2
        - json5:             0.9.11
        - jsonpointer:       2.3
        - jsonschema:        4.17.3
        - jupyter-client:    8.0.3
        - jupyter-core:      5.2.0
        - jupyter-events:    0.6.3
        - jupyter-server:    2.3.0
        - jupyter-server-fileid: 0.7.0
        - jupyter-server-terminals: 0.4.4
        - jupyter-server-ydoc: 0.6.1
        - jupyter-telemetry: 0.1.0
        - jupyter-ydoc:      0.2.2
        - jupyterhub:        3.1.1
        - jupyterlab:        3.6.1
        - jupyterlab-pygments: 0.2.2
        - jupyterlab-server: 2.19.0
        - lightning-utilities: 0.7.0
        - mako:              1.2.4
        - markupsafe:        2.1.2
        - matplotlib-inline: 0.1.6
        - mistune:           2.0.5
        - multidict:         6.0.4
        - multiprocess:      0.70.14
        - nbclassic:         0.5.2
        - nbclient:          0.7.2
        - nbconvert:         7.2.9
        - nbformat:          5.7.3
        - nest-asyncio:      1.5.6
        - notebook:          6.5.2
        - notebook-shim:     0.2.2
        - numpy:             1.24.2
        - oauthlib:          3.2.2
        - packaging:         23.0
        - pamela:            1.0.0
        - pandas:            1.5.3
        - pandocfilters:     1.5.0
        - parso:             0.8.3
        - pexpect:           4.8.0
        - pickleshare:       0.7.5
        - pillow:            9.4.0
        - pip:               23.0.1
        - platformdirs:      3.0.0
        - prometheus-client: 0.16.0
        - prompt-toolkit:    3.0.36
        - protobuf:          3.20.3
        - psutil:            5.9.4
        - ptyprocess:        0.7.0
        - pure-eval:         0.2.2
        - pyarrow:           11.0.0
        - pycparser:         2.21
        - pygments:          2.14.0
        - pyopenssl:         23.0.0
        - pyrsistent:        0.19.3
        - python-dateutil:   2.8.2
        - python-json-logger: 2.0.6
        - pytorch-lightning: 1.9.2
        - pytz:              2022.7.1
        - pyyaml:            6.0
        - pyzmq:             25.0.0
        - regex:             2022.10.31
        - requests:          2.28.2
        - responses:         0.18.0
        - rfc3339-validator: 0.1.4
        - rfc3986-validator: 0.1.1
        - ruamel.yaml:       0.17.21
        - ruamel.yaml.clib:  0.2.7
        - send2trash:        1.8.0
        - setuptools:        67.3.2
        - six:               1.16.0
        - sniffio:           1.3.0
        - soupsieve:         2.4
        - sqlalchemy:        2.0.4
        - stack-data:        0.6.2
        - tensorboardx:      2.6
        - terminado:         0.17.1
        - tinycss2:          1.2.1
        - tokenizers:        0.13.2
        - tomli:             2.0.1
        - torch:             1.13.1+cu116
        - torchaudio:        0.13.1+cu116
        - torchmetrics:      0.11.1
        - torchvision:       0.14.1+cu116
        - tornado:           6.2
        - tqdm:              4.64.1
        - traitlets:         5.9.0
        - transformers:      4.26.1
        - typing-extensions: 4.5.0
        - uri-template:      1.2.0
        - urllib3:           1.26.14
        - wcwidth:           0.2.6
        - webcolors:         1.12
        - webencodings:      0.5.1
        - websocket-client:  1.5.1
        - wheel:             0.38.4
        - xxhash:            3.2.0
        - y-py:              0.5.9
        - yarl:              1.8.2
        - ypy-websocket:     0.8.2
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                - ELF
        - processor:         x86_64
        - python:            3.10.9
        - version:           #1 SMP Thu Nov 17 16:37:07 EST 2022

More info

No response

cc @awaelchli

Start with the provided MWE, especially get_trainer, train_dataloader, val_dataloader, and validation_step, and compare DDP runs with validation enabled, disabled, and different validation batch sizes. Check when training losses first diverge after validation and verify the model state before and after validation. Done means the cause is identified and the inconsistent training behavior is reproduced by a regression test or otherwise corrected.