Lightning-AI / Lightning-AI/pytorch-lightning

EMA validation swap models before first update_parameters

Open Beginner friendly
#21,724 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug ver: 2.6.x
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

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.