fix: Optimiser schema was ignoring kwargs

[HCookie]

Description

Allow extra inputs in optimiser schema for use in setting betas, and other keys. Before it was being ignored.

[icedoom888]

@anaprietonem Testing this right now.
Shoudln't all of these be patched as well?
https://github.com/search?q=repo%3Aecmwf%2Fanemoi-core%20PydanticBaseModel&type=code

Also why not revert back to inherit from BaseModel in anemoi.utils? What was the logic of the switch?

[icedoom888]

As expected, this fixed the Betas overwrite, but not all the other ones: normalizers, scalers, graph attributes and layer_kernels.
The should all be patched.

Comparison Fixed and No config_validation: https://mlflow.ecmwf.int/#/compare-runs?runs=[%22e638d811e39f41728168852c81ae2119%22,%22bbc6f6560e1c492c8190fb6f9bf6c4bb%22]&experiments=[%22109%22]

Coimparison Fixed and config_validation: https://mlflow.ecmwf.int/#/compare-runs?runs=[%2236480c41c7904cbcbfbcbfa0a83a5a6b%22,%22bbc6f6560e1c492c8190fb6f9bf6c4bb%22]&experiments=[%22109%22] (only diff is Betas)

[HCookie]

@icedoom888 Yes, likely all of those should be patched. I believe the reasoning here is to allow any kwargs, whereas the utils basemodel explicitly forbids that.

[icedoom888]

@HCookie I can confirm that inheriting from the anemoi.utils schema causes issues due to forbedden extra inputs.
I am now experimenting with applying this fix in all other instances of PydanticBaseModel ,

[anaprietonem]

@icedoom888 do you already have a list of the schemas that needs patching? as Harrison said some of this was related to the flexibility for passing kwargs, but this is not something that we may want to enable for each of the schemas we have.

[icedoom888]

Issue seems to be solved when patching all schemas:
https://mlflow.ecmwf.int/#/metric?runs=[%22e638d811e39f41728168852c81ae2119%22,%2236480c41c7904cbcbfbcbfa0a83a5a6b%22,%22bbc6f6560e1c492c8190fb6f9bf6c4bb%22,%22a3a86d94bc544bef972453ee7a29423b%22]&metric=%22train_mse_loss_step%22&experiments=[%22109%22]&plot_metric_keys=%5B%22train_mse_loss_step%22%5D&plot_layout={}&x_axis=step&y_axis_scale=linear&line_smoothness=1&show_point=false&deselected_curves=[]&last_linear_y_axis_range=[]

@anaprietonem The full list is:

• UnvalidatedBaseSchema in base_schema.py
• DataSchema in data.py
• DatasetSchema in dataloader.py
• InputSchema and LogsSchema in system.py
• OptimizerSchema in training.py

These are all the ones i patched in the latest run.

[JPXKQX]

We currently don’t log or print the optimizer configuration. It would be helpful to include this in the training logs (e.g. optimizer type and key hyperparameters).

Since we are changing the default value of optimizer.betas (when config_validatios is set to True), we should explicitly document this in the changelog to make the change visible to users.

[VeraChristina]

LGTM, thanks for fixing this!

[icedoom888]

Do it @HCookie