torch.compile doesn't work with self.toggle_optimizer() in lightning module

[Mo-Junyang]

Bug description

I simply add a self.toggle_optimizer(self.optimizers()) in training_step,

class SimpleModel(L.LightningModule):
    """Simple model for testing"""
    
    def __init__(self):
        super().__init__()
        self.layer = torchvision.models.resnet50()
        # self.layer = nn.Linear(10, 10)
    def forward(self, x):
        return self.layer(x, x)[0]
    
    def training_step(self, batch, batch_idx):
        self.toggle_optimizer(self.optimizers())
        x, y = batch
        y_hat = self(x)
        loss = F.mse_loss(y_hat, y)
        self.log('train_loss', loss)
        return loss
    
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.mse_loss(y_hat, y)
        self.log('val_loss', loss)
        return loss
    
    def on_train_epoch_end(self) -> None:
        print(0)
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

and run

    model = SimpleModel()
    model = torch.compile(model)
    trainer.fit(model, train_loader, val_loader)

then the error comes:

    param.requires_grad = param_requires_grad_state[param]
                          ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^
KeyError: Parameter containing:
tensor([ 0.0567,  0.0312,  0.0305,  0.0453, -0.0708, -0.0092,  0.0487,  0.0671,
         0.0399, -0.0994,  0.0946,  0.0910,  0.0989,  0.0787, -0.0801, -0.0465,
        -0.0522, -0.0650, -0.0257, -0.0062,  0.0660, -0.0539,  0.0658,  0.0259,
        -0.0534, -0.0563, -0.0235,  0.0693, -0.0789,  0.0502, -0.0854,  0.0091])
Epoch 0:   0%|          | 0/125 [00:16<?, ?it/s] 

But I do need toggle_optimizer in GAN training. Can this be fixed?

What version are you seeing the problem on?

v2.5

Reproduced in studio

No response

How to reproduce the bug

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):

More info

No response

cc @ethanwharris

[CodeVishal-17]

This looks like an incompatibility between torch.compile and LightningModule.toggle_optimizer, likely due to dynamic requires_grad mutation during graph capture—happy to open a PR or help test a fix if there’s a preferred direction.

[deependujha]

Hi @Mo-Junyang, thanks for raising the issue.

Sample complete code to reproduce the bug:

import torch
import lightning as L
from torch.utils.data import DataLoader, TensorDataset
import torchvision.models as models
import torch.nn.functional as F

# ------------------
# data (fake image data)
# ------------------
x = torch.randn(64, 3, 224, 224)   # image-like input
y = torch.randint(0, 2, (64,))
loader = DataLoader(TensorDataset(x, y), batch_size=4, shuffle=True)


# ------------------
# model
# ------------------
class TwoOptResNet(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.automatic_optimization = False

        self.backbone = models.resnet50(weights=None)
        self.backbone.fc = torch.nn.Identity()  # remove classifier
        self.head = torch.nn.Linear(2048, 2)

    def forward(self, x):
        feats = self.backbone(x)
        return self.head(feats)

    def training_step(self, batch, batch_idx):
        opt_backbone, opt_head = self.optimizers()
        x, y = batch

        # ---- backbone optimizer ----
        self.toggle_optimizer(opt_backbone)

        logits = self(x)
        loss_backbone = F.cross_entropy(logits, y)

        opt_backbone.zero_grad()
        self.manual_backward(loss_backbone)
        opt_backbone.step()

        self.untoggle_optimizer(opt_backbone)

        # ---- head optimizer ----
        self.toggle_optimizer(opt_head)

        logits = self(x)
        loss_head = F.cross_entropy(logits, y)

        opt_head.zero_grad()
        self.manual_backward(loss_head)
        opt_head.step()

        self.untoggle_optimizer(opt_head)

        self.log_dict(
            {
                "loss_backbone": loss_backbone,
                "loss_head": loss_head,
            }
        )

    def configure_optimizers(self):
        opt_backbone = torch.optim.Adam(self.backbone.parameters(), lr=1e-4)
        opt_head = torch.optim.Adam(self.head.parameters(), lr=1e-3)
        return [opt_backbone, opt_head]


# ------------------
# trainer
# ------------------
trainer = L.Trainer(
    max_epochs=1,
    accelerator="auto",
    devices=1,
    log_every_n_steps=1,
)

model = TwoOptResNet()

# ⚠️ THIS IS THE PROBLEMATIC PART
model = torch.compile(model)

trainer.fit(model, loader)

---

here's what ChatGPT says:

• torch.compile: static assumptions about parameters + backward graphs
• toggle_optimizer: dynamic parameter mutation (requires grad)

toggle_optimizer and torch.compile(model) are architecturally incompatible.

• Either don't use torch.compile
• or, compile submodules, not the complete LightningModule

[Mo-Junyang]

Hi, @deependujha . I copy your code(before edited) and still get some errors.

I don't know what happened here

For your information, my packages version is as followed:
torch==2.6.0+cu126
lightning==2.5.2

Thanks for the recommendation. I am now successfully running the code by only compile the network but the LightningModule. Since the Pytorchlightning Document gave the exsample that compile the whold Module, I am wondering if the speed will be faster when compiling the whole module, including the optimizer.

[deependujha]

Thanks for checking and sharing the details.

I ran the earlier example on macOS (CPU/MPS), which is likely why it appeared to work there. CUDA tends to exercise stricter paths in torch.compile, so the issue shows up more reliably.

In principle, compiling the full LightningModule could enable more aggressive graph capture, but self.toggle_optimizer() dynamically changes requires_grad after compilation, which violates assumptions made by torch.compile / AOTAutograd. That’s what leads to the failures you’re seeing.

At the moment, compiling only the network modules (e.g. generator / discriminator) but not the LightningModule itself is the most reliable approach. I’m not aware of a clean way to make full-module compilation safe with optimizer toggling without changes upstream in PyTorch.