Repository: mlflow/mlflow

Issue:
[BUG] PyTorch Lightning autologging doesn't handle step number correctly with multiple optimizers
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Where did you encounter this bug?

Local machine

### Willingness to contribute

Yes. I can contribute a fix for this bug independently.

### MLflow version

- Client: 2.8.1
- Tracking server: 2.8.1

### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux Fedora 39
- **Python version**: 3.11
- **yarn version, if running the dev UI**: NA
- **lightning version**: 2.1.2

### Describe the problem

When using multiple optimizers with PyTorch lightning (when training a GAN for example) and `log_every_step` > 1, metrics logged per step might never be logged to mlflow or might be logged at a different frequency than expected, depending on the parameters used.

PyTorch Lightning increments the global step once per optimizer, so for example with 2 optimizers, the step number seen in the `__MLflowPLCallback.on_train_batch_end` method will be 2, 4, 6 etc, and the `(step + 1) % self.log_every_n_step == 0` condition will never be met if `log_every_step` is an even number.

I've verified that this behaviour is the same when using manual optimization and the old approach used in PyTorch lightning < 2.0.0 (tested with 1.9.5) where the optimizer index is passed to the `training_step` method.

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
NA
```

### Code to reproduce issue

The following code for training a GAN was adapted from https://github.com/nocotan/pytorch-lightning-gans/blob/master/models/dcgan.py and updated to work with manual optimization for PyTorch Lightning >= 2.0. I haven't verified that the training logic is still correct but it reproduces the issue. No metrics are logged at all unless `log_every_n_step` is changed to 1.

```python
import os
from argparse import ArgumentParser, Namespace

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST

import lightning as L
import mlflow.pytorch


