Lightning-AI / Lightning-AI/pytorch-lightning
BaseFinetuning re-freezes already unfrozen modules when resuming from a checkpoint
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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