Bug with `torch.compile` and `self.logger.experiment.log`

[mishooax]

Bug description

Running with lightning 2.0 and pytorch 2.0.

Hi, I'm torch.compiling my lightning model - while I do see a non-negligible (~ 20%!) speedup in training, torch.dynamo errors out in the validation loop when it encounters a self.logger.experiment.log call to Weights&Biases (see below)

Is there anything i can do to stop torch.dynamo from choking on this bit of code? Maybe some decorator that tells JIT to ignore the offending self.logger.experiment.log call? Many thanks!

How to reproduce the bug

# torch.compile call
# type(model) == LightningModel
model = torch.compile(model, mode="default", backend="inductor", fullgraph=False)
# the model should log an image to wandb through a self.logger.experiment.log call

Error messages and logs

[2023-03-28 10:14:51,227] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)
   function: '<graph break in draw_text>' (/python3.9/site-packages/matplotlib/backends/backend_agg.py:212)
   reasons:  d == 2.0
to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.
[2023-03-28 10:14:51,298] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)
   function: '<graph break in draw_text>' (python3.9/site-packages/matplotlib/backends/backend_agg.py:212)
   reasons:  d == 2.0
to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.
[2023-03-28 10:14:51,621] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)
   function: '<graph break in draw_text>' (python3.9/site-packages/matplotlib/backends/backend_agg.py:212)
   reasons:  d == 2.0
to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.
[2023-03-28 10:14:51,968] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)
   function: '<graph break in draw_text>' (python3.9/site-packages/matplotlib/backends/backend_agg.py:212)
   reasons:  d == 2.0
to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.
[2023-03-28 10:14:52,445] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)
   function: '<graph break in draw_text>' (python3.9/site-packages/matplotlib/backends/backend_agg.py:204)
   reasons:  s == 't850 target'
