#21901·pytorch-lightning

BaseFinetuning re-freezes already unfrozen modules when resuming from a checkpoint

Author: yenturCreated Aug 14, 2026Updated Aug 14, 2026

Labelsver: 2.7.x

Bug description

Resuming a BaseFinetuning run from a checkpoint silently re-freezes the modules that were already unfrozen before the interruption. The affected parameters stay in the optimizer with their correct learning rate, but requires_grad is False, so they receive no gradients and never update again. There is no warning, and LearningRateMonitor still reports a learning rate for the group.

The cause is a half-completed restore:

BaseFinetuning.setup() calls self.freeze_before_training(pl_module) unconditionally, on a resumed run as well as a fresh one. This has to happen in setup() because it must run before configure_optimizers().
BaseFinetuning.on_fit_start() then restores optimizer.param_groups from _internal_optimizer_metadata, but it only restores group membership. It never restores param.requires_grad.

BackboneFinetuning.finetune_function() unfreezes only on the exact epoch epoch == unfreeze_backbone_at_epoch; for later epochs it just rescales the learning rate. So once the unfreeze epoch has passed, nothing puts requires_grad back, and the backbone stays frozen for the rest of training.

What version are you seeing the problem on?

master

How to reproduce the bug

python

import torch
from torch import nn

from lightning.pytorch import Trainer
from lightning.pytorch.callbacks import BackboneFinetuning, ModelCheckpoint
from lightning.pytorch.demos.boring_classes import BoringModel


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

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

    def configure_optimizers(self):
        # only the head to start with; the callback adds the backbone when it unfreezes it
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def fit(max_epochs, callbacks, ckpt_path=None):
    model = Model()
    Trainer(
        default_root_dir="/tmp/bb",
        max_epochs=max_epochs,
        limit_train_batches=2,
        limit_val_batches=0,
        num_sanity_val_steps=0,
        enable_progress_bar=False,
        enable_model_summary=False,
        logger=False,
        callbacks=callbacks,
    ).fit(model, ckpt_path=ckpt_path)
    return model


ckpt = ModelCheckpoint(dirpath="/tmp/bb", save_last=True)

straight = fit(4, [BackboneFinetuning(unfreeze_backbone_at_epoch=1)])
print("uninterrupted, backbone trainable:", straight.backbone.weight.requires_grad)

fit(2, [BackboneFinetuning(unfreeze_backbone_at_epoch=1), ckpt])
resumed = fit(4, [BackboneFinetuning(unfreeze_backbone_at_epoch=1)], ckpt_path=ckpt.last_model_path)
print("resumed,       backbone trainable:", resumed.backbone.weight.requires_grad)

Error messages and logs

uninterrupted, backbone trainable: True
resumed,       backbone trainable: False

Tracking requires_grad and whether the backbone weights actually change, per epoch, over the same 4 epochs:

 epoch                uninterrupted           resumed at epoch 2
       requires_grad  weights_moved requires_grad  weights_moved
     0         False          False         False          False
     1          True           True          True           True
     2          True           True         False          False
     3          True           True         False          False

With a staged schedule that unfreezes block i at epoch i, resuming at epoch 2 re-freezes the blocks unfrozen in the earlier epochs:

uninterrupted requires_grad per block: [True, True, True, False]
resumed       requires_grad per block: [False, False, True, False]

Environment

Current environment

#- PyTorch Lightning Version: 2.6.2 (master, fcef40451)
#- PyTorch Version: 2.13.0
#- Python version: 3.11
#- OS: macOS 15 (arm64)
#- CUDA/cuDNN version: none, CPU
#- How you installed Lightning: source, editable

More info

tests/tests_pytorch/callbacks/test_finetuning_callback.py::test_callbacks_restore_backbone already exercises this resume path but makes no assertion after the second fit(), which is why the suite stays green.

Note that the bug only shows when the unfrozen modules were added to the optimizer as a new param group, which is the pattern the BackboneFinetuning docstring shows. If configure_optimizers() already passes every parameter to the optimizer, unfreeze_and_add_param_group() drops them in filter_on_optimizer() and no new group is recorded, so there is nothing in the saved metadata to restore from.

I have a fix and will open a PR shortly.

Source: Lightning-AI/pytorch-lightning

View original on GitHub View discussion on GitHub