feat(training,rollout)!: Rollout Schedulers

[HCookie]

Closes #14

Rollout Schedulers

Expand the ways to describe rollout, and provide an interface to schedule updates

New default rollout config

rollout:
  _target_: anemoi.training.schedulers.rollout.stepped.EpochStepped
  minimum: 1
  maximum: 12
  # increase rollout every n epochs
  every_n_epochs: 1
  # Control the incrementing of the rollout window
  increment:
    step:
      0: 0
      200000: 1 # After 200k steps, increment by 1 every 1 epoch

Can step by epoch, step, and control the increment based on either the step or epoch.

Additonally, formally add random steppers.

Todo

• Integrate with the data loader
• Ensure that the randomness is seeded appropriately
• Randomness broadcast
• Ensure restartability
• Ability to change config

Tested

Tested with restart, and with change in config after restart.

---

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

---

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

---

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

[anaprietonem]

Started to go through the PR and left some comments! I still need to understand better some of the functionality so hope te questions makes sense. Thanks for this Harrison!

[FussyDuck]

All committers have signed the CLA.

[ssmmnn11]

nice, please see comments

[HCookie]

@ssmmnn11

can you check that the runs are still shuffled differently for every epoch?

They are not, the indices are the same for each epoch.
However, I have been thinking further on the need to reload, I think it may be possible to do it without it, but I'll need to figure out how to broadcast a value to all the dataloader workers, as this is why they do not update. Rebuilding them obviously fixed that but causes other issues.

I don't believe it is as simple as setting persistent_workers=False, as this also causes the issue with no shuffling.

[ssmmnn11]

persistent_workers=False means we have an initial seed and then just draw randon numbers from the generator, hence there should not be any repetition.

What we could do: before we did not have the epoch information in the dataset. This should now be the case and you could simply add the epoch to the random seed:

anemoi-core/training/src/anemoi/training/data/dataset.py

Line 224 in 83d72e1

base_seed = get_base_seed()

something like: base_seed = get_base_seed() + self.epoch * 100

though we should probably rename base_seed here to dataset_seed otherwise it might be confusing -> we call refer to the base seed elsewhere in the code.