Hey @senarvi, this looks great!

I saw you already added support for saving and resuming which is great. There are many scenarios there (save every n steps, time-based, every epoch, etc) let's make sure we cover them all (for inspiration, we added quite a few tests here #20379)

we could still have different callbacks ("StepwiseAveragingCallback" and "EpochwiseAveragingCallback")

No I think it's better to have one with configurable averaging flags, more lightning-esque

Constructs the AveragedModel with use_buffers=True, so that an extra step is not needed for updating the batch normalization statistics. StochasticWeightAveraging performs an extra step in the end. Consequently the implementation is significantly more complex and it's difficult to make sure that it works in all cases. Should we add this as an option in this class too?

I think this is ok, but my doubt with forcing use_buffers to be true is what happens when a user has a module with buffers in it that are not meant to be averaged. I guess at that point they will probably be the same over time (e.g. the RoPE cache), but that's not really a guarantee.

Wdyt about this? I don't necessarily want to make the implementation more complex, so this is just for discussion.

Updates the average model after every step. StochasticWeightAveraging updates the average model after every epoch, and I recall that the original paper updated it only at certain points (the learning rate minima). I guess it would be nice to be able to select whether the average model will be updated after every step, after every epoch, or after certain epochs. Then we would need only one callback and we could remove the StochasticWeightAveraging callback, but would it make this class too complex?

It would be nice to make it configurable, and probably users will want to get to some minimum and then start averaging. The criteria to do so may be very bespoke, so maybe allowing the user to implement a custom hook to decide whether to start averaging or whether to average at a given step would be super handy. Otherwise I'm expecting users will train for some time, save a checkpoint, then reload with this callback added to the trainer and start averaging. Which is totally fine but it requires you to stop and resume.

Regarding removing the StochasticWeightAveraging callback, I don't necessarily see that happening. We have a pretty strong commitment to backward compatibility at this point, so keeping that in with a notice to just use this one will not hurt.

I think this is ok, but my doubt with forcing use_buffers to be true is what happens when a user has a module with buffers in it that are not meant to be averaged. I guess at that point they will probably be the same over time (e.g. the RoPE cache), but that's not really a guarantee.

That's a good point. I don't know what would be a good solution.

Updates the average model after every step. StochasticWeightAveraging updates the average model after every epoch, and I recall that the original paper updated it only at certain points (the learning rate minima). I guess it would be nice to be able to select whether the average model will be updated after every step, after every epoch, or after certain epochs. Then we would need only one callback and we could remove the StochasticWeightAveraging callback, but would it make this class too complex?

It would be nice to make it configurable, and probably users will want to get to some minimum and then start averaging. The criteria to do so may be very bespoke, so maybe allowing the user to implement a custom hook to decide whether to start averaging or whether to average at a given step would be super handy. Otherwise I'm expecting users will train for some time, save a checkpoint, then reload with this callback added to the trainer and start averaging. Which is totally fine but it requires you to stop and resume.

That's an interesting idea. We could have the user pass a function update_on_step(global_step) or update_on_epoch(epoch) that returns a boolean. After each optimizer step and after each epoch we would call the function to check whether we should update the average model.

It seems that AveragedModel will copy the current model parameters when called the first time, and update the average on subsequent calls. This means that the first average is computed when update_on_step() or update_on_epoch() returns True for the second time. I don't see a better alternative.

I checked how StochasticWeightAveraging does this and I think it doesn't work correctly. It only ever updates the average model parameters in on_train_epoch_start(), so the average is not updated after the last epoch. Just shows why I'd like to keep the logic as simple as possible.

Hi, I have a couple questions.

1. You added the on_validation_epoch_start and on_validation_epoch_end hooks to swap the weights, but shouldn't the same happen for test?
2. In my current workflow I have a separate script that does the model exporting to ONNX. It's short, and really the only Lightning specific thing is the MyLightningModule.load_from_checkpoint(...) method. Since the averaged weights are a part of the callback, I would have to instantiate the trainer for the weights to be loaded. And even then, I wouldn't have a function I could call to explicitly swap the weights (since _swap_weights is private and not really accessible). So, my question is, can we have some sort of an API, outside of the trainer, that can load the averaged weights instead of the regular weights? Perhaps adding some sort of a parameter to the load_from_checkpoint method?

