Refactor Callbacks

[HCookie]

• Split into seperate files
• Use list in config to add callbacks
• Provide legacy config enabled approach
• Fix ruff issues

New Usage

Set config.diagnostics.callbacks to a list of callback names to include

Closes #59

---

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

[FussyDuck]

All committers have signed the CLA.

[HCookie]

At the moment, this is the proposed refactor, I am yet to complete an exhaustive test of the changes

[JesperDramsch]

Great work, thank you for taking this on.

I was thinking that it might be nice to make this fully configurable through instantiate.

For example, no one is really using the stochastic weight averaging as far as I know, so having specific config entries for this is a bit of feature bloat.

Then the list of callbacks would just look like this:

callbacks:
  swa: _target_: pytorch_lightning.callbacks.stochastic_weight_avg.StochasticWeightAveraging
          swa_lr: 1e-4
          swa_epoch_start: 123
          annealing_epochs: 5
          annealing_strategy: cos
          device: null
  blabla: _target_: blabla_callback
             blabla: bla

This makes it more extensible and actually reduces some of or less used config entries.

Additionally, we can keep the standard callbacks, like model checkpoints as "permanent callback" (I don't think we have to make everything optional).

One idea I also had is that we could make a special list for "plot_callbacks" in the same style. Then we can easily keep the super convenient "plots.enabled = False" as a shortcut to disable them?

[JesperDramsch]

Hi @HCookie, thanks for taking on the callbacks!

It's already much better, great work on that. I think we can take the refactor even further and make the callbacks (almost?) fully modular, which would be incredible for future extensibility.

One comment regarding the file names. So far we haven't been using <xyz>-ing.py as language. Especially "checkpointing" would be confusing with activation checkpointing (although that is and will stay confusing honestly). Can we rename these please?

JesperDramsch is currently absent

[JPXKQX]

A long-awaited refactor. Thanks Harrison for the amazing work! I have tested with 1 GPU and 2 GPU (with num_gpus_per_model=2), and all plots in "detailed" option are produced as expected.

[sahahner]

As long as dataloader.validation_rollout=1, which is the default, this callback only increases the runtime without providing any additional plots. Should we move it into rollout_eval.yaml?

[HCookie]

We could provide a rollout plots configuration?
Addressed in f1d883f