to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.
Traceback (most recent call last):
  File "python3.9/site-packages/torch/_dynamo/convert_frame.py", line 324, in _compile
    out_code = transform_code_object(code, transform)
  File "python3.9/site-packages/torch/_dynamo/bytecode_transformation.py", line 445, in transform_code_object
    transformations(instructions, code_options)
  File "python3.9/site-packages/torch/_dynamo/convert_frame.py", line 311, in transform
    tracer.run()
  File "python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1726, in run
    super().run()
  File "python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 576, in run
    and self.step()
  File "python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 540, in step
    getattr(self, inst.opname)(inst)
  File "python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1030, in LOAD_ATTR
    result = BuiltinVariable(getattr).call_function(
  File "python3.9/site-packages/torch/_dynamo/variables/builtin.py", line 566, in call_function
    result = handler(tx, *args, **kwargs)
  File "python3.9/site-packages/torch/_dynamo/variables/builtin.py", line 958, in call_getattr
    obj.var_getattr(tx, name).clone(source=source).add_options(options)
  File "python3.9/site-packages/torch/_dynamo/variables/user_defined.py", line 319, in var_getattr
    return variables.UserMethodVariable(
  File "python3.9/site-packages/torch/_dynamo/variables/functions.py", line 291, in call_function
    return super().call_function(tx, args, kwargs)
  File "python3.9/site-packages/torch/_dynamo/variables/functions.py", line 259, in call_function
    return super().call_function(tx, args, kwargs)
  File "python3.9/site-packages/torch/_dynamo/variables/functions.py", line 92, in call_function
    return tx.inline_user_function_return(
  File "/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 510, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1806, in inline_call
Traceback (most recent call last):
  File "python3.9/site-packages/torch/_dynamo/convert_frame.py", line 324, in _compile
    return cls.inline_call_(parent, func, args, kwargs)
  File "python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1862, in inline_call_
    out_code = transform_code_object(code, transform)
  File "python3.9/site-packages/torch/_dynamo/bytecode_transformation.py", line 445, in transform_code_object
    tracer.run()
  File "python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 576, in run
    transformations(instructions, code_options)
  File "python3.9/site-packages/torch/_dynamo/convert_frame.py", line 311, in transform
    and self.step()
  File "python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 540, in step
    tracer.run()
  File "python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1726, in run
    getattr(self, inst.opname)(inst)
  File "python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1975, in LOAD_CLOSURE
    super().run()
  File "python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 576, in run
    self.push(self.closure_cells[inst.argval])
KeyError: fn

...

from user code:
   File "trainer.py", line 214, in <graph break in _output_figure>
    self.logger.experiment.log({exp_log_tag: wandb.Image(save_path)})
  File "python3.9/site-packages/lightning_fabric/loggers/logger.py", line 114, in experiment
    def get_experiment() -> Callable:

Set torch._dynamo.config.verbose=True for more information

Environment

Current environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): LightningModule
#- PyTorch Lightning Version (e.g., 1.5.0): 2.0
#- Lightning App Version (e.g., 0.5.2): n/a
#- PyTorch Version (e.g., 2.0): 2.0
#- Python version (e.g., 3.9): 3.9
#- OS (e.g., Linux): Linux
#- CUDA/cuDNN version: 11.8
#- GPU models and configuration: A100
#- How you installed Lightning(`conda`, `pip`, source): conda
#- Running environment of LightningApp (e.g. local, cloud): n/a

More info

No response

cc @carmocca

[mishooax]

minimal example:

import os

import torch
import wandb
from lightning.pytorch import LightningModule, Trainer
from lightning.pytorch.loggers.wandb import WandbLogger
from torch.utils.data import DataLoader

from lightning.pytorch.demos.boring_classes import RandomDataset


def get_logger():
    os.environ["WANDB_MODE"] = "dryrun"
    return WandbLogger(
        project="...",
        entity="...",
    )

class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def validation_step(self, batch, batch_idx):
        dummy_image = torch.randn(batch.shape[1], batch.shape[1]).numpy()
        self.logger.experiment.log({"image": wandb.Image(dummy_image)})


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

    model = BoringModel()
    model = torch.compile(model)
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_val_batches=1,
        enable_model_summary=False,
        enable_progress_bar=False,
        enable_checkpointing=False,
        logger=get_logger(),
    )
    trainer.validate(model, dataloaders=test_data)


if __name__ == "__main__":
    run()

Error:

from user code:
   File "compile_test.py", line 45, in <graph break in validation_step>
    self.logger.experiment.log({"image": wandb.Image(dummy_image)})
  File "/lib/python3.9/site-packages/lightning_fabric/loggers/logger.py", line 114, in experiment
    def get_experiment() -> Callable:

Set torch._dynamo.config.verbose=True for more information
You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "compile_test.py", line 78, in <module>
    run()
  File "compile_test.py", line 73, in run
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
  File "/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 520, in fit
    call._call_and_handle_interrupt(
  File "/lib/python3.9/site-packages/pytorch_lightning/trainer/call.py", line 44, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 559, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 935, in _run
    results = self._run_stage()
  File "/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 978, in _run_stage
    self.fit_loop.run()
  File "/lib/python3.9/site-packages/pytorch_lightning/loops/fit_loop.py", line 201, in run
    self.advance()
  File "/lib/python3.9/site-packages/pytorch_lightning/loops/fit_loop.py", line 354, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/lib/python3.9/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 134, in run
    self.on_advance_end()
  File "/lib/python3.9/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 248, in on_advance_end
    self.val_loop.run()
  File "/lib/python3.9/site-packages/pytorch_lightning/loops/utilities.py", line 174, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/lib/python3.9/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 115, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx)
  File "/lib/python3.9/site-packages/pytorch_lightning/loops/evaluation_loop.py", line 375, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_kwargs.values())
  File "/lib/python3.9/site-packages/pytorch_lightning/trainer/call.py", line 288, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/lib/python3.9/site-packages/pytorch_lightning/strategies/strategy.py", line 378, in validation_step
    return self.model.validation_step(*args, **kwargs)
  File "/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "compile_test.py", line 43, in validation_step
    loss = self(batch).sum()
  File "compile_test.py", line 44, in <graph break in validation_step>
    dummy_image = torch.randn(batch.shape[1], batch.shape[1]).numpy()
  File "/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 337, in catch_errors
    return callback(frame, cache_size, hooks)
  File "/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 404, in _convert_frame
    result = inner_convert(frame, cache_size, hooks)
  File "/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 104, in _fn
    return fn(*args, **kwargs)
  File "/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 262, in _convert_frame_assert
    return _compile(
  File "/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 163, in time_wrapper
    r = func(*args, **kwargs)
  File "/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 394, in _compile
    raise InternalTorchDynamoError() from e
torch._dynamo.exc.InternalTorchDynamoError

[carmocca]

I opened #17216 with a fix. Can you give it a try?

[mishooax]

@carmocca Many thanks for looking into this. Hmm... so I patched .../lib/python3.9/site-packages/lightning_fabric/loggers/logger.py in my lightning 2.0 install - as in your PR:

def rank_zero_experiment(fn: Callable) -> Callable:
    """Returns the real experiment on rank 0 and otherwise the _DummyExperiment."""

    @wraps(fn)
    @rank_zero_only
    def experiment(self) -> Union[Any, _DummyExperiment]:  # type: ignore[no-untyped-def]
        """
        Note:
            ``self`` is a custom logger instance. The loggers typically wrap an ``experiment`` method
            with a ``@rank_zero_experiment`` decorator.

            ``Union[Any, _DummyExperiment]`` is used because the wrapped hooks have several return
            types that are specific to the custom logger. The return type here can be considered as
            ``Union[return type of logger.experiment, _DummyExperiment]``.
        """
        return fn(self) or _DummyExperiment()

The simple test doesn't crash after this change, but I'm also seeing

Exception ignored in: <generator object iter_slices at 0x147d8c1116d0>                                                                                                                                                                                                                        
Traceback (most recent call last):
  File ".../lib/python3.9/site-packages/requests/utils.py", line 575, in iter_slices
    def iter_slices(string, slice_length):
SystemError: error return without exception set

My "real" code still crashes, this time with a different error (related to the iter_slices exception above):

2023-03-28 17:17:09,129] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)
   function: '<graph break in draw_text>' (<...>/lib/python3.9/site-packages/matplotlib/backends/backend_agg.py:212)
   reasons:  d == 2.0