Hi @cyanic-selkie

During training (stage=fit), the actual LightningModule is what we update using the optimizer (I call it the current model) and an AveragedModel is maintained in the background (I call it the average model).

I assume that validation is only called during training. Before and after validation we swap the current model and the average model, so the average model will be validated.

When saving a checkpoint, we save the average model parameters in the state_dict. So if you later load the checkpoint without WeightAveraging callback and run a test or export to ONNX, you will be using the average parameters.

When training ends, we copy the average model parameters to the current model. So if you run a test or export to ONNX after training, you will be using the average parameters.

That's the idea at least. I'm not confident that I have thought about every possible corner case. It would be great if you could test that it works in your case.

@senarvi Ah! Thanks for the clarification, I should've checked the code out more carefully. I tried your branch out on a quantization aware training enabled model with ONNX export at the end and everything is working beautifully! I hope this gets merged quickly.

The user can now provide either the update_on_step or the update_on_epoch argument. (In theory also both.) It should be a function that takes the step/epoch number and returns True if the average model should be updated at that point of time.

For example:

update_on_step = lambda x: x > 100

or

update_on_epoch = lambda x: x in (3, 5, 7)

Using update_on_epoch, SWA should be possible. I added one unit test for SWA.

I tested EMA in an actual learning task and it gave an improvement, so I'm starting to be more confident that this works.

I think the biggest question that is still left is whether it's a problem that we force use_buffers=True. It would be nice if we could provide the option to instead call update_bn() after training and we wouldn't have to duplicate any of that code. That function takes a data loader and iterates through the data. I can imagine that passing the Trainer's data loader might not work in all cases. We could also leave calling this function to the user.

StochasticWeightAveraging increments the number of epochs in on_fit_start() and during the extra epoch disables the backward pass. I could also copy the code from that class, but there are some details that I don't understand, and I'm not that excited of copying code that I don't fully understand.

@tchaton I think you contributed the StochasticWeightAveraging callback, maybe you have some insight?

Is there anything blocking this from being merged?

I marked this ready for review. There were no comments whether it's a problem that we force use_buffers=True. Would it make sense to merge this now and perhaps introduce such option later based on the feedback that we receive?

Codecov Report

Attention: Patch coverage is 94.68085% with 5 lines in your changes missing coverage. Please review.

Project coverage is 79%. Comparing base (831870a) to head (5deb0bb).

❗ There is a different number of reports uploaded between BASE (831870a) and HEAD (5deb0bb). Click for more details.

HEAD has 349 uploads less than BASE

Flag	BASE (831870a)	HEAD (5deb0bb)
cpu	105	27
python3.10	24	6
lightning_fabric	26	0
pytest	57	0
python	12	3
python3.12	10	3
python3.12.7	35	9
lightning	60	15
python3.11	24	6
gpu	4	0
pytorch2.1	12	6
pytorch_lightning	23	12
pytest-full	52	27
pytorch2.2.2	6	3
pytorch2.3	6	3
pytorch2.5	6	3
pytorch2.6	6	3
pytorch2.4.1	6	3
pytorch2.5.1	5	3
pytorch2.7	5	3

@@            Coverage Diff            @@
##           master   #20545     +/-   ##
=========================================
- Coverage      87%      79%     -8%     
=========================================
  Files         268      266      -2     
  Lines       23449    23488     +39     
=========================================
- Hits        20389    18475   -1914     
- Misses       3060     5013   +1953     

BTW: I think it's totally fine to merge this as is and open an issue to gather discussions about averaging buffers.

The other question I have (for the future) is related to fitting both models on GPU. It may make sense to give the ability to keep the AveragedModel on a different device (e.g. cpu) to keep the callback usable with larger models.

The other question I have (for the future) is related to fitting both models on GPU. It may make sense to give the ability to keep the AveragedModel on a different device (e.g. cpu) to keep the callback usable with larger models.

There's a device argument already, and actually the default is cpu - as with StochasticWeightAveraging.

