Lightning-AI / Lightning-AI/pytorch-lightning
EMA validation swap models before first update_parameters
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
With `EMAWeightAveraging` and a delayed start (e.g. update_starting_at_step: 1000), validation still swaps `pl_module` and `_average_model` on every val epoch (e.g. val every 100 steps), even while `self._average_model.n_averaged == 0`. The `_average_model` is only a deepcopy(pl_module) from setup() and has never been updated. Retrieval metrics before `update_starting_at_step` near 0.
Example for EMA init:
```yaml:
use_buffers: false
decay: 0.999
update_every_n_steps: 1
update_starting_at_step: 1000
update_starting_at_epoch: -1
```
Suggested fix: in `WeightAveraging` [on_validation_epoch_start](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/callbacks/weight_averaging.py#L216) / [on_validation_epoch_end](https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/callbacks/weight_averaging.py#L230), only call `_swap_models` when `self._average_model.n_averaged > 0`.
### What version are you seeing the problem on?
v2.6
### Reproduced in studio
_No response_
### How to reproduce the bug
```python
```
### Error messages and logs
_No response_
### Environment
_No response_
### More info
_No response_
cc @ethanwharris
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
Start in src/lightning/pytorch/callbacks/weight_averaging.py at on_validation_epoch_start and on_validation_epoch_end, then inspect how _average_model.n_averaged is initialized and updated. Reproduce the delayed-start configuration from the issue and verify validation does not swap the models before averaging begins, while swapping still occurs after n_averaged becomes positive.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100