class Generator(nn.Module):
    def __init__(self, latent_dim, img_shape):
        super().__init__()
        self.img_shape = img_shape

        self.init_size = img_shape[1] // 4
        self.l1 = nn.Sequential(nn.Linear(latent_dim, 128 * self.init_size ** 2))

        self.conv_blocks = nn.Sequential(
            nn.BatchNorm2d(128),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(128, 128, 3, stride=1, padding=1),
            nn.BatchNorm2d(128, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(128, 64, 3, stride=1, padding=1),
            nn.BatchNorm2d(64, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, img_shape[0], 3, stride=1, padding=1),
            nn.Tanh(),
        )

    def forward(self, z):
        out = self.l1(z)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img


class Discriminator(nn.Module):
    def __init__(self, img_shape):
        super(Discriminator, self).__init__()

        def discriminator_block(in_feat, out_feat, bn=True):
            block = [nn.Conv2d(in_feat, out_feat, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]
            if bn:
                block.append(nn.BatchNorm2d(out_feat, 0.8))
            return block

        self.model = nn.Sequential(
            *discriminator_block(img_shape[0], 16, bn=False),
            *discriminator_block(16, 32),
            *discriminator_block(32, 64),
            *discriminator_block(64, 128),
        )

        # The height and width of downsampled image
        ds_size = img_shape[1] // 2 ** 4
        self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid())

    def forward(self, img):
        out = self.model(img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)

        return validity


class DCGAN(L.LightningModule):
    def __init__(self,
                 latent_dim: int = 100,
                 lr: float = 0.0002,
                 b1: float = 0.5,
                 b2: float = 0.999,
                 batch_size: int = 64,
                 **kwargs):
        super().__init__()
        self.save_hyperparameters()

        self.latent_dim = latent_dim
        self.lr = lr
        self.b1 = b1
        self.b2 = b2
        self.batch_size = batch_size

        # networks
        img_shape = (1, 32, 32)
        self.generator = Generator(latent_dim=self.latent_dim, img_shape=img_shape)
        self.discriminator = Discriminator(img_shape=img_shape)

        self.validation_z = torch.randn(8, self.latent_dim)

        self.example_input_array = torch.zeros(2, self.latent_dim)

        self.automatic_optimization = False

    def forward(self, z):
        return self.generator(z)

    def adversarial_loss(self, y_hat, y):
        return F.binary_cross_entropy(y_hat, y)

    def training_step(self, batch, batch_idx):
        g_opt, d_opt = self.optimizers()

        imgs, _ = batch

        z = torch.randn(imgs.shape[0], self.latent_dim)
        z = z.type_as(imgs)

        # train generator

        valid = torch.ones(imgs.size(0), 1)
        valid = valid.type_as(imgs)

        g_loss = self.adversarial_loss(self.discriminator(self(z)), valid)

        g_opt.zero_grad()
        self.manual_backward(g_loss)
        g_opt.step()

        # train discriminator

        valid = torch.ones(imgs.size(0), 1)
        valid = valid.type_as(imgs)

        real_loss = self.adversarial_loss(self.discriminator(imgs), valid)

        fake = torch.zeros(imgs.size(0), 1)
        fake = fake.type_as(imgs)

        fake_loss = self.adversarial_loss(
            self.discriminator(self(z).detach()), fake)

        d_loss = (real_loss + fake_loss) / 2

        d_opt.zero_grad()
        self.manual_backward(d_loss)
        d_opt.step()

        self.log_dict({
            'g_loss': g_loss,
            'd_loss': d_loss,
        }, on_step=True)

    def configure_optimizers(self):
        lr = self.lr
        b1 = self.b1
        b2 = self.b2

        opt_g = torch.optim.Adam(self.generator.parameters(), lr=lr, betas=(b1, b2))
        opt_d = torch.optim.Adam(self.discriminator.parameters(), lr=lr, betas=(b1, b2))
        return [opt_g, opt_d], []

    def train_dataloader(self):
        transform = transforms.Compose([
            transforms.Resize((32, 32)),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ])

        dataset = MNIST(os.getcwd(), train=True, download=True, transform=transform)
        return DataLoader(dataset, batch_size=self.batch_size, num_workers=0)


def main(args: Namespace) -> None:
    mlflow.pytorch.autolog(log_every_n_step=50)

    model = DCGAN(**vars(args))

    devices = "auto" if args.gpus == 0 else args.gpus
    accelerator = "gpu" if args.gpus > 0 else "cpu"
    trainer = L.Trainer(
            accelerator=accelerator,
            devices=devices,
            max_epochs=4,
            log_every_n_steps=50)

    trainer.fit(model)


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument("--gpus", type=int, default=0, help="number of GPUs")
    parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
    parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
    parser.add_argument("--b1", type=float, default=0.5,
                        help="adam: decay of first order momentum of gradient")
    parser.add_argument("--b2", type=float, default=0.999,
                        help="adam: decay of first order momentum of gradient")
    parser.add_argument("--latent_dim", type=int, default=100,
                        help="dimensionality of the latent space")

    hparams = parser.parse_args()

    main(hparams)
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
NA
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
NA
```


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/gateway`: AI Gateway service, Gateway client APIs, third-party Gateway integrations
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [X] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:

Resolving patch:
diff --git a/mlflow/pytorch/__init__.py b/mlflow/pytorch/__init__.py
--- a/mlflow/pytorch/__init__.py
+++ b/mlflow/pytorch/__init__.py
@@ -927,9 +927,12 @@ def autolog(
 
     :param log_every_n_epoch: If specified, logs metrics once every `n` epochs. By default, metrics
                        are logged after every epoch.
-    :param log_every_n_step: If specified, logs batch metrics once every `n` global step.
+    :param log_every_n_step: If specified, logs batch metrics once every `n` training step.
                        By default, metrics are not logged for steps. Note that setting this to 1 can
                        cause performance issues and is not recommended.
+                       Metrics are logged against Lightning's global step number,
+                       and when multiple optimizers are used it is assumed that all optimizers
+                       are stepped in each training step.
     :param log_models: If ``True``, trained models are logged as MLflow model artifacts.
                        If ``False``, trained models are not logged.
     :param log_datasets: If ``True``, dataset information is logged to MLflow Tracking.
diff --git a/mlflow/pytorch/_lightning_autolog.py b/mlflow/pytorch/_lightning_autolog.py
--- a/mlflow/pytorch/_lightning_autolog.py
+++ b/mlflow/pytorch/_lightning_autolog.py
@@ -86,6 +86,7 @@ def __init__(
         self.log_models = log_models
         self.log_every_n_epoch = log_every_n_epoch
         self.log_every_n_step = log_every_n_step
+        self._global_steps_per_training_step = 1
         # Sets for tracking which metrics are logged on steps and which are logged on epochs
         self._step_metrics = set()
         self._epoch_metrics = set()
@@ -212,7 +213,7 @@ def on_train_batch_end(
         ]
         self._step_metrics.update(name for (name, _) in metric_items)
         step = trainer.global_step
-        if (step + 1) % self.log_every_n_step == 0:
+        if ((step // self._global_steps_per_training_step) + 1) % self.log_every_n_step == 0:
             self._log_metrics(trainer, step, metric_items)
 
     @rank_zero_only
@@ -236,6 +237,10 @@ def on_train_start(self, trainer, pl_module):
         #    to rename the optimizer parameters based on keys in default dictionary.
 
         if hasattr(trainer, "optimizers"):
+            # Lightning >= 1.6.0 increments the global step every time an optimizer is stepped.
+            # We assume every optimizer will be stepped in each training step.
+            if _pl_version >= Version("1.6.0"):
+                self._global_steps_per_training_step = len(trainer.optimizers)
             optimizer = trainer.optimizers[0]
             params["optimizer_name"] = _get_optimizer_name(optimizer)
 


Test patch:
diff --git a/tests/pytorch/iris.py b/tests/pytorch/iris.py
--- a/tests/pytorch/iris.py
+++ b/tests/pytorch/iris.py
@@ -106,5 +106,38 @@ def test_step(self, batch, batch_idx):
         self.log("test_acc", self.test_acc.compute())
 
 
+class IrisClassificationMultiOptimizer(IrisClassificationBase):
+    """
+    Contrived lightning module that uses multiple optimizers. In real-world scenarios
+    multiple optimizers might be used for Generative Adversarial Networks (GANs).
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.automatic_optimization = False
+
+    def training_step(self, batch, batch_idx):
+        opt1, opt2 = self.optimizers()
+
+        x, y = batch
+        logits = self.forward(x)
+        loss = self.cross_entropy_loss(logits, y)
+
+        opt1.zero_grad()
+        opt2.zero_grad()
+
+        self.manual_backward(loss)
+
+        opt1.step()
+        opt2.step()
+
+        self.log("loss", loss, on_epoch=True, on_step=True)
+
+    def configure_optimizers(self):
+        opt1 = torch.optim.Adam(self.parameters(), 0.01)
+        opt2 = torch.optim.Adam(self.parameters(), 0.01)
+        return [opt1, opt2]
+
+
 if __name__ == "__main__":
     pass
diff --git a/tests/pytorch/test_pytorch_autolog.py b/tests/pytorch/test_pytorch_autolog.py
--- a/tests/pytorch/test_pytorch_autolog.py
+++ b/tests/pytorch/test_pytorch_autolog.py
@@ -1,7 +1,11 @@
 import pytest
 import pytorch_lightning as pl
 import torch
-from iris import IrisClassification, IrisClassificationWithoutValidation
+from iris import (
+    IrisClassification,
+    IrisClassificationMultiOptimizer,
+    IrisClassificationWithoutValidation,
+)
 from iris_data_module import IrisDataModule, IrisDataModuleWithoutValidation
 from packaging.version import Version
 from pytorch_lightning.callbacks import ModelCheckpoint
@@ -57,6 +61,20 @@ def pytorch_model_with_steps_logged(request):
     return trainer, run, log_every_n_epoch, log_every_n_step
 
 
+@pytest.fixture(params=[(1, 1), (1, 10), (2, 1)])
+def pytorch_multi_optimizer_model(request):
+    log_every_n_epoch, log_every_n_step = request.param
+    mlflow.pytorch.autolog(log_every_n_epoch=log_every_n_epoch, log_every_n_step=log_every_n_step)
+    model = IrisClassificationMultiOptimizer()
+    dm = IrisDataModule()
+    dm.setup(stage="fit")
+    trainer = pl.Trainer(max_epochs=NUM_EPOCHS)
+    trainer.fit(model, dm)
+    client = MlflowClient()
+    run = client.get_run(client.search_runs(["0"])[0].info.run_id)
+    return trainer, run, log_every_n_epoch, log_every_n_step
+
+
 @pytest.mark.parametrize("log_models", [True, False])
 def test_pytorch_autolog_log_models_configuration(log_models):
     mlflow.pytorch.autolog(log_models=log_models)
@@ -168,6 +186,30 @@ def test_pytorch_autolog_logging_forked_metrics_on_step_and_epoch(
         ), f"Expected {expected_len} values for {metric_key}, got {len(metric_history)}"
 
 
+@pytest.mark.skipif(
+    Version(pl.__version__) < Version("1.1.0"),
+    reason="Access to metrics from the current step is only possible since PyTorch-lightning 1.1.0 "
+    "when LoggerConnector.cached_results was added",
+)
+def test_pytorch_autolog_log_on_step_with_multiple_optimizers(
+    pytorch_multi_optimizer_model,
+):
+    trainer, run, log_every_n_epoch, log_every_n_step = pytorch_multi_optimizer_model
+    num_logged_steps = NUM_EPOCHS * len(trainer.train_dataloader) // log_every_n_step
+    num_logged_epochs = NUM_EPOCHS // log_every_n_epoch
+
+    client = MlflowClient()
+    for metric_key, expected_len in [
+        ("loss", num_logged_epochs),
+        ("loss_step", num_logged_steps),
+    ]:
+        assert metric_key in run.data.metrics, f"Missing {metric_key} in metrics"
+        metric_history = client.get_metric_history(run.info.run_id, metric_key)
+        assert (
+            len(metric_history) == expected_len
+        ), f"Expected {expected_len} values for {metric_key}, got {len(metric_history)}"
+
+
 @pytest.mark.skipif(
     Version(pl.__version__) >= Version("1.1.0"),
     reason="Logging step metrics is supported since PyTorch-Lightning 1.1.0",