Hi! Thanks for this great PR. The current implementation only leverages avg_fn argument should it also consider the in-place version multi_avg_fn ?

Hi! Thanks for this great PR. The current implementation only leverages avg_fn argument should it also consider the in-place version multi_avg_fn ?

I think we could just pass **averaged_model_kwargs. I'll look into it over the weekend.

@Borda should be good now

Sweet, EMA finally merged! Great work, @senarvi!

Congrats <3

One small caveat for folks to be aware of when using this callback: If your LightningModule reaches the end of your test epoch and then calls on_train_end(), the callback will swap your current (unaveraged) model weights with your EMA (averaged) model weights prior to (potentially) saving a final model checkpoint. If such a checkpoint is saved, it will contain your averaged model weights, not the unaveraged weights. This is fine for most downstream use cases, assuming the model is truly finished training. However, if you later decide to load such a checkpoint and finetune the model for additional training epochs, your training losses will likely look fine, but your validation/test performance might suffer.

I've seen this myself, where loading a checkpoint after it has reached its trainer.max_epochs=20000 and then manually training for longer (after setting a larger value of trainer.max_epochs=40000) resulted in normal training losses with no disruption, yet I noticed a huge drop in validation performance as shown in the image below. This is likely because, from then on, the callback is treating the averaged weights as the unaveraged model weights.

@amorehead based on your comment, I am going to assume this is not something that should lead to a code change but is it something that is worth documenting at least?

@SkafteNicki precisely

Hi @amorehead, I came across your comment and I am now concerned my weights are not correctly averaged when I decide to train for a few more epochs by increasing max_epochs. Are you saying that the averaged and non-averaged weights are not correctly loaded from the checkpoint when resuming training like this? If so isn't that something which needs fixing?

Hi, @philgzl. Yes, I've observed this behavior myself. However, I'm not sure it can be considered a bug, since reaching trainer.max_epochs should indicate to the callback that training has finished and that it can save the "final" (irreversible) checkpoint for production/inference. If you exceed the original max_epochs, you are then using the callback in an undefined manner.

Thanks. So if I understand correctly this is all because the non-averaged weights are not saved at the end of training.

This seems like a quite drastic behavior for just reducing the checkpoint size. Would it make sense to at least add an option to save the non-averaged weights, so users can resume training or inspect the non-averaged weights? I would personally use that option by default in all my projects.

Just to be clear, both the averaged and non-averaged weights are saved during training such that training can be correctly resumed in case of a crash, right?

What about the interaction with a ModelCheckpoint which tracks e.g. validation loss? Is a best checkpoint saved mid-training considered ready for "production", as in it does not contain the non-averaged weights? Or do I need to manually to load the averaged weights after loading such a checkpoint?

Would it make sense to at least add an option to save the non-averaged weights

I think so, perhaps setting the default to not save them to save space.

Both the averaged and non-averaged weights are saved during training such that training can be correctly resumed in case of a crash, right?

Yes, you can tell if you compare the size of your checkpoints mid-training and at the end of training. The latter should be smaller because it only contains the averaged weights.

Is a best checkpoint saved mid-training considered ready for "production", as in it does not contain the non-averaged weights?

No, mid-training checkpoints contain both the averaged weights (n.b., stored in the default model state_dict key of the checkpoint) and non-averaged weights (n.b., stored in a separate, dedicated key). As such, these mid-training checkpoints will be larger than necessary for downstream deployment without manual editing. However, it would be impossible to correctly resume training without saving both copies of the weights, so this is a necessary tradeoff.

Do I need to manually load the averaged weights after loading such a checkpoint?

The averaged weights will be loaded by default via the model checkpoint's state_dict key (n.b., which PyTorch assumes are your non-averaged, regular weights). That's why it's necessary when resuming training to reuse the same WeightAveragingCallback, because it will run logic to correctly extract both the averaged and non-averaged weights from your model checkpoints when resuming. Otherwise, if you directly run torch.load_state_dict (or Lightning's equivalent) on your checkpoints, you will be loading the averaged weights.

Great, thanks a lot for the clear explanation @amorehead!