fix: Kcrps fixes for Pydantic schemas #189
Conversation
anaprietonem
changed the title
| num_gpus_per_model: PositiveInt = Field(example=2) | ||
| "Number of GPUs per model." | ||
| # TODO(Ana): check why the read_group_size is not passed in the config | ||
| kwargs: dict[str, Any] = Field(default_factory=dict) |
There was a problem hiding this comment.
Just as note, we'll have to keep in mind that the kwargs will not be validated. We might want to look at including them in the config validation at a later stage.
| "Sanity check runs n batches of val before starting the training routine." | ||
| gradient_clip: GradientClip | ||
| "Config for gradient clipping." | ||
| forecaster: Any # TODO: Fix this | ||
| forecaster: Any # TODO(Simon): Fix this |
There was a problem hiding this comment.
Is there a plan to create a schema for the forecaster at some point?
| @@ -25,7 +27,7 @@ model: | |||
| grid_skip: 0 # Which of the input indices to use as residual connection, null if none. | |||
|
|
|||
| training: | |||
| task: anemoi.training.train.interpolator.GraphInterpolator | |||
| model_class: anemoi.training.train.interpolator.GraphInterpolator | |||
There was a problem hiding this comment.
I think model_task is clearer, would need changing across multiple configs.
| @@ -58,7 +58,7 @@ processor: | |||
| mlp_hidden_ratio: 4 # GraphTransformer or Transformer only | |||
| num_heads: 16 # GraphTransformer or Transformer only | |||
| cpu_offload: ${model.cpu_offload} | |||
| qk_norm: False | |||
| qk_norm: True | |||
There was a problem hiding this comment.
Need to check this with Simon and Helen. Looking at the settings, the qv_norm is kept to true for the processor in both the transformer (ens and single) and graph_transformer (ens). So that's why I thought this could be an error and needed correction. I will check with them before merging
Description
Type of Change
Issue Number
Code Compatibility
Code Performance and Testing
Dependencies
Documentation
Additional Notes