Feature/encoder and decoder healpix cell parallelism

config/config_dinov2.yml

@@ -58,7 +58,12 @@ healpix_level: 5
 with_mixed_precision: True
 with_flash_attention: True
 compile_model: False
-with_fsdp: False
+distributed:
+  data_parallel:
+    with_ddp: False
+    with_fsdp: False
+  spatial_parallel:
+    size: 1
 ddp_find_unused_parameters: False
 attention_dtype: bf16
 mixed_precision_dtype: bf16
@@ -92,7 +97,6 @@ general:
   world_size: ???
 
   # local_rank, 
-  # with_ddp,
   # data_path_*, 
   # model_path, 
   # run_path, 

config/config_era5_georing_avhrr.yml

@@ -70,7 +70,12 @@ rope_2D: False
 with_mixed_precision: True
 with_flash_attention: True
 compile_model: False
-with_fsdp: True
+distributed:
+  data_parallel:
+    with_ddp: False
+    with_fsdp: True
+  spatial_parallel:
+    size: 1
 attention_dtype: bf16
 mixed_precision_dtype: bf16
 mlp_norm_eps: 1e-5
@@ -103,7 +108,6 @@ general:
   world_size: ???
 
   # local_rank, 
-  # with_ddp,
   # data_path_*, 
   # model_path, 
   # run_path, 

config/config_era5_georing_avhrr_masking.yml

@@ -70,7 +70,12 @@ rope_2D: False
 with_mixed_precision: True
 with_flash_attention: True
 compile_model: False
-with_fsdp: True
+distributed:
+  data_parallel:
+    with_ddp: False
+    with_fsdp: True
+  spatial_parallel:
+    size: 1
 attention_dtype: bf16
 mixed_precision_dtype: bf16
 mlp_norm_eps: 1e-5
@@ -103,7 +108,6 @@ general:
   world_size: ???
 
   # local_rank, 
-  # with_ddp,
   # data_path_*, 
   # model_path, 
   # run_path, 

config/config_forecasting.yml

@@ -70,7 +70,12 @@ rope_2D: False
 with_mixed_precision: True
 with_flash_attention: True
 compile_model: False
-with_fsdp: True
+distributed:
+  data_parallel:
+    with_ddp: False
+    with_fsdp: True
+  spatial_parallel:
+    size: 1
 attention_dtype: bf16
 mixed_precision_dtype: bf16
 mlp_norm_eps: 1e-5
@@ -103,7 +108,6 @@ general:
   world_size: ???
 
   # local_rank, 
-  # with_ddp,
   # data_path_*, 
   # model_path, 
   # run_path, 

config/config_forecasting_eerie.yml

@@ -70,7 +70,12 @@ rope_2D: False
 with_mixed_precision: True
 with_flash_attention: True
 compile_model: False
-with_fsdp: True
+distributed:
+  data_parallel:
+    with_ddp: False
+    with_fsdp: True
+  spatial_parallel:
+    size: 1
 attention_dtype: bf16
 mixed_precision_dtype: bf16
 mlp_norm_eps: 1e-5
@@ -103,7 +108,6 @@ general:
   world_size: ???
 
   # local_rank, 
-  # with_ddp,
   # data_path_*, 
   # model_path, 
   # run_path, 

config/config_forecasting_finetuning.yml

@@ -17,7 +17,6 @@ general:
   world_size: ???
 
   # local_rank, 
-  # with_ddp,
   # data_path_*, 
   # model_path, 
   # run_path, 
@@ -57,4 +56,3 @@ training_config:
       offset: 1
       num_steps: 8
       policy: "fixed"
-

config/config_jepa.yml

@@ -73,7 +73,12 @@ rope_2D: False
 with_mixed_precision: True
 with_flash_attention: True
 compile_model: False
-with_fsdp: False
+distributed:
+  data_parallel:
+    with_ddp: False
+    with_fsdp: False
+  spatial_parallel:
+    size: 1
 ddp_find_unused_parameters: False
 attention_dtype: bf16
 mixed_precision_dtype: bf16
@@ -106,7 +111,6 @@ general:
   world_size: ???
 
   # local_rank, 
