feat (training): Configurable weight averaging (exponential moving average (EMA), ...)

[jakob-schloer]

Weight averaging is a technique to improve model generalization by averaging model weights during training. In the new pytorch lightning version (2.6.0) different versions of weight averaging are now supported.

Here, I refactor the existing SWA implementation to support general weight averaging methods through the config. For example, EMA training can be used by adding the following to the config.training:

         weight_averaging:
            _target_: pytorch_lightning.callbacks.EMAWeightAveraging
            decay: 0.999

Note:

• I've deliberately deleted the previous version of weight initialization to keep the code cleaner.
• Old configs should still work, since defaults are null if weight_averaging does not exist, but this is still a breaking change if swa was used before.
• To test the code with lighting==2.6, this PR fix: progress bar with lightning==2.6.0 #739 needs to be merged first

As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/

By opening this pull request, I affirm that all authors agree to the Contributor License Agreement.

---

📚 Documentation preview 📚: https://anemoi-training--743.org.readthedocs.build/en/743/

---

📚 Documentation preview 📚: https://anemoi-graphs--743.org.readthedocs.build/en/743/

---

📚 Documentation preview 📚: https://anemoi-models--743.org.readthedocs.build/en/743/

[anaprietonem]

@jakob-schloer - How you're getting on with this PR? Are you planning to work on this soon?