to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.
[2023-03-28 17:17:09,207] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)
   function: '<graph break in draw_text>' (<...>/lib/python3.9/site-packages/matplotlib/backends/backend_agg.py:212)
   reasons:  d == 2.0
to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.
[2023-03-28 17:17:09,210] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)
   function: '<graph break in draw_text>' (<...>/lib/python3.9/site-packages/matplotlib/backends/backend_agg.py:212)
   reasons:  d == 2.0
to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.
[2023-03-28 17:17:10,351] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)
   function: '<graph break in draw_text>' (<...>/lib/python3.9/site-packages/matplotlib/backends/backend_agg.py:204)
   reasons:  s == 't850 target'
to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.
[2023-03-28 17:17:10,468] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)
   function: '<graph break in draw_text>' (<...>/lib/python3.9/site-packages/matplotlib/backends/backend_agg.py:212)
   reasons:  d == 2.0
to diagnose recompilation issues, see https://pytorch.org/docs/master/dynamo/troubleshooting.html.

....

  File "<...>train/trainer.py", line 214, in <graph break in _output_figure>
    self.logger.experiment.log({exp_log_tag: wandb.Image(save_path)})
AttributeError: 'NoneType' object has no attribute 'log'
    self._plot_sample(batch_idx, rstep, x[..., : self.feature_dim], y[..., : self.feature_dim], y_hat)
  File "<...>train/trainer.py", line 195, in _plot_sample
    self.logger.experiment.log({exp_log_tag: wandb.Image(save_path)})
AttributeError: 'NoneType' object has no attribute 'log'

...

  File "<...>train/trainer.py", line 214, in <graph break in _output_figure>
    self.logger.experiment.log({exp_log_tag: wandb.Image(save_path)})
AttributeError: 'NoneType' object has no attribute 'log'
Exception ignored in: <generator object iter_slices at 0x145a9c9cde40>
Traceback (most recent call last):
  File "<...>/lib/python3.9/site-packages/requests/utils.py", line 575, in iter_slices
    def iter_slices(string, slice_length):
SystemError: error return without exception set

[carmocca]

Can you provide a new repro for that issue?

As you see, torch.compile is very fragile at the moment. You can work around many of these issues by just compiling your nn.Module (example).

[mishooax]

Can you provide a new repro for that issue?

@carmocca - I can try to come up with one, yes.

How about the iters_slices exception? Any idea what may be causing that? I am seeing it with the minimal example I posted above, even after including your patch. Many thanks.

[carmocca]

I'm not able to reproduce the iter_slices issue. I'm using Python 3.10.6