-  # with_ddp,
   # data_path_*, 
   # model_path, 
   # run_path, 

config/config_jepa_finetuning.yml

@@ -29,7 +29,12 @@ with_step_conditioning: True # False
 with_mixed_precision: True
 with_flash_attention: True
 compile_model: False
-with_fsdp: False
+distributed:
+  data_parallel:
+    with_ddp: False
+    with_fsdp: False
+  spatial_parallel:
+    size: 1
 ddp_find_unused_parameters: False #True
 attention_dtype: bf16
 mixed_precision_dtype: bf16
@@ -54,7 +59,6 @@ general:
   world_size: ???
 
   # local_rank, 
-  # with_ddp,
   # data_path_*, 
   # model_path, 
   # run_path, 

config/config_jepa_forecasting_finetuning.yml

@@ -24,7 +24,12 @@ healpix_level: 5
 with_mixed_precision: True
 with_flash_attention: True
 compile_model: False
-with_fsdp: False
+distributed:
+  data_parallel:
+    with_ddp: False
+    with_fsdp: False
+  spatial_parallel:
+    size: 1
 ddp_find_unused_parameters: False #True
 attention_dtype: bf16
 mixed_precision_dtype: bf16
@@ -58,7 +63,6 @@ general:
   world_size: ???
 
   # local_rank, 
-  # with_ddp,
   # data_path_*, 
   # model_path, 
   # run_path, 

config/config_operan_georing_avhrr_forecasting_lowres.yml

@@ -71,7 +71,12 @@ rope_2D: False
 with_mixed_precision: True
 with_flash_attention: True
 compile_model: False
-with_fsdp: True
+distributed:
+  data_parallel:
+    with_ddp: False
+    with_fsdp: True
+  spatial_parallel:
+    size: 1
 attention_dtype: bf16
 mixed_precision_dtype: bf16
 mlp_norm_eps: 1e-5
@@ -105,7 +110,6 @@ general:
   world_size: ???
 
   # local_rank, 
-  # with_ddp,
   # data_path_*, 
   # model_path, 
   # run_path, 

config/config_physical_jepa.yml

@@ -68,7 +68,12 @@ healpix_level: 5
 with_mixed_precision: True
 with_flash_attention: True
 compile_model: False
-with_fsdp: False
+distributed:
+  data_parallel:
+    with_ddp: False
+    with_fsdp: False
+  spatial_parallel:
+    size: 1
 ddp_find_unused_parameters: False
 attention_dtype: bf16
 mixed_precision_dtype: bf16
@@ -101,7 +106,6 @@ general:
   world_size: ???
 
   # local_rank, 
-  # with_ddp,
   # data_path_*, 
   # model_path, 
   # run_path, 

config/default_config.yml

@@ -66,6 +66,17 @@ num_register_tokens: 0
 
 healpix_level: 5
 
+distributed:
+  data_parallel:
+    # Set at runtime from the launched world size.
+    with_ddp: False
+    with_fsdp: True
+    find_unused_parameters: True
+  spatial_parallel:
+    # Split the model's HEALPix/location dimension across this many consecutive
+    # distributed ranks.
+    size: 4
+
 # Use 2D RoPE instead of traditional global positional encoding
 # When True: uses 2D RoPE based on healpix cell coordinates (lat/lon)
 # When False: uses traditional pe_global positional encoding
@@ -74,7 +85,6 @@ 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
@@ -108,7 +118,6 @@ general:
   world_size: ???
 
   # local_rank, 
-  # with_ddp,
   # data_path_*, 
   # model_path, 
   # run_path, 

config/spatial_parallel_4.yml

@@ -0,0 +1,10 @@
+# (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.
+
+# Four consecutive ranks share one batch and each process one quarter of the
+# HEALPix/location dimension through local assimilation.
+distributed:
+  spatial_parallel:
+    size: 4

config/spatial_parallel_8.yml

@@ -0,0 +1,10 @@
+# (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.
+
+# Eight consecutive ranks share one batch and each process one eighth of the
+# HEALPix/location dimension through local assimilation.
+distributed:
+  spatial_parallel:
+    size: 8