Issue with logs when using torch.compile

[Forbu]

Bug description

I experience a bug with logging when using torch compile.

Traceback (most recent call last):
  File "/app/data/src_stardust/prediction/training/train_gnn_realdata.py", line 586, in <module>
    trainer.fit(modele, train_loader, validation_loader, ckpt_path=last_checkpoint)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 545, in fit
    call._call_and_handle_interrupt(
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 44, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 581, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 990, in _run
    results = self._run_stage()
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 1034, in _run_stage
    self._run_sanity_check()
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 1063, in _run_sanity_check
    val_loop.run()
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/loops/utilities.py", line 181, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 134, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 391, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_args)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 309, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/strategies/strategy.py", line 403, in validation_step
    return self.lightning_module.validation_step(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 328, in _fn
    return fn(*args, **kwargs)
  File "/app/data/src_stardust/prediction/training/train_gnn_realdata.py", line 296, in validation_step
    ) = st.preprocessing_batch_real(
  File "/app/data/src_stardust/prediction/training/train_gnn_realdata.py", line 344, in <resume in validation_step>
    self.log(
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/core/module.py", line 412, in log
    apply_to_collection(value, dict, self.__check_not_nested, name)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/core/module.py", line 413, in <resume in log>
    apply_to_collection(
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/core/module.py", line 500, in <resume in log>
    results.log(
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py", line 390, in log
    meta.sync = _Sync(_should=sync_dist, fn=sync_dist_fn, _group=sync_dist_group, rank_zero_only=rank_zero_only)
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py", line 398, in <resume in log>
    raise MisconfigurationException(
  File "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py", line 398, in <resume in log>
    raise MisconfigurationException(
lightning.fabric.utilities.exceptions.MisconfigurationException: You called `self.log(val_aff_db, ...)` twice in `validation_step` with different arguments. This is not allowed

In my validation_step method, I only have one self.log statement :

self.log(
                name="validation_metric",
                value=myvalue,
                on_step=True,
                on_epoch=True,
                logger=True,
                batch_size=1,
                prog_bar=False,
                reduce_fx="mean",
                enable_graph=False,
                add_dataloader_idx=False,
                metric_attribute=None,
                rank_zero_only=True,
                sync_dist=False,
            )

This issue only happen using torch.compile() on the whole lightning module.
I have the latest lightning version 2.1
(and torch 2.1)

What version are you seeing the problem on?

master

How to reproduce the bug

I will try to use the simple BoringModule from the test to easily reproduce the bug (I will do that in the comment later).

Environment

Current environment

- Lightning Component (e.g. Trainer, LightningModule):
- PyTorch Lightning Version (e.g., 2.1):
- PyTorch Version (e.g., 2.1):
- I installed lightning with pip

More info

This bug only happen when I compile my module :

myligthningmodule = torch.compile(myligthningmodule)

cc @carmocca

[apardyl]

+1 to this issue, same version of lightning and pytorch. Seems that the _Metadata class https://github.com/Lightning-AI/lightning/blob/e50b68aae37be119272668860472da25a6b697b6/src/lightning/pytorch/trainer/connectors/logger_connector/result.py#L107 should define a custom __eq__.

[carmocca]

For the moment, and as a workaround, I would strongly suggest that you simply torch.compile the underlying nn.Module instead of the LightningModule

[Forbu]

@carmocca Yeah this is exactly what I did ! simple workaround

[carmocca]

Here's a repro

import os

import torch
from lightning.pytorch import LightningModule, Trainer
from torch.utils.data import DataLoader, Dataset

from lightning.pytorch.demos.boring_classes import RandomDataset


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)
        return {"loss": loss}

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=2,
        limit_val_batches=0,
        num_sanity_val_steps=0,
        max_epochs=2,
        enable_model_summary=False,
    )
    cmodel = torch.compile(model)
    trainer.fit(cmodel, train_dataloaders=train_data)


if __name__ == "__main__":
    run()

Torch compile seems to have multiple issues with compiling the logging logic. Simply doing

import torch

from lightning.pytorch.trainer.connectors.logger_connector.result import _Metadata


def log():
    meta = _Metadata("a", "loss")
    meta2 = _Metadata("a", "loss")
    if meta != meta2:
        raise ValueError("Foo")

log = torch.compile(log)
log()

already fails. There's no plan to fix this soon. It will require finding workarounds as well as reporting and fixing these limitations on PyTorch.

[carmocca]

For anybody desperate, this is a valid workaround:

diff --git a/src/lightning/pytorch/trainer/connectors/logger_connector/result.py b/src/lightning/pytorch/trainer/connectors/logger_connector/result.py
index 81630362a..4d0b88208 100644
--- a/src/lightning/pytorch/trainer/connectors/logger_connector/result.py
+++ b/src/lightning/pytorch/trainer/connectors/logger_connector/result.py
@@ -398,10 +398,12 @@ class _ResultCollection(dict):
             self.register_key(key, meta, value)
 
         # check the stored metadata and the current one match
-        elif meta != self[key].meta:
-            raise MisconfigurationException(
-                f"You called `self.log({name}, ...)` twice in `{fx}` with different arguments. This is not allowed"
-            )
+        else:
+            torch._dynamo.graph_break()
+            if meta != self[key].meta:
+                raise MisconfigurationException(
+                    f"You called `self.log({name}, ...)` twice in `{fx}` with different arguments. This is not allowed"
+                )
 
         batch_size = self._extract_batch_size(self[key], batch_size, meta)
         self.update_metrics(key, value, batch_size)