config_forecasting.yml

# (C) Copyright 2025 WeatherGenerator contributors.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
#
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.

embed_orientation: "channels"
embed_unembed_mode: "block"
embed_dropout_rate: 0.1

ae_local_dim_embed: 2048
ae_local_num_blocks: 0
ae_local_num_heads: 16
ae_local_dropout_rate: 0.1
ae_local_with_qk_lnorm: True

ae_local_num_queries: 1
ae_local_queries_per_cell: False
ae_adapter_num_heads: 16
ae_adapter_embed: 128
ae_adapter_with_qk_lnorm: True
ae_adapter_with_residual: True
ae_adapter_dropout_rate: 0.1

ae_global_dim_embed: 2048
ae_global_num_blocks: 4
ae_global_num_heads: 32
ae_global_dropout_rate: 0.1
ae_global_with_qk_lnorm: True
# TODO: switching to < 1 triggers triton-related issues.
# See https://github.com/ecmwf/WeatherGenerator/issues/1050
ae_global_att_dense_rate: 1.0
ae_global_block_factor: 64
ae_global_mlp_hidden_factor: 2
ae_global_trailing_layer_norm: False

ae_aggregation_num_blocks: 0
ae_aggregation_num_heads: 32
ae_aggregation_dropout_rate: 0.1
ae_aggregation_with_qk_lnorm: True
ae_aggregation_att_dense_rate: 1.0
ae_aggregation_block_factor: 64
ae_aggregation_mlp_hidden_factor: 2

decoder_type: PerceiverIOCoordConditioning #  Main options PerceiverIOCoordConditioning or Linear
pred_adapter_kv: False
pred_self_attention: True
pred_dyadic_dims: False
pred_mlp_adaln: True
num_class_tokens: 0
num_register_tokens: 0

# number of steps offset applied to first target window; if set to zero and forecast_steps=0 then
# one is training an auto-encoder
fe_num_blocks: 16
fe_num_heads: 16
fe_dropout_rate: 0.1
fe_with_qk_lnorm: True
fe_layer_norm_after_blocks: [7]  # Index starts at 0. Thus, [3] adds a LayerNorm after the fourth layer
fe_impute_latent_noise_std: 1e-4
# currently fixed to 1.0 (due to limitations with flex_attention and triton)
forecast_att_dense_rate: 1.0

healpix_level: 5

rope_2D: False

with_mixed_precision: True
with_flash_attention: True
compile_model: False
with_fsdp: True
attention_dtype: bf16
mixed_precision_dtype: bf16
mlp_norm_eps: 1e-5
norm_eps: 1e-4

latent_noise_kl_weight: 0.0 # 1e-5
latent_noise_gamma: 2.0
latent_noise_saturate_encodings: 5 
latent_noise_use_additive_noise: False
latent_noise_deterministic_latents: True 

freeze_modules: ""
load_chkpt: {}

norm_type: "LayerNorm"

#####################################

streams_directory: "./config/streams/era5_1deg_forecasting/"
streams: ???

# type of zarr_store
zarr_store: "zip" # "zarr" for LocalStore, "zip" for ZipStore

general:

  # mutable parameters
  istep: 0
  rank: ???
  world_size: ???

  # local_rank, 
  # with_ddp,
  # data_path_*, 
  # model_path, 
  # run_path, 
  # path_shared_
  
  multiprocessing_method: "fork"

  desc: ""
  run_id: ???
  run_history: []

# logging frequency in the training loop (in number of batches)
train_logging:
  terminal: 10
  metrics: 20
  checkpoint: 250

# parameters for data loading
data_loading :

  num_workers: 12
  rng_seed: ???
  repeat_data_in_mini_epoch : False


# config for training
training_config:
  
  # training_mode: "masking", "student_teacher", "latent_loss"
  training_mode: ["masking"]

  num_mini_epochs: 64
  samples_per_mini_epoch: 4096
  shuffle: True

  start_date: 1979-01-01T00:00
  end_date: 2022-12-31T00:00

  time_window_step: 06:00:00
  time_window_len: 06:00:00
  
  learning_rate_scheduling :
    lr_start: 1e-6
    lr_max: 5e-5
    lr_final_decay: 2e-6
    lr_final: 0.0
    num_steps_warmup: 256
    num_steps_cooldown: 512
    policy_warmup: "cosine"
    policy_decay: "constant"
    policy_cooldown: "linear"
    parallel_scaling_policy: "sqrt"

  optimizer:
    grad_clip: 1.0
    weight_decay: 0.1
    log_grad_norms: False
    adamw :
      # parameters are scaled by number of DDP workers
      beta1 : 0.98125 # == 0.85 on 2 nodes x 4 gpus
      beta2 : 0.9875  # == 0.90 on 2 nodes x 4 gpus
      eps : 2e-08

  losses : {
    "physical": {
        type: LossPhysical,
        loss_fcts: { "mse": { }, },
        },
     "cosine_matching": {
        type: LossLatentCosineMatching,
        weight: 1.0,
        target_and_aux_calc: "Physical",
        loss_fcts: {
            "params": {
                cosine_low: 0.68,
                cosine_high: 0.78,
            },
        },
        },
    }
  
  model_input: {
    "forecasting" : {
      # masking strategy: "random", "healpix", "forecast"
      masking_strategy: "forecast",
      },
    }

  forecast :
      time_step: 06:00:00
      offset: 1
      num_steps: 3
      policy: "fixed"


# validation config; full validation config is merge of training and validation config
validation_config: 

  samples_per_mini_epoch: 256
  shuffle: False

  start_date: 2023-10-01T00:00
  end_date: 2023-12-31T00:00
  
  # whether to track the exponential moving average of weights for validation
  validate_with_ema: 
    enabled : True
    ema_ramp_up_ratio: 0.09
    ema_halflife_in_thousands: 600

  # parameters for validation samples that are written to disk
  output : {
    # number of samples that are written
    num_samples: 0,
    # write samples in normalized model space
    normalized_samples: False,
    # output streams to write; default all
    streams: null,
    }

  # run validation before training starts (mainly for model development)
  validate_before_training: False


# test config; full test config is merge of validation and test config
# test config is used by default when running inference

# Tags for experiment tracking
# These tags will be logged in MLFlow along with completed runs for train, eval, val
# The tags are free-form, with the following rules:
# - tags should be primitive types (strings, numbers, booleans). NO lists or dictionaries
# - tags should not duplicate existing config entries.
# - try to reuse existing tags where possible. MLFlow does not like having too many unique tags
# - do not use long strings in values (less than 20 characters is a good rule of thumb, we may enforce this in the future)
wgtags:
  # The name of the organization of the person running the experiment.
  # This may be autofilled in the future. Expected values are lowercase strings 
  # e.g. "ecmwf", "cmcc", "metnor", "jsc", "escience"
  org: null
  # The Github issue corresponding to this run (number such as 1234)
  # Github issues are the central point when running experiment and contain 
  # links to hedgedocs, code branches, pull requests etc.
  # It is recommended to associate a run with a Github issue.
  issue: null
  # The name of the experiment. This is a distinctive codename for the experiment campaign being run.
  # This is expected to be the primary tag for comparing experiments in MLFlow, along with the
  # issue number.
  # Expected values are lowercase strings with no spaces, just underscores:
  # Examples: "rollout_ablation_grid"  
  exp: null
  # *** Experiment-specific tags ***
  # All extra tags (including lists, dictionaries, etc.) are treated 
  # as strings by mlflow, so treat all extra tags as simple string key: value pairs.
  grid: null