Disable compilation for training_step, validation_step, etc.

[guarin]

Hi! I am training a model with Fabric and would like to compile only the forward method of the model and skip compilation for the training_step, validation_step, etc. The reason for this is that training_step contains hard to compile logic like metrics tracking and has dynamic inputs. forward however is safe to compile.

It seems that the forward wrapping magic in Fabric automatically results in training_step to be compiled as well and I haven't managed to exclude training_step from compilation. I tried wrapping it in @torch.compiler.disable() but that didn't work. Calling model.training_step = torch.compiler.disable(model.training_step, recursive=False) doesn't work if called before fabric.setup and results in a recursion depth exceeded error if called after fabric.setup.

Any hints on how to achieve this or what the best practice is for this scenario?

Code to reproduce:

from __future__ import annotations

import torch
from lightning_fabric import Fabric
from pytorch_lightning.demos.boring_classes import BoringModel, RandomDataset
from torch.utils.data import DataLoader
import uuid

class MyModel(BoringModel):
    def training_step(self, batch, filename) -> torch.Tensor:
        output = self(batch)
        loss = output.sum()
        return loss


def main() -> None:
    # Configuration
    batch_size = 32
    max_steps = 20

    # Initialize Fabric
    fabric = Fabric(accelerator="auto", devices="auto")
    fabric.launch()

    # Create model and optimizer
    model = MyModel()
    optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)

    # Setup with Fabric
    model = torch.compile(model)
    model, optimizer = fabric.setup(model, optimizer)

    # Create dataset
    dummy_dataset = RandomDataset(size=32, length=100)
    train_loader = DataLoader(dataset=dummy_dataset, batch_size=batch_size, shuffle=True)
    train_loader = fabric.setup_dataloaders(train_loader)

    # Training loop
    model.train()
    global_step = 0

    while global_step < max_steps:
        for batch in train_loader:
            if global_step >= max_steps:
                break
            fabric.print(f"Step: {global_step}")
            output = model.training_step(batch, filename=str(uuid.uuid4()))
            loss = output.sum()
            fabric.backward(loss)
            optimizer.step()
            optimizer.zero_grad()
            global_step += 1


if __name__ == "__main__":
    main()

Output:

Step: 1
Step: 2
Step: 3
Step: 4
Step: 5
Step: 6
Step: 7
[rank1]:W0306 11:23:42.278000 32517 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:1676] [0/8] torch._dynamo hit config.recompile_limit (8)
[rank1]:W0306 11:23:42.278000 32517 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:1676] [0/8]    function: 'inner' (.venv/lib/python3.10/site-packages/torch/_dynamo/external_utils.py:67)
[rank1]:W0306 11:23:42.278000 32517 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:1676] [0/8]    last reason: 0/7: kwargs['filename'] == 'd747cced-2800-4811-a060-b8a54021ecf3'
[rank1]:W0306 11:23:42.278000 32517 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:1676] [0/8] To log all recompilation reasons, use TORCH_LOGS="recompiles".
[rank1]:W0306 11:23:42.278000 32517 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:1676] [0/8] To diagnose recompilation issues, see https://pytorch.org/docs/main/compile/programming_model.recompilation.html
Step: 8
[rank0]:W0306 11:23:42.280000 32441 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:1676] [0/8] torch._dynamo hit config.recompile_limit (8)
[rank0]:W0306 11:23:42.280000 32441 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:1676] [0/8]    function: 'inner' (.venv/lib/python3.10/site-packages/torch/_dynamo/external_utils.py:67)
[rank0]:W0306 11:23:42.280000 32441 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:1676] [0/8]    last reason: 0/7: kwargs['filename'] == '5aeb7ef4-5387-4621-953e-7bb1154a7f3d'
[rank0]:W0306 11:23:42.280000 32441 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:1676] [0/8] To log all recompilation reasons, use TORCH_LOGS="recompiles".
[rank0]:W0306 11:23:42.280000 32441 .venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:1676] [0/8] To diagnose recompilation issues, see https://pytorch.org/docs/main/compile/programming_model.recompilation.html
Step: 9
Step: 10
Step: 11
Step: 12
Step: 13
Step: 14
Step: 15
Step: 16
Step: 17
Step: 18
Step: 19

Here the compile limit is reached because filename is different in every training step.

What does work is calling torch.compile only on the forward:

model.forward = torch.compile(model.forward)
model, optimizer = fabric.setup(model, optimizer)

But then fabric.setup doesn't re-apply compilation because not the whole model was compiled. This is at least my understanding.

[gitsult4n]

torch.compiler.disable() on training_step cannot win here because Fabric hijacks that method by name before your decorator gets a say

_LIGHTNING_MODULE_STEP_METHODS = ("training_step", "validation_step", "test_step", "predict_step")

is seeded straight into _FabricModule._forward_methods and __getattr__ routes every one of them through _redirection_through_forward whenever _forward_module != _original_module

fabric.setup makes them different on purpose when the model is compiled
it unwraps your OptimizedModule keeps the clean copy as original_module recompiles the strategy wrapped module and only then builds _FabricModule So calling model.training_step(...) does not run your method at all it monkey patches original_module.forward to re enter training_step and calls the compiled forward instead That is also why decorating the step recurses the patch and the compiled call are the same trip

Two ways out

Compile only forward before fabric.setup

model.forward = torch.compile(model.forward)
model, optimizer = fabric.setup(model, optimizer)

_unwrap_compiled then finds no OptimizedModule on the LightningModule itself _to_compiled is skipped and the redirection never lands in a compiled region

Or keep compiling the wrapped module but stop Fabric from reapplying it

fabric.setup(model, optimizer, _reapply_compile=False)

documented for exactly this case compiling the DDP or FSDP wrapped module without Fabric refusing it