on_step logging not working as expected/described
🐛 Bug
When training a model with the MLFlowLogger, on_step logging in training_step() does not appear to log metrics as frequently as expected.  See complete example below.
To Reproduce
Code sample
Here is a complete working example that generates the described behavior.  This example is derived from the code in Lightning in 2 steps.  The only changes are adding the logging details, setting a mini-batch size, and subsetting the MNIST dataset to decrease running time.
import os
import torch
import torch.nn.functional as F
from torchvision.datasets import MNIST
from torchvision import transforms
from torch.utils.data import DataLoader, Subset
import pytorch_lightning as pl
from torch.utils.data import random_split
import numpy as np

class LitModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer_1 = torch.nn.Linear(28 * 28, 128)
        self.layer_2 = torch.nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.layer_1(x)
        x = F.relu(x)
        x = self.layer_2(x)
        return x

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        result = pl.TrainResult(loss)
        result.log('train_loss', loss, on_step=True)
        return result

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        result = pl.EvalResult(checkpoint_on=loss)
        result.log('val_loss', loss)
        return result


# dataloaders
dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
data_subset = Subset(dataset, np.random.choice(len(dataset), 1000, replace=False))
train, val = random_split(data_subset, [700, 300])
train_loader = DataLoader(train, batch_size=4)
val_loader = DataLoader(val)

# init model
model = LitModel()

logger = pl.loggers.mlflow.MLFlowLogger()

trainer = pl.Trainer(logger=logger, max_epochs=20)
trainer.fit(model, train_loader, val_loader)

Expected behavior
Training for 20 epochs with a training dataset of 700 samples and a mini-batch size of 4, we would expect training_step() to run 3,500 times ((700 / 4) * 20).  Assuming I correctly understand the semantics of logging on_step, I would also expect the training loss to be logged 3,500 times.  However, when I inspect the results in the MLflow UI or directly examine the logged values in the metrics/ folder within mlruns/, only 60 loss values are logged for the entirety of the run.  I've confirmed that training_step() does, in fact, get called 3,500 times over the course of the run.
Environment

CUDA:

GPU:
available:         False
version:           10.2


Packages:

numpy:             1.19.2
pyTorch_debug:     False
pyTorch_version:   1.6.0
pytorch-lightning: 0.9.0
tqdm:              4.50.0


System:

OS:                Linux
architecture:

64bit
ELF


processor:         x86_64
python:            3.6.9



Additional context
I am still learning pytorch-lightning, so I will readily admit that I might not understand how this is actually supposed to work.  However, I'll note that the documentation describes the semantics of on_step as "logs the metric at that step in training", which suggests the expected behavior described above.