Feature/encoder spatial parallel

config/default_config.yml

@@ -66,6 +66,10 @@ num_register_tokens: 0
 
 healpix_level: 5
 
+# Split the encoder's HEALPix/location dimension across this many consecutive
+# distributed ranks through local assimilation.
+encoder_spatial_parallel_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

config/encoder_spatial_parallel_4.yml

@@ -0,0 +1,8 @@
+# (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.
+encoder_spatial_parallel_size: 4

config/encoder_spatial_parallel_8.yml

@@ -0,0 +1,8 @@
+# (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.
+encoder_spatial_parallel_size: 8