Use torch.compile with optimizer and LightningModule

[Tomas542]

Since PyTorch 2.2.0 the new feature got introduced - compiling optimizer.step() (ref). What is the proper way to use it with LightningModule?

As far as I understand, I need to overwrite train_step as in Manual Optimization guide, like:

class MyModel(LightningModule):
    def __init__(self):
        super().__init__()
        # Important: This property activates manual optimization.
        self.automatic_optimization = False

    def training_step(self, batch, batch_idx):
        opt = self.optimizers()
        opt.zero_grad()
        loss = self.compute_loss(batch)
        self.manual_backward(loss)
        self._opt_step(opt)

    @torch.compile(fullgraph=False)
    def _opt_step(self, opt):
        opt.step()

Or I need to unwrap optimizer first?

class Model(LightningModule):
    def __init__(self):
        super().__init__()
        self.automatic_optimization = False
        ...

    def training_step(self, batch, batch_idx):
        optimizer = self.optimizers()
        optimizer = optimizer.optimizer
        ...
        self._opt_step(optimizer)

Also, question about compiling optimizer with scheduler - how to use it with above examples?