Checkpointing and Early Stopping fail to work correctly when increasing number of train batches (in some cases)
🐛 Bug
( Preface: I created a complete minimal example for this bug report that unfortunately didn't end up reproducing the behavior, but I still think it might be useful to mention this nevertheless ).
The symptom is that when I leave everything else the same but increase the number of my training batches from 1000 to 5000, both checkpointing and early stopping completely fail to work correctly. As verified by creating a minimal example with a different simpler model it's not so much the number of batches but perhaps somehow related to the time it takes for an epoch to run, maybe. Here is a more detailed description:
Setup
In the LightningModule:
 def training_step(self, batch, _) -> Tensor:
        """ Perform a single step in the training loop """
        loss, nll = self.shared_step(batch, self.hparams.teacher_forcing)
        loss_with_reg = loss + self.reg(self.process)
        logs = {"loss_no_reg": loss, "loss_with_reg": loss_with_reg,
                    "nll": nll}

        self.log_dict(logs, on_epoch=True)
        return loss_with_reg

    def validation_step(self, batch: types.DataSplit, batch_idx) ->None:
        """ Perform an evaluation step """
        nll = torch.tensor(float(0)).to(batch.context.device)
        losses = [40, 20, 30, 10, 1, 0.9, 1, 1, 90, 100]
        loss = torch.tensor(float(losses[self.current_epoch])).to(batch.context.device)
        logs = {"val_loss": loss, "val_nll": nll}
        self.log_dict(logs)

By construction, epoch 5 should be the best model, and early stopping should trigger on epoch 8.
Experiment setup:
outroot = Path(args.out_dir)
logger = TestTubeLogger(save_dir=str(outroot / "logs"))

ckpt_filepath = "{}/{{epoch}}-{{val_loss:.2f}}".format(
        str(outroot / "logs" / "checkpoints"))
checkpoint_callback = ModelCheckpoint(
    filepath=ckpt_filepath, save_top_k=1, monitor="val_loss",
    verbose=True
)

early_stop = EarlyStopping(monitor="val_loss", verbose=True)

trainer = Trainer.from_argparse_args(
    args, logger=logger, checkpoint_callback=checkpoint_callback,
    early_stop_callback=early_stop
)

trainer.fit(model, datamodule=dm)

Behavior
Run 1 - All correct.
Okay, so with that, if I run with anywhere between 10 to a 1000 training batches, things work perfectly:
❯ python -m run.run_synthetic_social --gpus 1 ... --max_epochs 10  --limit_train_batches 10 --limit_val_batches 5
...
home/chirag/miniconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:37: UserWarning: The validation_epoch_end should not return anything as of 9.1.to log, use self.log(...) or self.write(...) directly in the LightningModule
  warnings.warn(*args, **kwargs)
Epoch 0:  73%|██████████████████████        | 11/15 [00:00<00:00, 12.14it/s, loss=443490368.000, v_num=0]Epoch 0: val_loss reached 40.00000 (best 40.00000), saving model to /home/chirag/Projects/mingle/social-processes/artefacts/exp/dev_run/logs/checkpoints/epoch=0-val_loss=40.00.ckpt as top 1
Epoch 1:  80%|████████████████████████      | 12/15 [00:01<00:00, 11.90it/s, loss=372233728.000, v_num=0]Epoch 1: val_loss reached 20.00000 (best 20.00000), saving model to /home/chirag/Projects/mingle/social-processes/artefacts/exp/dev_run/logs/checkpoints/epoch=1-val_loss=20.00.ckpt as top 1
Epoch 2:  80%|████████████████████████      | 12/15 [00:01<00:00, 10.89it/s, loss=197302992.000, v_num=0]Epoch 2: val_loss was not in top 1                                         | 1/5 [00:00<00:00,  7.89it/s]
Epoch 3:  80%|████████████████████████      | 12/15 [00:01<00:00, 11.46it/s, loss=108065304.000, v_num=0]Epoch 3: val_loss reached 10.00000 (best 10.00000), saving model to /home/chirag/Projects/mingle/social-processes/artefacts/exp/dev_run/logs/checkpoints/epoch=3-val_loss=10.00.ckpt as top 1
Epoch 4:  80%|████████████████████████      | 12/15 [00:00<00:00, 12.58it/s, loss=104392560.000, v_num=0]Epoch 4: val_loss reached 1.00000 (best 1.00000), saving model to /home/chirag/Projects/mingle/social-processes/artefacts/exp/dev_run/logs/checkpoints/epoch=4-val_loss=1.00.ckpt as top 1
Epoch 5:  80%|████████████████████████▊      | 12/15 [00:00<00:00, 13.81it/s, loss=64182412.000, v_num=0]Epoch 5: val_loss reached 0.90000 (best 0.90000), saving model to /home/chirag/Projects/mingle/social-processes/artefacts/exp/dev_run/logs/checkpoints/epoch=5-val_loss=0.90.ckpt as top 1
Epoch 6:  80%|████████████████████████▊      | 12/15 [00:01<00:00, 11.82it/s, loss=65260504.000, v_num=0]Epoch 6: val_loss was not in top 1                                         | 1/5 [00:00<00:00,  6.81it/s]
Epoch 7:  80%|████████████████████████      | 12/15 [00:01<00:00, 11.47it/s, loss=105555992.000, v_num=0]Epoch 7: val_loss was not in top 1                                         | 1/5 [00:00<00:00,  7.17it/s]
Epoch 8:  80%|████████████████████████      | 12/15 [00:01<00:00, 11.64it/s, loss=113607824.000, v_num=0]Epoch 8: val_loss was not in top 1                                         | 1/5 [00:00<00:00,  6.88it/s]
Epoch 8: 100%|██████████████████████████████| 15/15 [00:01<00:00, 13.69it/s, loss=113607824.000, v_num=0Epoch 00009: early stopping triggered.                                                                    
Epoch 8: 100%|██████████████████████████████| 15/15 [00:01<00:00, 13.17it/s, loss=113607824.000, v_num=0]

The checkpoint updates correctly on disk during training, and the last one is correctly named epoch=5-val_loss=0.90.ckpt.
Run 2 - Problematic.
If I increase the number of training batches to 5000, early stopping is triggered after the first 3 epochs, the checkpoints are not created live (only after early stopping has ended), and has the completely incorrect name epoch=3-val_loss=40.00.ckpt. Note that the epoch 3 val loss by construction should be  10, while it's picked up the epoch 0 loss.
❯ python -m run.run_synthetic_social --gpus 1 ... --max_epochs 10  --limit_train_batches 5000 --limit_val_batches 5
...
home/chirag/miniconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:37: UserWarning: The validation_epoch_end should not return anything as of 9.1.to log, use self.log(...) or self.write(...) directly in the LightningModule
  warnings.warn(*args, **kwargs)
Epoch 3:  20%|█████▉                       | 1026/5005 [01:22<05:18, 12.49it/s, loss=341889.875, v_num=0]Epoch 3: val_loss reached 40.00000 (best 40.00000), saving model to /home/chirag/Projects/mingle/social-processes/artefacts/exp/dev_run/logs/checkpoints/epoch=3-val_loss=40.00.ckpt as top 1
Epoch 00004: early stopping triggered.
Epoch 3:  20%|█████▉                       | 1026/5005 [01:22<05:19, 12.44it/s, loss=341889.875, v_num=0]

To Reproduce / Code Sample

I do have an isolated minimal code sample but unfortunately it works okay as expected, even with 15000 training batches. It's a much simpler model, since I'm using a variation of the LitModel, and I've tried to keep dimensions of tensors similar to my actual problem, so I don't know where the problem is right now. Here is the minimal code sample nevertheless:
https://gist.github.com/chiragraman/16b1a89787df0c517b8dfffae5c3d591
Expected behavior
The expected behavior in Run 2 above is to match the behavior in Run 1.
Environment

CUDA:
- GPU:
- Quadro P4000
- available:         True
- version:           10.2
Packages:
- numpy:             1.19.1
- pyTorch_debug:     False
- pyTorch_version:   1.6.0
- pytorch-lightning: 0.9.1rc4
- tqdm:              4.48.2
System:
- OS:                Linux
- architecture:
- 64bit
- ELF
- processor:         x86_64
- python:            3.8.5
- version:           #113-Ubuntu SMP Thu Jul 9 23:41:39 UTC 2020

Additional context
Also, minor side note: I'm getting a warning about UserWarning: The validation_epoch_end should not return anything as of 9.1.to log, use self.log(...) or self.write(...) directly in the LightningModule when I haven't implemented validation_epoch_end at all, and am not returning anything from validation_step