Lightning-AI / Lightning-AI/pytorch-lightning

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

Open
#21,901 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

ver: 2.7.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### 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.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Read BaseFinetuning.setup() and on_fit_start(), then inspect the existing resume coverage in tests/tests_pytorch/callbacks/test_finetuning_callback.py::test_callbacks_restore_backbone. Reproduce the two-stage BackboneFinetuning run and extend the test so resumed modules remain trainable and continue updating after the unfreeze epoch.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, testing-qa
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.