diff --git a/config/config_dinov2.yml b/config/config_dinov2.yml index 2ec2c67015..849243c409 100644 --- a/config/config_dinov2.yml +++ b/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, diff --git a/config/config_era5_georing_avhrr.yml b/config/config_era5_georing_avhrr.yml index 2472d12a30..e3a9dd241b 100644 --- a/config/config_era5_georing_avhrr.yml +++ b/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, diff --git a/config/config_era5_georing_avhrr_masking.yml b/config/config_era5_georing_avhrr_masking.yml index 09447c0e93..189f60eb19 100644 --- a/config/config_era5_georing_avhrr_masking.yml +++ b/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, diff --git a/config/config_forecasting.yml b/config/config_forecasting.yml index 3491da6df7..b72e784de0 100644 --- a/config/config_forecasting.yml +++ b/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, diff --git a/config/config_forecasting_eerie.yml b/config/config_forecasting_eerie.yml index 88d4a9c1a8..17bd3a0fb7 100644 --- a/config/config_forecasting_eerie.yml +++ b/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, diff --git a/config/config_forecasting_finetuning.yml b/config/config_forecasting_finetuning.yml index d3b4c24411..3898569c1c 100644 --- a/config/config_forecasting_finetuning.yml +++ b/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" - diff --git a/config/config_jepa.yml b/config/config_jepa.yml index 197f3574ab..2a163c622a 100644 --- a/config/config_jepa.yml +++ b/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, diff --git a/config/config_jepa_finetuning.yml b/config/config_jepa_finetuning.yml index c4921d5cd3..434c991c29 100644 --- a/config/config_jepa_finetuning.yml +++ b/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, diff --git a/config/config_jepa_forecasting_finetuning.yml b/config/config_jepa_forecasting_finetuning.yml index 36de995bf4..d734d4d27b 100644 --- a/config/config_jepa_forecasting_finetuning.yml +++ b/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, diff --git a/config/config_operan_georing_avhrr_forecasting_lowres.yml b/config/config_operan_georing_avhrr_forecasting_lowres.yml index 5115f6310b..3084fcc68d 100644 --- a/config/config_operan_georing_avhrr_forecasting_lowres.yml +++ b/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, diff --git a/config/config_physical_jepa.yml b/config/config_physical_jepa.yml index a0f35ca029..7392e00bc6 100644 --- a/config/config_physical_jepa.yml +++ b/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, diff --git a/config/default_config.yml b/config/default_config.yml index 53f1ffc582..d116166d3b 100644 --- a/config/default_config.yml +++ b/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, diff --git a/config/spatial_parallel_4.yml b/config/spatial_parallel_4.yml new file mode 100644 index 0000000000..a14fcb29a7 --- /dev/null +++ b/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 diff --git a/config/spatial_parallel_8.yml b/config/spatial_parallel_8.yml new file mode 100644 index 0000000000..a54e542343 --- /dev/null +++ b/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 diff --git a/docs/decoder_spatial_parallelism.md b/docs/decoder_spatial_parallelism.md new file mode 100644 index 0000000000..006b98f8e3 --- /dev/null +++ b/docs/decoder_spatial_parallelism.md @@ -0,0 +1,421 @@ +# Decoder spatial parallelism + +This document describes the HEALPix cell-parallel decoder introduced by commits: + +| Commit | Purpose | +| --- | --- | +| `d753973b` | Parallelize decoder work by HEALPix cell | +| `6f087772` | Fix the Linear path, synchronize NaN handling, and remove a redundant lens collective | + +The decoder feature extends the topology and contiguous cell ownership described in +[`encoder_spatial_parallelism.md`](encoder_spatial_parallelism.md). It deliberately reuses +`distributed.spatial_parallel.size`; there is no separate decoder-parallel configuration. + +## Motivation + +Before this change, every rank in an encoder spatial group decoded every target coordinate. +The ranks consumed the same sample and held the same global latent tensor, so this duplicated +coordinate embedding, prediction-engine, and prediction-head work across the group. + +Decoder spatial parallelism distributes that work by output HEALPix cell: + +1. every spatial rank retains the encoder's contiguous cell range; +2. target coordinates are restricted to cells in that range; +3. each owned cell receives its one-ring HEALPix neighborhood: itself and eight neighbors; +4. the rank decodes only its local target coordinates; +5. variable-length local predictions are gathered differentiably; +6. gathered rank shards are restored to sample-major, global cell order. + +The complete post-forecast latent tensor is still replicated on each spatial rank. The feature +reduces decoder computation and intermediate activations; it does not shard the forecast engine +or its output tensor. + +## Shared topology and configuration + +The decoder reuses the encoder's `SpatialParallelContext` during model creation: + +```text +model.spatial_parallel = encoder.spatial_parallel +``` + +The relevant configuration remains: + +```yaml +distributed: + spatial_parallel: + size: 4 +``` + +For HEALPix level 5 and four spatial ranks: + +| Spatial rank | Decoder-owned cells | Cells/rank | +| ---: | ---: | ---: | +| 0 | `[0, 3072)` | 3072 | +| 1 | `[3072, 6144)` | 3072 | +| 2 | `[6144, 9216)` | 3072 | +| 3 | `[9216, 12288)` | 3072 | + +Ranks in one spatial group consume the same batch. Different spatial groups remain data-parallel +and must not exchange decoder results with one another. For that reason, decoder collectives +always use `spatial_parallel.group`, not the default world process group. + +With spatial size one, the decoder retains the original single-rank behavior and performs no +prediction gather. + +## End-to-end data flow + +For each decoded forecast step and output stream: + +```text +global latent after assimilation/forecasting + [batch, global_cells, queries_per_cell, latent_dim] + │ + ├─ select each owned cell's self-plus-eight-neighbor indices + ▼ +rank-local packed neighborhoods + [batch × local_cells × 9 × queries_per_cell, latent_dim] + │ + ├─ slice globally constructed target coordinates by owned cell range + ▼ +rank-local packed target coordinates and per-cell lengths + │ + ▼ +coordinate embedding → prediction engine → prediction head + │ + ▼ +rank-local variable-length predictions + │ + ├─ pad to the largest rank-local prediction count + ├─ autograd-aware all-gather within the spatial group + ├─ remove padding using globally known per-cell target lengths + ▼ +sample-major predictions in global HEALPix cell order +``` + +This process is repeated independently for every stream and every step on which the model calls +`predict_decoders()`. + +## Nine-cell HEALPix neighborhoods + +`ModelParams.hp_nbours` has shape: + +```text +[num_healpix_cells, 9] +``` + +For each cell, column zero is the cell itself and the remaining columns are its eight one-ring +neighbors. HEALPix cells with a missing neighbor use the center cell as the replacement, so every +row has exactly nine valid indices. + +`select_healpix_neighborhood_shard()` validates the global lookup shape, restricts its rows to +the rank-owned interval, and indexes the globally available latent tensor. The packed result +preserves this logical order: + +```text +sample + → owned cell + → 9 neighborhood cells + → query token +``` + +The implementation currently gathers the global latent at the encoder boundary and executes the +forecast engine globally on every spatial rank. Consequently, neighborhood construction is a +local tensor selection; it does not require a halo exchange between ranks. + +For `PerceiverIOCoordConditioning`, each local target cell attends to the corresponding packed +nine-cell neighborhood. The default configuration uses one latent query per HEALPix cell. + +## Target-coordinate sharding + +Targets remain globally constructed on every rank in a spatial group. For one stream and forecast +step, each sample provides: + +```text +target_coords = packed target coordinates in cell-major order +target_coords_lens = number of target coordinates in every global cell +``` + +Target counts vary by sample, stream, and cell. Directly slicing the packed coordinate tensor by +cell index would therefore be incorrect. `select_packed_cell_shard()` expands the local cell mask +by the per-cell lengths and returns complete coordinate segments for: + +```text +[spatial_parallel.cell_start, spatial_parallel.cell_end) +``` + +The local lengths retain `(sample, local_cell)` order and are passed to the variable-length +prediction engine. + +## Decoder paths + +### PerceiverIOCoordConditioning + +The default and shipped configurations use `PerceiverIOCoordConditioning`. Its local prediction +engine receives: + +- locally owned target-coordinate embeddings; +- the packed nine-cell latent neighborhoods; +- one target length per local cell; +- one latent-neighborhood length per local cell; +- the local raw target coordinates used for conditioning. + +The stream-specific ensemble prediction head then maps prediction tokens to physical channels. + +### Linear + +`Linear` remains a documented decoder option. `BilinearDecoder` expects one latent row for every +cell length entry, so spatial execution passes only: + +```text +tokens[:, spatial_parallel.cell_start:spatial_parallel.cell_end] +``` + +rather than the global latent tensor. The current Linear implementation requires: + +```yaml +ae_local_num_queries: 1 +``` + +and raises `NotImplementedError` for any other query count. + +## Empty local target domains and FSDP + +A stream can have global targets while one spatial rank owns no target coordinates. That rank +must not skip FSDP-wrapped decoder modules while its peers enter them, because mismatched module +execution can deadlock forward or backward collectives. + +For an empty local target domain, the decoder therefore: + +1. creates one zero-valued dummy coordinate with the stream's coordinate width; +2. assigns that coordinate to one local cell for the variable-length call; +3. executes coordinate embedding, the prediction engine, and the prediction head; +4. discards the dummy prediction with an empty slice; +5. participates in the prediction gather with a valid empty local shard. + +This preserves module-call order without adding a physical prediction. + +If a stream has no target coordinates anywhere in the global batch, every rank skips that stream +before entering the local decoder path. + +## Group-consistent NaN handling + +After target-coordinate embedding, every spatial rank computes a local NaN flag. The flags are +reduced with `MAX` over the decoder spatial group before any rank enters the prediction engine or +prediction head. + +If any rank reports a NaN, every rank raises `FloatingPointError`. Continuing with an empty result +would be unsafe because: + +- ranks would enter FSDP-wrapped modules in different orders; +- gathered prediction lengths would no longer match the global target lengths; +- a later `torch.split()` could fail or silently misalign outputs. + +## Variable-length prediction gather + +### Why padding is required + +`torch.distributed.nn.functional.all_gather` requires the same tensor shape on every rank, while +the number of target coordinates per rank can differ. Each rank derives every shard's packed +prediction count from the globally available `target_coords_lens`, then pads its local prediction +tensor to the largest count. + +No collective is required for target lengths: every rank already has the same global target lens, +and `split_cell_lens_by_shard()` reconstructs the contiguous per-rank slices locally. + +The padded predictions are gathered with the autograd-aware collective: + +```python +gathered_pred = all_gather( + pred, + group=spatial_parallel.group, +) +``` + +Using the functional distributed gather preserves gradients from the reconstructed global loss +to the decoder computation on the rank that owned each target coordinate. + +### Why gathered tensors cannot simply be flattened + +Each rank packs all local samples before the gather: + +```text +rank 0: [sample 0 local cells, sample 1 local cells, ...] +rank 1: [sample 0 local cells, sample 1 local cells, ...] +``` + +The gather returns rank-major order. Flattening it would produce: + +```text +rank 0 sample 0 +rank 0 sample 1 +rank 1 sample 0 +rank 1 sample 1 +``` + +The output and loss code expect sample-major order: + +```text +sample 0 rank 0 cells +sample 0 rank 1 cells +sample 1 rank 0 cells +sample 1 rank 1 cells +``` + +Flattening would also retain rank padding. It is only sufficient in the special case of local +batch size one with equal prediction counts on every rank. + +### Reassembly + +`reassemble_packed_cell_shards()` uses each rank's per-cell target lengths to: + +1. calculate the valid prediction segment for each `(rank, sample)` pair; +2. remove padding; +3. concatenate rank segments for one sample in spatial-rank order; +4. concatenate the reconstructed samples in batch order. + +Because spatial ranks own consecutive cell intervals, spatial-rank order is also global HEALPix +cell order. The reconstructed tensor has shape: + +```text +[ensemble_size, total_global_target_coordinates, output_channels] +``` + +It is finally split by the original per-sample coordinate counts and stored in `ModelOutput`. + +## Autoregressive rollout behavior + +The decoder is invoked inside the loop over `batch.get_output_idxs()`. With: + +```yaml +forecast: + num_steps: 2 + offset: 1 +``` + +the output indices are `[1, 2]`. The forecast engine advances the latent state and the decoder +runs at both steps. Therefore the prediction gather also occurs once per decoded stream at each +of those two steps. + +When pushforward training is enabled, intermediate steps advance the forecast state without +decoding. Only the final gradient-bearing step executes the decoder and its gather. + +## What is local and what remains global + +| Stage/data | Spatially sharded? | +| --- | --- | +| Target reads and target-coordinate construction | No | +| Post-encoder/global-assimilation latent | No | +| Forecast engine and forecast latent | No | +| Selection of target coordinates by output cell | Yes | +| Selection of nine-cell neighborhoods | Yes | +| Coordinate embedding | Yes | +| Prediction engine | Yes | +| Prediction head | Yes | +| Reconstructed physical prediction | No | +| Physical loss after reconstruction | No | + +The principal benefit is reduced decoder computation and activation memory. The principal added +cost is one differentiable prediction gather per decoded stream and forecast step, plus a small +NaN-flag reduction. + +## Ordering guarantees + +The implementation relies on these invariants: + +1. all ranks in a spatial group process the same batch and target tensors; +2. decoder ownership exactly matches encoder ownership; +3. each rank owns a consecutive range of global nested HEALPix cell IDs; +4. target coordinates are packed in `(sample, cell)` order; +5. local coordinate selection preserves increasing global cell order; +6. process-group rank order matches cell-range order; +7. per-rank lengths and predictions use identical packing order; +8. reassembly removes padding before restoring sample-major order. + +Violating any of these assumptions can produce a correct gather shape with incorrectly aligned +predictions, so ordering tests are as important as shape tests. + +## Runtime verification + +For HEALPix level 5, spatial size 4, local batch size `B`, and one query per cell, verify: + +```text +global latent cells: 12288 +local decoder cells/rank: 3072 +neighborhoods per local cell: 9 +packed neighborhood rows/rank: B × 3072 × 9 +``` + +For each stream and decoded step: + +- the sum of local target counts across spatial ranks equals the global target count; +- each local coordinate belongs to the rank's cell interval; +- gathered output length equals the sum of the original per-sample coordinate lengths; +- predictions have the same ordering and values as spatial size one, within numerical tolerance; +- backward reaches decoder parameters and the owning rank's local prediction graph. + +Peak-memory comparisons should record every spatial rank. Sparse target streams can produce +imbalanced decoder work and memory even though cell ownership is equal. + +## Troubleshooting + +### Hang in decoder forward or backward + +Verify that: + +- all group ranks process the same stream and forecast step; +- empty local target ranks execute the dummy path; +- NaN handling is reduced over the spatial group; +- decoder collectives use the encoder spatial process group; +- no rank independently skips a globally non-empty target stream. + +### Prediction gather shape mismatch + +Check: + +- identical ensemble size and channel count on every rank; +- correct global `target_coords_lens` on every rank; +- local prediction length equals the sum of the rank's local lens slice; +- local padding uses the maximum valid prediction count across the group; +- every rank enters one gather per stream and decoded step. + +### Output ordering mismatch for batch size greater than one + +Do not flatten the rank-major gather result. Use `reassemble_packed_cell_shards()` to transpose +the logical rank/sample order and remove padding. + +### Linear decoder failure + +Verify that the Linear path receives only the local latent cell slice and that +`ae_local_num_queries` is one. + +### No decoder speedup + +Confirm that target coordinates exist across the global domain and inspect target counts per +rank. The forecast engine and global latent are still replicated, so workloads dominated by +forecasting or global assimilation will not scale with decoder spatial size. + +## File-level implementation map + +| File | Responsibility | +| --- | --- | +| `config/default_config.yml` | Shared `distributed.spatial_parallel` settings and decoder type | +| `src/weathergen/model/model.py` | Local neighborhood/target selection, decoder execution, NaN synchronization, and prediction gather | +| `src/weathergen/model/spatial_parallel.py` | Packed cell selection, neighborhood selection, lens splitting, and prediction reassembly | +| `src/weathergen/model/encoder.py` | Spatial group and contiguous ownership reused by the decoder | +| `src/weathergen/utils/distributed.py` | Shared spatial topology context, validation, and process-group construction | +| `tests/test_spatial_parallel.py` | Neighborhood selection, lens slicing, and result-order tests | + +## Review checklist + +- [ ] The decoder reuses `distributed.spatial_parallel.size` and the shared spatial process group. +- [ ] Cell ownership is identical in the encoder and decoder. +- [ ] Every local cell receives exactly nine neighborhood indices. +- [ ] Target-coordinate selection respects variable per-cell lengths. +- [ ] Linear decoding uses only local latent cells and one query per cell. +- [ ] Empty local target ranks enter all FSDP-wrapped decoder modules. +- [ ] NaN detection fails collectively before prediction modules. +- [ ] Prediction gathering remains autograd-aware. +- [ ] Target lens are derived locally instead of gathered redundantly. +- [ ] Padding is removed before output reconstruction. +- [ ] Multi-sample batches are restored to sample-major order. +- [ ] Every decoded autoregressive step performs matching collectives on all group ranks. diff --git a/docs/encoder_spatial_parallelism.md b/docs/encoder_spatial_parallelism.md new file mode 100644 index 0000000000..413a9d0893 --- /dev/null +++ b/docs/encoder_spatial_parallelism.md @@ -0,0 +1,828 @@ +# Encoder spatial parallelism + +This document describes the encoder spatial-parallel implementation introduced by +[pkuyyj/WeatherGenerator PR #1](https://github.com/pkuyyj/WeatherGenerator/pull/1). +It covers the two feature commits: + +| Commit | Purpose | +| --- | --- | +| `15dcff5a` | Add HEALPix encoder spatial parallelism | +| `db4b477d` | Build encoder inputs on rank-local HEALPix domains | + +Upstream evaluation changes and branch-synchronization-only changes in the PR history are +outside the scope of this document. + +## Motivation + +WeatherGenerator receives several heterogeneous input streams. A stream can be globally +regular, such as ERA5, or contain a variable number of observations, such as +`METOP_ABC_AVHRR_IASI`. Before this change, every distributed rank embedded every source +token and performed local assimilation for every HEALPix cell. + +Embedding and local assimilation are spatially independent until the local cell +representations are projected into the global latent representation. The implementation +therefore distributes these stages over a HEALPix domain: + +1. ranks in one spatial group consume the same training sample; +2. every rank owns a disjoint set of complete HEALPix cells; +3. each of the nine streams is filtered and tokenized independently for that local domain; +4. embedding, local assimilation, and local-to-global projection operate on local data; +5. fixed-size per-cell latent representations are gathered in global cell order; +6. query aggregation and global assimilation continue with the reconstructed global tensor. + +This reduces source-token and local-encoder activation memory. It does not shard the complete +model or the complete training step. + +## Terminology + +| Term | Meaning | +| --- | --- | +| `world_size` | Total number of distributed ranks | +| `distributed.spatial_parallel.size` | Number of ranks cooperating on one encoder input | +| spatial group | Consecutive ranks that consume the same sample and partition its HEALPix domain | +| data-parallel rank | Index of a spatial group in the global job | +| spatial rank | Rank index within a spatial group | +| data HEALPix level | The configured `healpix_level`, normally level 5 | +| packed tokens | Variable-length tensor containing only existing stream tokens | +| dense cell tensor | Fixed-size tensor with one position for every owned HEALPix cell | + +For global rank `r` and spatial-parallel size `S`: + +```text +data_parallel_rank = r // S +spatial_rank = r % S +``` + +Ranks `[gS, gS + 1, ..., gS + S - 1]` form spatial group `g`. + +For example, with eight total ranks and `S = 4`: + +```text +spatial group 0: global ranks [0, 1, 2, 3] +spatial group 1: global ranks [4, 5, 6, 7] +``` + +The four ranks in each group consume the same sample. The two groups remain data-parallel +with respect to each other. + +## HEALPix partitioning + +### Number of cells + +At HEALPix level `L`, the number of cells is: + +```text +N(L) = 12 × 4^L +``` + +At the default data level `L = 5`: + +```text +N(5) = 12 × 4^5 = 12,288 +``` + +### Contiguous equal ownership + +The configured data-level cells are divided directly into equal, consecutive rank-local +intervals: + +```text +cells_per_rank = N(L) / S + +cell_start = spatial_rank × cells_per_rank +cell_end = cell_start + cells_per_rank +``` + +`N(L)` must be divisible by the spatial-parallel size. For level 5 and four spatial ranks: + +| Spatial rank | Level-5 cells | Number of level-5 cells | +| ---: | ---: | ---: | +| 0 | `[0, 3072)` | 3072 | +| 1 | `[3072, 6144)` | 3072 | +| 2 | `[6144, 9216)` | 3072 | +| 3 | `[9216, 12288)` | 3072 | + +Common group sizes produce: + +| Spatial ranks | Level-5 cells/rank | +| ---: | ---: | +| 4 | 3072 | +| 8 | 1536 | +| 16 | 768 | +| 32 | 384 | + +### Why consecutive ranges are valid + +The input mapping uses nested HEALPix ordering: + +```python +ang2pix(2**healpix_level, theta, phi, nest=True) +``` + +Nested ordering gives every data-level cell a stable integer ID. Assigning consecutive ID +intervals means that concatenating rank-local tensors in spatial-rank order reconstructs the +original global cell order. No coarser parent-cell constraint is required. + +## Distributed topology and process groups + +`get_spatial_parallel_size()` validates that: + +- the configured spatial size is at least one; +- it does not exceed `world_size`; +- `world_size` is divisible by the spatial size. + +`get_spatial_parallel_group()` creates process groups from consecutive global ranks. +All global ranks create the groups in the same order, and each process caches the group that +contains it. + +`SpatialParallelContext` centralizes the resulting process group, spatial rank and size, +effective DDP rank and world size, and the rank-local HEALPix interval. The data sampler, +encoder, and decoder model consume this shared context instead of independently deriving or +copying the same topology fields. + +With spatial size one, no additional process group is created and the feature reduces to the +single-rank compatibility behavior. + +## End-to-end data flow + +The source path can be summarized as: + +```text +source readers + │ + ├─ same sample on every rank in a spatial group + │ + ▼ +map every source location to nested HEALPix cell ID + │ + ├─ filter cell_start <= cell_id < cell_end independently for each stream + │ + ▼ +construct local stream tokens and local per-cell token counts + │ + ▼ +embed every stream independently on its owning rank + │ + ▼ +scatter stream-major embeddings into local cell-major packed order + │ + ▼ +local assimilation + │ + ▼ +local-to-global projection + │ + ▼ +restore fixed-size dense local-cell tensor + │ + ▼ +autograd-aware all-gather in spatial-rank order + │ + ▼ +global cell tensor [0, N(healpix_level)) + │ + ▼ +query aggregation and global assimilation +``` + +### What is local and what remains global + +| Stage/data | Spatially sharded? | +| --- | --- | +| Source storage reads | No; each spatial rank currently reads the same source sample | +| Source coordinate-to-HEALPix mapping | Computed independently on every spatial rank | +| Source filtering and token construction | Yes | +| Per-stream embedding | Yes | +| Local assimilation | Yes | +| Local-to-global projection | Yes | +| Dense per-cell latent after gather | No | +| Query aggregation | No | +| Global assimilation | No | +| Target construction | No | +| Prediction | Yes; see [`decoder_spatial_parallelism.md`](decoder_spatial_parallelism.md) | +| Physical loss after prediction gather | No | +| Model parameters/FSDP state | Controlled separately by FSDP | + +The implementation reduces GPU source-token and activation memory, but does not yet perform +distributed source I/O. Raw source reads and the coordinate mapping are replicated inside a +spatial group. + +## Rank-local stream construction + +### Filtering mask + +Every stream is treated independently. After converting its coordinates to HEALPix IDs, the +rank-local point mask is: + +```python +local_domain_mask = (cell_ids >= cell_start) & (cell_ids < cell_end) +``` + +`np.flatnonzero()` returns the original indices of retained points. The retained points are +then grouped by local HEALPix cell. The output is a list of length +`cell_end - cell_start`; empty cells remain present as empty index arrays. + +This happens for every input stream, including: + +- `METOP_ABC_AVHRR_IASI`; +- `ERA5_in`; +- `ERA5`; +- `METEOSAT_SEVIRI_IR`; +- `GOES_ABI_IR`; +- `HIMAWARI_AHI_IR`; +- `GOES_ABI_VIS`; +- `HIMAWARI_AHI_VIS`; +- `SurfaceCombined`. + +Each stream can have a different number of retained observations and tokens on a rank. + +### Local tokenization + +`TokenizerMasking` receives `source_cell_start` and `source_cell_end`. Source windows call +`tokenize_space()` or `tokenize_spacetime()` with this interval. Target windows continue to +use the global interval. + +For each local source cell: + +1. locations are ordered by colatitude (`theta`); +2. ordered locations are split into patches of the stream's configured `token_size`; +3. the final patch is padded when required; +4. masking selects retained patches without changing their relative cell order; +5. `source_tokens_lens` records the number of retained patches in every local cell. + +`StreamData` is initialized with: + +```text +source HEALPix cells = spatial_parallel.local_num_cells +target HEALPix cells = num_healpix_cells +``` + +Consequently, for level 5 and four-way spatial parallelism: + +```python +batch.source_samples.tokens_lens.shape[-1] == 3072 +batch.target_samples.tokens_lens.shape[-1] == 12288 +``` + +### Variable-length METOP observations + +METOP is a useful example because its number of observations varies by time and domain. + +Before tokenization, METOP-A, METOP-B, and METOP-C reader rows are concatenated in configured +reader order. Within each reader, source rows retain their storage order unless +`shuffle_source` is enabled. + +For multiple source windows, `StreamData` stores step 0 as the newest window, followed by older +windows. This step order is retained when the embedding engine concatenates the stream inputs. + +Tokenization deliberately converts this raw row order into cell-major order: + +```text +local HEALPix cell 0 + locations ordered north-to-south + patch 0: up to 512 locations + patch 1: up to 512 locations + ... +local HEALPix cell 1 + ... +``` + +A tensor such as: + +```text +source_tokens_cells[step].shape = (N_rank, 512, 30) +``` + +has: + +- a variable `N_rank`, the number of retained METOP patch tokens on this rank; +- 512 padded locations per patch; +- 30 encoded features per location. + +`N_rank` does not need to be equal across spatial ranks. The variable-length METOP tensor is +never directly passed to a fixed-shape all-gather. + +The temporary raw-row indices used by tokenization are not retained after source construction. +The implementation preserves equivalence with the original cell-major model order, but does +not provide an inverse map from assimilated latents back to original METOP-A/B/C storage rows. + +## Embedding order + +The embedding engine processes streams in configuration order. For each stream, it concatenates +source tokens over input steps and batch samples and applies the stream-specific embedding +network. + +At this point, embeddings are stream-major. `get_scatter_idxs_vectorized()` uses +`batch.tokens_lens` to scatter them into packed cell-major order. Within each +`(input_step, sample, cell)` position, streams retain configuration order and each stream +retains its per-cell patch order. + +The positional-encoding index is derived from the same per-cell counts. Therefore the packed +embedding tensor, its cell lengths, and its positional encodings remain aligned. + +If a rank owns no observations from any stream for a sample, the embedding engine returns an +empty tensor rather than failing. The rank must still participate in later synchronized model +and distributed operations. + +## Local assimilation and empty domains + +`cell_lens_local` describes the packed local token tensor in: + +```text +(input_step, sample, local_cell) +``` + +order. Local assimilation uses these lengths to derive cumulative packed-token boundaries. +Cell boundaries are never inferred by evenly slicing the token dimension, because token counts +vary by stream, time, cell, and rank. + +When spatial parallelism is enabled, one local rank shard is processed as one local-assimilation +chunk. This keeps the number and order of calls to FSDP-wrapped local modules consistent across +spatial ranks. + +A local domain can contain no observations. Skipping FSDP modules on that rank would cause +ranks to enter FSDP collectives in different orders and can deadlock backward. For an empty +chunk, the implementation therefore: + +1. creates a one-token zero-valued dummy input; +2. calls the local assimilation engine; +3. calls latent interpolation; +4. calls the local-to-global adapter; +5. multiplies the result by zero and attaches it to the real output. + +The zero dependency has no numerical effect, but preserves the forward/backward graph and +ensures that FSDP hooks execute consistently on every spatial rank. + +## Local-to-global boundary and all-gather + +The gather occurs after: + +```text +stream embedding + → local assimilation + → latent interpolation + → local-to-global projection +``` + +It occurs before: + +```text +query aggregation + → global assimilation +``` + +This is the principal synchronization boundary of the feature. + +### Why variable-length source tensors can be gathered safely + +The variable-length source tensors are first projected into a fixed number of query latents per +cell. Before gathering, every rank restores a dense local tensor: + +```text +[ + input_steps × samples, + spatial_parallel.local_num_cells, + local_queries_per_cell, + global_embedding_dimension, +] +``` + +Non-empty cells receive the projected local result. Empty cells retain their initialized latent +slot. All ranks in a spatial group therefore have the same gather shape even when their raw +observation and patch-token counts differ substantially. + +The implementation uses `torch.distributed.nn.functional.all_gather`, not the non-autograd +collective, so gradients propagate from global processing back into the owning rank's local +encoder path. + +### Reconstruction order + +The spatial process group is created from increasing, consecutive global ranks. The gather +returns tensors in process-group rank order: + +```text +[spatial rank 0, spatial rank 1, ..., spatial rank S - 1] +``` + +Each rank's local tensor is already ordered by increasing cell ID within its consecutive range. +Concatenating gathered tensors along the cell dimension therefore reconstructs: + +```text +global cell 0, global cell 1, ..., global cell N - 1 +``` + +For four ranks at level 5: + +```text +gather result = + rank 0 cells [0, 3072) + + rank 1 cells [3072, 6144) + + rank 2 cells [6144, 9216) + + rank 3 cells [9216, 12288) +``` + +The rank-local `tokens_lens` tensors are gathered and concatenated along their cell dimension in +the same rank order. The reconstructed global cell mask is therefore aligned with the dense +global latent tensor. + +The implementation then packs non-empty global cells for query aggregation and restores the +global dense layout afterward. + +## Ordering guarantees + +The feature guarantees model-equivalent cell ordering: + +1. nested HEALPix maps every source location to a deterministic global cell ID; +2. each global cell belongs to exactly one spatial rank; +3. rank-local cell lists are ordered by increasing global cell ID; +4. local tokenization matches the corresponding slice of global tokenization; +5. stream embeddings are deterministically scattered using `tokens_lens`; +6. dense local latent position `j` maps to global cell `cell_start + j`; +7. all-gather results are concatenated in spatial-rank order; +8. gathered lengths and gathered latents use the same order. + +This does not mean that the original source storage-row order survives tokenization. The model's +canonical order is cell-major, followed by per-cell location/patch order. That was already the +order consumed by local assimilation before spatial parallelism. + +The local/global construction equivalence test builds the global cell lists and the four local +cell lists independently, concatenates the local lists, and checks exact array equality for every +cell. + +### Sort stability caveat + +On NumPy 2.x, cell grouping explicitly requests a stable sort. On NumPy 1.x, the compatibility +fallback currently uses NumPy's default `argsort`. The subsequent PyTorch colatitude sort is +stable, so ordinary observations remain deterministic, but exact relative ordering of duplicate +cell IDs with identical colatitudes is not formally guaranteed on the NumPy 1.x fallback. + +If exact duplicate-location ordering is required across NumPy versions, use: + +```python +np.argsort(local_cell_ids, kind="stable") +``` + +and retain an explicit raw observation identifier for inverse mapping. + +## Backward-compatible full-grid input + +The encoder accepts either: + +- a rank-local source batch with `spatial_parallel.local_num_cells`; or +- a legacy global source batch with `num_healpix_cells`. + +For a local batch, the embedding result is already local. + +For a global batch, `select_packed_cell_shard()` selects complete cells from the packed token +tensor using `cell_lens`. This fallback preserves compatibility, but the selection occurs after +embedding, so it does not provide the embedding-memory reduction of rank-local source +construction. + +`select_packed_cell_shard()`: + +1. reshapes flattened cell lengths into complete HEALPix grids; +2. marks `[cell_start, cell_end)` in every input-step/sample row; +3. expands the cell mask by each cell's variable token count; +4. selects the corresponding packed tokens; +5. returns local cell lengths in unchanged packed order. + +It validates the cell range and verifies that the lengths describe the actual packed tensor. + +## Data-parallel training semantics + +Spatial ranks consume the same sample and must not be counted as independent data-parallel +replicas. + +The effective data-parallel world size is: + +```text +distributed.data_parallel.world_size = world_size / spatial_parallel_size +``` + +The effective batch size is: + +```text +effective_batch_size = + batch_size_per_spatial_group × distributed.data_parallel.world_size +``` + +The trainer uses this effective data-parallel size for: + +- total batch-size reporting; +- learning-rate scaling; +- scheduler construction; +- mini-epoch and continuation accounting. + +The sampler similarly maps global ranks in the same spatial group to one data-parallel rank, so +they receive the same workset and random seed. Loader-worker IDs and mini-epoch indices still +differentiate independent workers and epochs. + +When present in a saved run, `distributed.spatial_parallel.size_original` records the historical +spatial size so effective-batch calculations remain consistent. If the key is absent, the +implementation defaults it to the current spatial size; continuation from a pre-feature run +therefore requires checking this value explicitly. + +## Configuration + +The relevant option is: + +```yaml +distributed: + data_parallel: + with_ddp: false + with_fsdp: true + find_unused_parameters: true + spatial_parallel: + size: 4 +``` + +`distributed.spatial_parallel.size` controls the number of ranks per spatial group. The +`distributed.data_parallel` subsection can be passed directly to model initialization. + +Ready-to-use overrides are provided: + +```text +config/spatial_parallel_4.yml +config/spatial_parallel_8.yml +``` + +For a four-rank job in which all ranks cooperate on one encoder sample: + +```yaml +distributed: + spatial_parallel: + size: 4 +``` + +The expected derived values are: + +```text +world_size: 4 +distributed.data_parallel.world_size: 1 +local level-5 cells/rank: 3072 +``` + +For an eight-rank job with two four-rank spatial groups: + +```text +world_size: 8 +distributed.spatial_parallel.size: 4 +distributed.data_parallel.world_size: 2 +``` + +### Continuing an existing run + +A continued run can inherit a saved configuration that predates this feature. Pass the spatial +configuration explicitly when the intended current run should use spatial parallelism. + +Verify the effective, merged configuration rather than relying only on the repository default. +In a four-rank spatial-only run, `distributed.data_parallel.world_size` must be 1, not 4. + +## Runtime verification + +### Startup log + +Every spatial rank logs its fine-cell range. For level 5 and four ranks, expect: + +```text +Encoder spatial rank 0/4 constructs source HEALPix cells [0, 3072) +Encoder spatial rank 1/4 constructs source HEALPix cells [3072, 6144) +Encoder spatial rank 2/4 constructs source HEALPix cells [6144, 9216) +Encoder spatial rank 3/4 constructs source HEALPix cells [9216, 12288) +``` + +### Tensor-shape checks + +For level 5 and four spatial ranks: + +```python +assert batch.source_samples.tokens_lens.shape[-1] == 3072 +assert batch.target_samples.tokens_lens.shape[-1] == 12288 +``` + +For a globally regular stream such as ERA5, each rank should retain approximately one quarter of +the source tokens. Exact equality depends on grid conventions and masks. + +For regional or observational streams such as METOP and geostationary satellite streams, token +counts can be strongly imbalanced. The important invariants are: + +- every retained source token belongs to the rank's cell interval; +- no source token belongs to two ranks; +- the union of rank-local cell intervals covers the full grid. + +### Memory checks + +Use PyTorch allocated-memory statistics when measuring the feature: + +```python +torch.cuda.reset_peak_memory_stats() + +# Run one representative training iteration. + +peak_allocated = torch.cuda.max_memory_allocated() / 2**30 +peak_reserved = torch.cuda.max_memory_reserved() / 2**30 +``` + +`nvidia-smi` primarily reflects CUDA memory reserved by PyTorch. Released tensor memory remains +in the caching allocator and may not visibly decrease even when live tensor memory does. + +Compare: + +- peak allocated memory; +- peak reserved memory; +- per-stage peaks; +- all spatial ranks, because token imbalance can make one domain more expensive than another. + +## Tests + +`tests/test_spatial_parallel.py` covers: + +| Test area | Invariant | +| --- | --- | +| Local construction equivalence | Concatenated local cell lists equal global construction cell-by-cell | +| Invalid local range | Out-of-range cell intervals are rejected | +| Packed-token selection | Complete cells are selected across multiple input-step/sample rows | +| Coverage and gradients | Shards cover every packed token exactly once and preserve gradients | +| Invalid packed ranges | Invalid cell intervals are rejected | +| Distributed size validation | Spatial groups must divide the distributed world | + +Recommended validation commands are: + +```bash +pytest -q tests/test_spatial_parallel.py +./scripts/actions.sh lint +./scripts/actions.sh unit-test +``` + +A multi-rank integration run remains necessary to validate FSDP collective ordering, runtime +memory, and real-stream token balance. + +## Expected memory behavior + +Memory savings are not expected to be exactly `1 / spatial_parallel_size`. + +Memory that should decrease includes: + +- rank-local source token tensors transferred to the GPU; +- stream-embedding activations; +- local-assimilation activations; +- local-to-global projection activations. + +Memory that remains replicated or becomes global includes: + +- model parameters, gradients, and optimizer state according to the FSDP configuration; +- source-reader CPU data before local filtering; +- target tensors and target-side computation; +- the dense gathered global latent tensor; +- query aggregation; +- global assimilation; +- forecast engine and prediction heads; +- CUDA allocator cache. + +The rank with the most observations can determine the job's usable batch size. Equal numbers of +HEALPix cells do not imply equal numbers of stream tokens. + +## Troubleshooting + +### Embedding memory does not decrease + +Check: + +```python +batch.source_samples.tokens_lens.shape[-1] +``` + +If it is 12,288 at level 5 in a four-way run, the data pipeline constructed a global source +batch. The encoder can still select a local shard after embedding, but embedding remains global. + +Verify: + +- `distributed.spatial_parallel.size: 4` is present in the effective run configuration; +- `distributed.data_parallel.world_size: 1` for a four-rank spatial-only run; +- all four ownership messages appear in the startup log; +- each source stream has a rank-local token count. + +### `nvidia-smi` remains high + +Compare `torch.cuda.max_memory_allocated()` with `torch.cuda.max_memory_reserved()`. A large +difference usually represents reusable cached allocator segments, not live model tensors. + +### A rank has no observations + +This is valid for regional streams and sparse samples. The embedding engine returns an empty +packed tensor, and the local assimilation path executes a zero-valued dummy dependency to keep +FSDP calls synchronized. + +### Gather shape mismatch + +The gather must receive: + +```text +[steps × samples, local_cells, queries_per_cell, embedding_dimension] +``` + +with identical `local_cells`, query count, dtype, and embedding dimension on all ranks in the +spatial group. Variable source-token counts must be resolved before this boundary. + +### Global ordering mismatch + +Verify together: + +- the identical contiguous local range computed by the sampler and encoder; +- nested HEALPix mapping (`nest=True`); +- local `tokens_lens` cell dimension; +- rank order inside the spatial process group; +- concatenation along the cell dimension; +- identical ordering for gathered lengths and gathered latents. + +Do not compare assimilated output against raw source-reader row order. Compare it against the +non-parallel model's cell-major token and latent order. + +### FSDP hang during backward + +Confirm that: + +- all ranks in a spatial group consume the same sample; +- all ranks enter the local FSDP modules in the same order; +- empty ranks execute the dummy dependency path; +- `world_size` is divisible by the spatial size; +- no rank independently skips an otherwise valid sample because its local source domain is empty. + +## File-level implementation map + +| File | Responsibility | +| --- | --- | +| `config/default_config.yml` | Default spatial size | +| `config/spatial_parallel_4.yml` | Four-rank override | +| `config/spatial_parallel_8.yml` | Eight-rank override | +| `src/weathergen/datasets/healpix_domain.py` | Rank-local point grouping | +| `src/weathergen/datasets/multi_stream_data_sampler.py` | Shared spatial-group samples, local ownership, and runtime logging | +| `src/weathergen/datasets/stream_data.py` | Separate source-local and target-global cell counts | +| `src/weathergen/datasets/tokenizer_masking.py` | Local source tokenization range | +| `src/weathergen/datasets/tokenizer_utils.py` | Nested HEALPix mapping and local token construction | +| `src/weathergen/model/engines.py` | Local per-stream embedding and empty-domain handling | +| `src/weathergen/model/spatial_parallel.py` | Legacy packed-token shard selection | +| `src/weathergen/model/encoder.py` | Local assimilation, local-to-global projection, and differentiable gather | +| `src/weathergen/train/trainer.py` | Effective data-parallel batch and scheduler semantics | +| `src/weathergen/utils/distributed.py` | Shared spatial topology context, validation, and process-group construction | +| `tests/test_spatial_parallel.py` | Cell ownership, ordering, coverage, validation, and gradient tests | + +## Commit-by-commit design history + +### `15dcff5a`: Add HEALPix encoder spatial parallelism + +This commit establishes the initial model-side and distributed design: + +- adds `distributed.spatial_parallel.size`; +- creates four-rank and eight-rank configuration overrides; +- creates consecutive-rank spatial process groups; +- makes spatial-group ranks consume the same data sample; +- changes effective batch and scheduler scaling from global world size to data-parallel group count; +- introduces `select_packed_cell_shard()` for variable-length cell-packed tensors; +- selects local cells before local assimilation; +- builds local latent query seeds and local positional encodings; +- performs local assimilation and local-to-global projection on each rank; +- restores dense local cells and gathers them before global processing; +- uses an autograd-aware gather; +- adds synchronized handling for empty local domains; +- adds initial unit tests for packed-cell coverage, gradients, and distributed-size validation. + +At this stage, embedding still runs on the global packed source tensor and the result is selected +after embedding. + +### `db4b477d`: Build encoder inputs on local HEALPix domains + +This commit moves the domain boundary into the data pipeline: + +- introduces rank-local HEALPix point grouping; +- passes source cell ranges into source tokenization; +- constructs only rank-local source cells and source token tensors; +- keeps targets global; +- adds separate source and target HEALPix dimensions to `StreamData`; +- makes embedding operate on local source tensors; +- gathers local `tokens_lens` to reconstruct the global mask; +- accepts both local and legacy global source batches; +- handles ranks with no embedded streams; +- validates local/global batch cell dimensions; +- adds exact local-construction-versus-global-slice tests. + +This is the commit that enables embedding activation memory to scale with the local observation +domain. + +## Review checklist + +Before merging or extending this feature, verify: + +- [ ] The effective spatial size is explicit in the run configuration. +- [ ] `world_size % distributed.spatial_parallel.size == 0`. +- [ ] The data-level HEALPix cell count is divisible by the spatial size. +- [ ] Every spatial group consumes identical samples. +- [ ] Every source stream constructs only local cells. +- [ ] Source and target cell dimensions remain intentionally different. +- [ ] Variable-length stream tokens are resolved into fixed-size cell latents before gather. +- [ ] Local lengths and local latents are gathered in identical rank order. +- [ ] The gather remains autograd-aware. +- [ ] Empty local domains enter all FSDP modules in consistent order. +- [ ] Effective batch size counts spatial groups, not individual spatial ranks. +- [ ] Tests cover exact ordering, token coverage, and gradients. +- [ ] Runtime validation compares allocated memory across all spatial ranks. diff --git a/src/weathergen/datasets/healpix_domain.py b/src/weathergen/datasets/healpix_domain.py new file mode 100644 index 0000000000..7d0a069e42 --- /dev/null +++ b/src/weathergen/datasets/healpix_domain.py @@ -0,0 +1,50 @@ +# (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. + +import numpy as np + + +def build_local_healpix_cell_splits( + cell_ids: np.typing.NDArray[np.integer], + num_cells: int, + cell_start: int, + cell_end: int, +) -> list[np.typing.NDArray[np.int64]]: + """Group original point indices for one consecutive HEALPix-cell domain.""" + + if not 0 <= cell_start < cell_end <= num_cells: + raise ValueError( + f"invalid HEALPix cell range [{cell_start}, {cell_end}) for {num_cells} cells" + ) + + # Domain-parallel filtering mask: this is applied independently to every + # stream immediately after its coordinates have been mapped to nested + # HEALPix cell IDs. + local_domain_mask = (cell_ids >= cell_start) & (cell_ids < cell_end) + local_point_idxs = np.flatnonzero(local_domain_mask) + local_cell_ids = cell_ids[local_point_idxs] + cell_splits = [np.array([], dtype=np.int64) for _ in range(cell_end - cell_start)] + if local_point_idxs.size == 0: + return cell_splits + + stable_args = {"stable": True} if int(np.__version__.split(".")[0]) >= 2 else {} + local_order = np.argsort(local_cell_ids, **stable_args) + sorted_point_idxs = local_point_idxs[local_order] + sorted_cell_ids = local_cell_ids[local_order] + split_offsets = np.flatnonzero(np.diff(sorted_cell_ids)) + point_idxs_by_occupied_cell = np.split(sorted_point_idxs, split_offsets + 1) + + for cell_id, point_idxs in zip( + np.unique(sorted_cell_ids), + point_idxs_by_occupied_cell, + strict=True, + ): + cell_splits[cell_id - cell_start] = point_idxs + + return cell_splits diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py index bd4a8ad87e..dd9d6e2c5e 100644 --- a/src/weathergen/datasets/multi_stream_data_sampler.py +++ b/src/weathergen/datasets/multi_stream_data_sampler.py @@ -34,7 +34,7 @@ ) from weathergen.readers_extra.registry import get_extra_reader from weathergen.train.utils import Stage, get_batch_size_from_config -from weathergen.utils.distributed import is_root +from weathergen.utils.distributed import SpatialParallelContext, is_root type AnyDataReader = DataReaderBase | DataReaderAnemoi | DataReaderObs type StreamName = str @@ -100,15 +100,32 @@ def __init__(self, cf: Config, mode_cfg: dict, stage: Stage): self.mini_epoch = 0 self.mask_value = 0.0 - self.rank = cf.rank - self.world_size = cf.world_size self.repeat_data = cf.data_loading.get("repeat_data_in_mini_epoch", False) # initialise healpic self.healpix_level = cf.healpix_level self.num_healpix_cells = 12 * 4**self.healpix_level + # Ranks in one spatial group must consume the same batch. Data + # parallelism therefore operates across groups, not individual ranks. + self.spatial_parallel = SpatialParallelContext.from_config( + cf, + self.num_healpix_cells, + ) self.masker = Masker(cf.healpix_level, stage, cf.streams, self.mode_cfg) - self.tokenizer = TokenizerMasking(cf.healpix_level, self.masker) + self.tokenizer = TokenizerMasking( + cf.healpix_level, + self.masker, + self.spatial_parallel.cell_start, + self.spatial_parallel.cell_end, + ) + if self.spatial_parallel.size > 1: + logger.info( + "Spatial rank %d/%d constructs source HEALPix cells [%d, %d)", + self.spatial_parallel.rank, + self.spatial_parallel.size, + self.spatial_parallel.cell_start, + self.spatial_parallel.cell_end, + ) forecast_cfg = FORECAST_DEFAULTS | OmegaConf.to_object(mode_cfg.get("forecast", {})) self.output_offset = forecast_cfg["offset"] @@ -190,11 +207,12 @@ def check_samples(self, fsm: int): epoch_len = self.samples_per_mini_epoch # ensure epoch_len is large enough to produce at least one batch per rank - min_samples = self.world_size * self.batch_size + min_samples = self.spatial_parallel.ddp_world_size * self.batch_size if epoch_len < min_samples: logger.warning( f"samples_per_mini_epoch={epoch_len} is too small for " - f"world_size={self.world_size} and batch_size={self.batch_size}. " + f"ddp_world_size={self.spatial_parallel.ddp_world_size} " + f"and batch_size={self.batch_size}. " f"samples_per_mini_epoch has to be equal to or larger than" f"world_size*batch_size to ensure that each rank can produce at least one sample. " f"Automatically increasing to {min_samples}." @@ -203,9 +221,11 @@ def check_samples(self, fsm: int): self.samples_per_mini_epoch = min_samples # adjust len to split loading across all workers and ensure it is multiple of batch_size - self.len = ((epoch_len // self.world_size) // self.batch_size) * self.batch_size + self.len = ( + (epoch_len // self.spatial_parallel.ddp_world_size) // self.batch_size + ) * self.batch_size - n_duplicates = self.len * self.world_size - available_samples + n_duplicates = self.len * self.spatial_parallel.ddp_world_size - available_samples if not self.repeat_data: assert n_duplicates <= 0 @@ -532,6 +552,7 @@ def _build_stream_data( num_steps_input, num_output_steps, self.num_healpix_cells, + source_num_healpix_cells=self.spatial_parallel.local_num_cells, ) stream_data = self._build_stream_data_input( @@ -703,7 +724,9 @@ def _get_batch(self, idx: int, num_forecast_steps: int): # tokenize windows # *_tokens = [ (cells_idx, cells_idx_lens), ... ] with length = #time_steps - input_tokens = self.tokenizer.get_tokens_windows(stream_info, input_data, True) + input_tokens = self.tokenizer.get_tokens_windows( + stream_info, input_data, True, local_source=True + ) output_tokens = self.tokenizer.get_tokens_windows(stream_info, output_data, False) for sidx, source_mask in enumerate(source_masks.masks): @@ -791,7 +814,13 @@ def __iter__(self) -> ModelBatch: # ensure the batch is valid, i.e. not completely empty and no NaN values # student teacher has no classical targets mode = self.mode_cfg.get("training_mode") - not_valid = batch.sources_empty() or batch.is_nan() + # A valid global sample may have no observations in one rank's + # local source domain. All spatial ranks must still emit the + # same sample and enter encoder collectives in lockstep. + local_sources_empty = batch.sources_empty() + not_valid = ( + local_sources_empty if self.spatial_parallel.size == 1 else False + ) or batch.is_nan() not_valid = not_valid or (batch.targets_empty() if "masking" in mode else False) # skip completely empty batch item or when all targets are empty -> no grad @@ -806,12 +835,13 @@ def __len__(self): return self.len def worker_workset(self): - local_start, local_end = self.rank * self.len, (self.rank + 1) * self.len + local_start = self.spatial_parallel.ddp_rank * self.len + local_end = (self.spatial_parallel.ddp_rank + 1) * self.len worker_info = torch.utils.data.get_worker_info() if worker_info is None: - assert self.world_size == 1, self.world_size + assert self.spatial_parallel.ddp_world_size == 1 iter_start = 0 iter_end = len(self) @@ -822,9 +852,8 @@ def worker_workset(self): # happens for each mini_epoch, for train and validation, and independently for each DDP # worker. After the bit-wise copy, the rng seed needs to be made unique for # DDP workers, loader process, mini_epoch. - dist = torch.distributed self.data_loader_rng_seed *= ( - (((dist.get_rank() + 1) * 73) if dist.is_initialized() else 1) + ((self.spatial_parallel.ddp_rank + 1) * 73) * ((worker_info.id + 1) * 37) * (self.mini_epoch + 13) * 7 @@ -836,7 +865,7 @@ def worker_workset(self): if worker_info.id + 1 == worker_info.num_workers: iter_end = local_end logger.info( - f"{self.rank}::{worker_info.id}" + f"{self.spatial_parallel.ddp_rank}::{worker_info.id}" + f" : dataset [{local_start},{local_end}) : [{iter_start},{iter_end})" ) diff --git a/src/weathergen/datasets/stream_data.py b/src/weathergen/datasets/stream_data.py index 4c3bb02b8b..6d8a2849ab 100644 --- a/src/weathergen/datasets/stream_data.py +++ b/src/weathergen/datasets/stream_data.py @@ -58,6 +58,7 @@ def __init__( input_steps: int, output_steps: int, healpix_cells: int, + source_num_healpix_cells: int | None = None, ) -> None: """ StreamData object @@ -72,7 +73,10 @@ def __init__( Number of output steps Note -- Last input step and first output step always overlap. healpix_cells : int - Number of healpix cells for source + Number of global healpix cells used for targets + source_num_healpix_cells : int | None + Number of rank-local healpix cells used for encoder sources. Defaults + to ``healpix_cells`` for non-spatial-parallel callers. Returns ------- @@ -84,6 +88,9 @@ def __init__( self.input_steps = input_steps self.output_steps = output_steps self.healpix_cells = healpix_cells + self.source_num_healpix_cells = ( + healpix_cells if source_num_healpix_cells is None else source_num_healpix_cells + ) self.source_is_spoof = [False for _ in range(self.input_steps)] self.target_is_spoof = [False for _ in range(self.output_steps)] @@ -104,7 +111,8 @@ def __init__( self.source_tokens_cells = [None for _ in range(self.input_steps)] # length of source tokens per cell (without padding) self.source_tokens_lens = [ - torch.zeros(self.healpix_cells, dtype=torch.int32) for _ in range(self.input_steps) + torch.zeros(self.source_num_healpix_cells, dtype=torch.int32) + for _ in range(self.input_steps) ] # unprocessed source (for logging) self.source_raw = [None for _ in range(self.input_steps)] diff --git a/src/weathergen/datasets/tokenizer_masking.py b/src/weathergen/datasets/tokenizer_masking.py index 7c033e398f..0491ef30d7 100644 --- a/src/weathergen/datasets/tokenizer_masking.py +++ b/src/weathergen/datasets/tokenizer_masking.py @@ -40,11 +40,28 @@ def readerdata_to_torch(rdata: IOReaderData) -> IOReaderData: class TokenizerMasking(Tokenizer): - def __init__(self, healpix_level: int, masker: Masker): + def __init__( + self, + healpix_level: int, + masker: Masker, + source_cell_start: int = 0, + source_cell_end: int | None = None, + ): super().__init__(healpix_level) self.masker = masker self.rng = None self.token_size = None + self.source_cell_start = source_cell_start + self.source_cell_end = ( + self.num_healpix_cells_source if source_cell_end is None else source_cell_end + ) + if not ( + 0 <= self.source_cell_start < self.source_cell_end <= self.num_healpix_cells_source + ): + raise ValueError( + f"invalid source HEALPix cell range " + f"[{self.source_cell_start}, {self.source_cell_end})" + ) def reset_rng(self, rng) -> None: """ @@ -53,7 +70,7 @@ def reset_rng(self, rng) -> None: self.masker.reset_rng(rng) self.rng = rng - def get_tokens_windows(self, stream_info, data, pad_tokens): + def get_tokens_windows(self, stream_info, data, pad_tokens, local_source=False): """ Tokenize data (to amortize over the different views that are generated) @@ -63,6 +80,8 @@ def get_tokens_windows(self, stream_info, data, pad_tokens): tok = tokenize_spacetime if tok_spacetime else tokenize_space hl = self.healpix_level token_size = stream_info["token_size"] + cell_start = self.source_cell_start if local_source else 0 + cell_end = self.source_cell_end if local_source else self.num_healpix_cells_source tokens = [] for rdata in data: @@ -72,7 +91,12 @@ def get_tokens_windows(self, stream_info, data, pad_tokens): continue # tokenize data idxs_cells, idxs_cells_lens = tok( - readerdata_to_torch(rdata), token_size, hl, pad_tokens + readerdata_to_torch(rdata), + token_size, + hl, + pad_tokens, + cell_start=cell_start, + cell_end=cell_end, ) tokens += [(idxs_cells, idxs_cells_lens)] @@ -129,6 +153,7 @@ def get_source( ): # create tokenization index (idxs_cells, idxs_cells_lens) = idxs_cells_data + cell_mask = cell_mask[self.source_cell_start : self.source_cell_end] # select strategy from XXX depending on stream and if student or teacher @@ -144,7 +169,7 @@ def get_source( stream_info["stream_id"], rdata, time_win, - self.hpy_verts_rots_source[-1], + self.hpy_verts_rots_source[-1][self.source_cell_start : self.source_cell_end], encode_times_source, ) diff --git a/src/weathergen/datasets/tokenizer_utils.py b/src/weathergen/datasets/tokenizer_utils.py index 1bd1d1722d..571d2d9cbd 100644 --- a/src/weathergen/datasets/tokenizer_utils.py +++ b/src/weathergen/datasets/tokenizer_utils.py @@ -5,6 +5,7 @@ from torch import Tensor from weathergen.common.io import IOReaderData +from weathergen.datasets.healpix_domain import build_local_healpix_cell_splits from weathergen.datasets.utils import ( locs_to_cell_coords_ctrs, locs_to_ctr_coords, @@ -12,10 +13,6 @@ s2tor3, ) -# on some clusters our numpy version is pinned to be 1.x.x where the np.argsort does not -# the stable=True argument -numpy_argsort_args = {"stable": True} if int(np.__version__.split(".")[0]) >= 2 else {} - def theta_phi_to_standard_coords(coords): thetas = ((90.0 - coords[:, 0]) / 180.0) * np.pi @@ -95,34 +92,46 @@ def encode_times_target(times, time_win) -> torch.tensor: return time_tensor + 0.5 -def hpy_cell_splits(coords: torch.tensor, hl: int): +def hpy_cell_splits( + coords: torch.tensor, + hl: int, + cell_start: int = 0, + cell_end: int | None = None, +): """Compute healpix cell id for each coordinate on given level hl Returns - hpy_idxs_ord_split : list of per cell indices into thetas,phis,posr3 + hpy_idxs_ord_split : list of per-local-cell indices into thetas,phis,posr3 thetas : thetas in rad phis : phis in rad """ + num_healpix_cells = 12 * 4**hl + cell_end = num_healpix_cells if cell_end is None else cell_end + thetas, phis = theta_phi_to_standard_coords(coords) # healpix cells for all points hpy_idxs = ang2pix(2**hl, thetas, phis, nest=True) - # extract information to split according to cells by first sorting and then finding split idxs - hpy_idxs_ord = np.argsort(hpy_idxs, **numpy_argsort_args) - splits = np.flatnonzero(np.diff(hpy_idxs[hpy_idxs_ord])) - - # extract per cell data - hpy_idxs_ord_temp = np.split(hpy_idxs_ord, splits + 1) - hpy_idxs_ord_split = [np.array([], dtype=np.int64) for _ in range(12 * 4**hl)] - # TODO: split smarter (with a augmented splits list?) so that this loop is not needed - for b, x in zip(np.unique(np.unique(hpy_idxs[hpy_idxs_ord])), hpy_idxs_ord_temp, strict=True): - hpy_idxs_ord_split[b] = x + # Nested HEALPix IDs make a consecutive interval a complete rank-local + # domain. The helper applies the point mask and builds only local cells. + hpy_idxs_ord_split = build_local_healpix_cell_splits( + hpy_idxs, + num_healpix_cells, + cell_start, + cell_end, + ) return (hpy_idxs_ord_split, thetas, phis) def hpy_splits( - coords: torch.Tensor, hl: int, token_size: int, pad_tokens: bool, offset_step: int = 0 + coords: torch.Tensor, + hl: int, + token_size: int, + pad_tokens: bool, + offset_step: int = 0, + cell_start: int = 0, + cell_end: int | None = None, ) -> tuple[list[torch.Tensor], list[torch.Tensor], torch.Tensor]: """Compute healpix cell for each data point and splitting information per cell; when the token_size is exceeded then splitting based on lat is used; @@ -135,7 +144,9 @@ def hpy_splits( """ # list of data points per healpix cell - (hpy_idxs_ord_split, thetas, phis) = hpy_cell_splits(coords, hl) + (hpy_idxs_ord_split, thetas, phis) = hpy_cell_splits( + coords, hl, cell_start=cell_start, cell_end=cell_end + ) # if token_size is exceeed split based on latitude # TODO: split by hierarchically traversing healpix scheme @@ -179,11 +190,21 @@ def tokenize_space( hl, pad_tokens=True, offset_step=0, + cell_start=0, + cell_end=None, ): """Process one window into tokens""" # idx_ord_lens is length is number of tokens per healpix cell - idxs_ord, idxs_ord_lens = hpy_splits(rdata.coords, hl, token_size, pad_tokens, offset_step) + idxs_ord, idxs_ord_lens = hpy_splits( + rdata.coords, + hl, + token_size, + pad_tokens, + offset_step, + cell_start, + cell_end, + ) return idxs_ord, idxs_ord_lens @@ -193,14 +214,17 @@ def tokenize_spacetime( token_size, hl, pad_tokens=True, + cell_start=0, + cell_end=None, ): """Tokenize respecting an intrinsic time step in the data, i.e. each time step is tokenized separately """ num_healpix_cells = 12 * 4**hl - idxs_cells = [[] for _ in range(num_healpix_cells)] - idxs_cells_lens = [[] for _ in range(num_healpix_cells)] + cell_end = num_healpix_cells if cell_end is None else cell_end + idxs_cells = [[] for _ in range(cell_end - cell_start)] + idxs_cells_lens = [[] for _ in range(cell_end - cell_start)] offset_step = 0 t_unique = np.unique(rdata.datetimes) @@ -210,7 +234,15 @@ def tokenize_spacetime( rdata_cur = IOReaderData( rdata.coords[mask], rdata.geoinfos[mask], rdata.data[mask], rdata.datetimes[mask] ) - idxs_cur, idxs_cur_lens = tokenize_space(rdata_cur, token_size, hl, pad_tokens, offset_step) + idxs_cur, idxs_cur_lens = tokenize_space( + rdata_cur, + token_size, + hl, + pad_tokens, + offset_step, + cell_start, + cell_end, + ) # collect data for all time steps idxs_cells = [t + tc for t, tc in zip(idxs_cells, idxs_cur, strict=True)] diff --git a/src/weathergen/model/encoder.py b/src/weathergen/model/encoder.py index 3e9f99a706..e1e319f85f 100644 --- a/src/weathergen/model/encoder.py +++ b/src/weathergen/model/encoder.py @@ -9,6 +9,7 @@ import torch from astropy_healpix import healpy +from torch.distributed.nn.functional import all_gather from torch.utils.checkpoint import checkpoint from weathergen.common.config import Config @@ -25,6 +26,8 @@ # from weathergen.model.model import ModelParams from weathergen.model.parametrised_prob_dist import LatentInterpolator from weathergen.model.positional_encoding import positional_encoding_harmonic +from weathergen.model.spatial_parallel import select_packed_cell_shard +from weathergen.utils.distributed import SpatialParallelContext class EncoderModule(torch.nn.Module): @@ -43,6 +46,11 @@ def __init__(self, cf: Config, sources_size, targets_num_channels, targets_coord self.healpix_level = cf.healpix_level self.num_healpix_cells = 12 * 4**self.healpix_level + self.spatial_parallel = SpatialParallelContext.from_config( + cf, + self.num_healpix_cells, + create_process_group=True, + ) self.cf = cf self.sources_size = sources_size @@ -117,6 +125,56 @@ def __init__(self, cf: Config, sources_size, targets_num_channels, targets_coord # global assimilation engine self.ae_global_engine = GlobalAssimilationEngine(cf, self.num_healpix_cells) + def _prepare_local_cell_tokens( + self, + stream_cell_tokens: torch.Tensor, + tokens_lens: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return packed tokens and per-cell lengths for this rank's HEALPix domain.""" + + cell_lens = torch.sum(tokens_lens, 2).flatten() + batch_num_cells = tokens_lens.shape[-1] + if batch_num_cells == self.spatial_parallel.local_num_cells: + # The data pipeline already constructed rank-local cells. + return stream_cell_tokens, cell_lens + + if batch_num_cells == self.num_healpix_cells: + # Backward-compatible path for batches containing the global grid. + return select_packed_cell_shard( + stream_cell_tokens, + cell_lens, + self.num_healpix_cells, + self.spatial_parallel.cell_start, + self.spatial_parallel.cell_end, + ) + + raise ValueError( + f"batch has {batch_num_cells} HEALPix cells; expected either " + f"{self.spatial_parallel.local_num_cells} local or " + f"{self.num_healpix_cells} global cells" + ) + + def _get_global_tokens_lens(self, tokens_lens: torch.Tensor) -> torch.Tensor: + """Return token lengths for the full HEALPix grid.""" + + batch_num_cells = tokens_lens.shape[-1] + if batch_num_cells == self.spatial_parallel.local_num_cells: + if self.spatial_parallel.size > 1: + return torch.cat( + all_gather(tokens_lens, group=self.spatial_parallel.group), + dim=-1, + ) + return tokens_lens + + if batch_num_cells == self.num_healpix_cells: + return tokens_lens + + raise ValueError( + f"batch has {batch_num_cells} HEALPix cells; expected either " + f"{self.spatial_parallel.local_num_cells} local or " + f"{self.num_healpix_cells} global cells" + ) + def forward(self, model_params, batch): """ Encoder forward @@ -125,9 +183,18 @@ def forward(self, model_params, batch): stream_cell_tokens = checkpoint( self.embed_engine, batch, model_params.pe_embed, use_reentrant=False ) + stream_cell_tokens, local_cell_lens = self._prepare_local_cell_tokens( + stream_cell_tokens, + batch.tokens_lens, + ) tokens_global, posteriors = checkpoint( - self.assimilate_local, model_params, stream_cell_tokens, batch, use_reentrant=False + self.assimilate_local, + model_params, + stream_cell_tokens, + batch, + local_cell_lens, + use_reentrant=False, ) tokens_global = checkpoint( @@ -153,7 +220,14 @@ def interpolate_latents(self, tokens: torch.Tensor) -> (torch.Tensor, torch.Tens return tokens, posteriors - def assimilate_local_project_chunked(self, tokens, tokens_global, cell_lens, q_cells_lens): + def assimilate_local_project_chunked( + self, + tokens, + tokens_global, + cell_lens, + q_cells_lens, + num_cells_per_sample=None, + ): """ Apply the local assimilation engine and then the local-to-global adapter using a chunking in the number of tokens @@ -163,14 +237,21 @@ def assimilate_local_project_chunked(self, tokens, tokens_global, cell_lens, q_c # combined cell lens for all tokens in batch across all input steps zero_pad = torch.zeros(1, device=tokens.device, dtype=torch.int32) - # subdivision factor for required splitting - clen = self.num_healpix_cells // (2 if self.cf.healpix_level <= 5 else 8) + # Spatial parallelism already reduces a rank to at most 1/8 of the + # HEALPix cells. Process that shard in one call so every spatial rank + # enters the local FSDP modules the same number of times. + num_cells_per_sample = num_cells_per_sample or self.num_healpix_cells + if self.spatial_parallel.size > 1: + clen = num_cells_per_sample + else: + clen = self.num_healpix_cells // (2 if self.cf.healpix_level <= 5 else 8) tokens_global_unmasked = [] posteriors = [] + empty_chunk_dependency = tokens_global.new_zeros(()) - for i in range(cell_lens.shape[0] // clen): - # make sure we properly catch all elements in last chunk - i_end = (i + 1) * clen if i < (cell_lens.shape[0] // clen) - 1 else cell_lens.shape[0] + num_chunks = (cell_lens.shape[0] + clen - 1) // clen + for i in range(num_chunks): + i_end = min((i + 1) * clen, cell_lens.shape[0]) l0, l1 = ( (0 if i == 0 else cell_lens[: i * clen].cumsum(0)[-1]), cell_lens[:i_end].cumsum(0)[-1], @@ -181,6 +262,20 @@ def assimilate_local_project_chunked(self, tokens, tokens_global, cell_lens, q_c # skip processing of the empty chunk in this case # Check if this chunk is empty if l0 == l1 or toks.shape[0] == 0: + # Spatial ranks must enter FSDP-wrapped local modules in the + # same order. Run a one-token zero-valued path and retain a + # zero dependency so its backward hooks also execute. + dummy_lens = torch.tensor([0, 1], device=tokens.device, dtype=torch.int32) + dummy_tokens = tokens.new_zeros((1, tokens.shape[-1])) + dummy_tokens = self.ae_local_engine(dummy_tokens, dummy_lens, use_reentrant=False) + dummy_tokens, _ = self.interpolate_latents(dummy_tokens) + dummy_global = self.ae_local_global_engine( + dummy_tokens, + tokens_global[i * clen : i * clen + 1], + dummy_lens, + dummy_lens, + ) + empty_chunk_dependency = empty_chunk_dependency + dummy_global.sum() * 0 continue toks_global = tokens_global[i * clen : i_end] @@ -210,10 +305,13 @@ def assimilate_local_project_chunked(self, tokens, tokens_global, cell_lens, q_c tokens_global_unmasked += [toks_global_unmasked] if len(tokens_global_unmasked) == 0: - assert False, "Not yet implemented" - tokens_global_unmasked = torch.cat(tokens_global_unmasked) + tokens_global_unmasked = tokens_global.new_empty( + (0, tokens_global.shape[-2], tokens_global.shape[-1]) + ) + else: + tokens_global_unmasked = torch.cat(tokens_global_unmasked) - return tokens_global_unmasked, posteriors + return tokens_global_unmasked, posteriors, empty_chunk_dependency def aggregation_engine_unmasked( self, @@ -273,7 +371,11 @@ def aggregation_engine_unmasked( return tokens_global_unmasked def assimilate_local( - self, model_params, tokens: torch.Tensor, batch: ModelBatch + self, + model_params, + tokens: torch.Tensor, + batch: ModelBatch, + cell_lens_local: torch.Tensor | None = None, ) -> torch.Tensor: """ Processes embedded tokens locally and prepares them for the global assimilation @@ -281,13 +383,14 @@ def assimilate_local( Args: model_params : Query and embedding parameters tokens : Input tokens to be processed by local assimilation - cell_lens : Used to identify range of tokens to use from generated tokens in cell - embedding + cell_lens_local : Per-cell token lengths for this rank's HEALPix domain Returns: Tokens for global assimilation """ - cell_lens = torch.sum(batch.tokens_lens, 2).flatten() + batch_num_cells = batch.tokens_lens.shape[-1] + tokens_lens_global = self._get_global_tokens_lens(batch.tokens_lens) + cell_lens = torch.sum(tokens_lens_global, 2).flatten() num_steps_input = batch.get_num_source_steps() rs = num_steps_input * len(batch) @@ -297,34 +400,79 @@ def assimilate_local( pos_enc = positional_encoding_harmonic tokens_global_register_class = pos_enc(self.q_cells.repeat(rs, num_extra_tokens, 1)) + # Direct calls retain the old API and perform the shard selection here. + # ``forward`` passes an already-sharded stream_cell_tokens tensor. + if cell_lens_local is None: + if batch_num_cells == self.spatial_parallel.local_num_cells: + cell_lens_local = torch.sum(batch.tokens_lens, 2).flatten() + else: + tokens, cell_lens_local = select_packed_cell_shard( + tokens, + cell_lens, + self.num_healpix_cells, + self.spatial_parallel.cell_start, + self.spatial_parallel.cell_end, + ) + + pe_global_local = model_params.pe_global[ + self.spatial_parallel.cell_start : self.spatial_parallel.cell_end + ] + # TODO: re-enable or remove ae_local_queries_per_cell if self.cf.ae_local_queries_per_cell: - tokens_global = (self.q_cells + model_params.pe_global).repeat(rs, 1, 1) + q_cells_local = self.q_cells[ + self.spatial_parallel.cell_start : self.spatial_parallel.cell_end + ] + tokens_global = (q_cells_local + pe_global_local).repeat(rs, 1, 1) else: - num_tokens = self.num_healpix_cells - tokens_global = self.q_cells.repeat(num_tokens, 1, 1) + model_params.pe_global + tokens_global = ( + self.q_cells.repeat(self.spatial_parallel.local_num_cells, 1, 1) + pe_global_local + ) tokens_global = tokens_global.repeat(rs, 1, 1) # apply local assimilation engine and project onto global latent vectors - tokens_global_unmasked, posteriors = self.assimilate_local_project_chunked( - tokens, tokens_global, cell_lens, model_params.q_cells_lens + tokens_global_unmasked, posteriors, empty_chunk_dependency = ( + self.assimilate_local_project_chunked( + tokens, + tokens_global, + cell_lens_local, + model_params.q_cells_lens, + self.spatial_parallel.local_num_cells, + ) ) + tokens_global = tokens_global + empty_chunk_dependency + + # Restore a dense local cell tensor before gathering. This gives every + # rank the same gather shape and preserves the full autograd graph. + local_mask = cell_lens_local.to(torch.bool) + tokens_global[local_mask] = tokens_global_unmasked.to(tokens_global.dtype) + tokens_global = tokens_global.reshape( + rs, + self.spatial_parallel.local_num_cells, + self.q_cells.shape[-2], + self.q_cells.shape[-1], + ) + if self.spatial_parallel.size > 1: + tokens_global = torch.cat( + all_gather(tokens_global, group=self.spatial_parallel.group), + dim=1, + ) + + # Recover packed, globally ordered non-empty cells for query aggregation. + cell_mask = cell_lens.reshape(rs, self.num_healpix_cells).to(torch.bool) + tokens_global_unmasked = tokens_global[cell_mask] # apply aggregation engine on unmasked tokens tokens_global_unmasked = self.aggregation_engine_unmasked( tokens_global_unmasked, tokens_global_register_class, - batch.tokens_lens, + tokens_lens_global, rope_cell_coords=model_params.rope_cell_coords, ) # final processing - tokens_global = ( - torch.permute(tokens_global, [1, 0, 2]) - .squeeze() - .reshape(rs, self.num_healpix_cells, -1) - ) + tokens_global = tokens_global.reshape(rs, self.num_healpix_cells, -1) # TODO, TODO, TODO: do we need this tokens_global = torch.cat([tokens_global_register_class, tokens_global], dim=1) diff --git a/src/weathergen/model/engines.py b/src/weathergen/model/engines.py index fde31213b6..501e0781ce 100644 --- a/src/weathergen/model/engines.py +++ b/src/weathergen/model/engines.py @@ -115,7 +115,10 @@ def forward(self, batch, pe_embed): ) " Increase ae_local_max_tokens_per_cell in config." - if batch.tokens_lens.shape[2] == 1: + if not x_embeds: + # This rank has no local observations. + return tokens_all + elif batch.tokens_lens.shape[2] == 1: # trivial with one stream tokens_all = torch.cat(x_embeds) diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py index f4035ea467..f2d7d73932 100644 --- a/src/weathergen/model/model.py +++ b/src/weathergen/model/model.py @@ -18,7 +18,9 @@ import astropy_healpix.healpy import numpy as np import torch +import torch.distributed as dist import torch.nn as nn +from torch.distributed.nn.functional import all_gather from torch.utils.checkpoint import checkpoint from weathergen.common.config import Config @@ -38,8 +40,14 @@ TargetPredictionEngineClassic, ) from weathergen.model.layers import MLP, NamedLinear +from weathergen.model.spatial_parallel import ( + reassemble_packed_cell_shards, + select_healpix_neighborhood_shard, + select_packed_cell_shard, + split_cell_lens_by_shard, +) from weathergen.model.utils import get_num_parameters -from weathergen.utils.distributed import is_root +from weathergen.utils.distributed import SpatialParallelContext, is_root from weathergen.utils.utils import get_dtype, is_stream_forcing logger = logging.getLogger(__name__) @@ -321,6 +329,9 @@ def __init__(self, cf: Config, sources_size, targets_num_channels, targets_coord self.sources_size = sources_size self.targets_num_channels = targets_num_channels self.targets_coords_size = targets_coords_size + self.target_coords_size_by_stream = dict( + zip(cf.streams.keys(), targets_coords_size, strict=True) + ) self.embed_target_coords = None self.encoder: EncoderModule | None = None @@ -329,6 +340,7 @@ def __init__(self, cf: Config, sources_size, targets_num_channels, targets_coord self.q_cells: torch.Tensor | None = None self.streams: dict[str, typing.Any] = cf.streams self.target_token_engines = None + self.spatial_parallel: SpatialParallelContext | None = None assert cf.get("forecast", {}).get("att_dense_rate", 1.0) == 1.0, ( "Local attention not adapted for register tokens" @@ -376,6 +388,8 @@ def create(self) -> "Model": self.encoder = EncoderModule( cf, self.sources_size, self.targets_num_channels, self.targets_coords_size ) + # Encoder and decoder share one homogeneous HEALPix decomposition. + self.spatial_parallel = self.encoder.spatial_parallel mode_cfg = cf.training_config if cf.fe_num_blocks > 0: @@ -758,18 +772,31 @@ def predict_decoders( # Empty dicts evaluate to False in python if not self.pred_heads: return output + assert self.spatial_parallel is not None + spatial_parallel = self.spatial_parallel # remove register and class tokens tokens = tokens[:, self.num_aux_tokens :] - # get 1-ring neighborhood for prediction + # Get the 1-ring neighbourhoods for the HEALPix cells owned by this + # spatial rank. Each row contains the cell itself plus eight neighbours. batch_size = len(batch) s = [batch_size, self.num_healpix_cells, self.cf.ae_local_num_queries, tokens.shape[-1]] - idxs = model_params.hp_nbours.unsqueeze(0).repeat((batch_size, 1, 1)).flatten(0, 1) - tokens_nbors = tokens.reshape(s).flatten(0, 1)[idxs.flatten()].flatten(0, 1) - # TODO: precompute in model_params? + local_hp_nbours = model_params.hp_nbours[ + spatial_parallel.cell_start : spatial_parallel.cell_end + ] + tokens_nbors = select_healpix_neighborhood_shard( + tokens.reshape(s), + model_params.hp_nbours, + spatial_parallel.cell_start, + spatial_parallel.cell_end, + ) + num_decoder_cells = batch_size * spatial_parallel.local_num_cells tokens_nbors_lens = torch.full( - (s[0] * s[1] + 1,), fill_value=9, dtype=torch.int32, device=tokens_nbors.device + (num_decoder_cells + 1,), + fill_value=local_hp_nbours.shape[1], + dtype=torch.int32, + device=tokens_nbors.device, ) tokens_nbors_lens[0] = 0 @@ -782,58 +809,123 @@ def predict_decoders( ] t_coords_lens = [len(t) for t in t_coords] t_coords = torch.cat(t_coords) - if len(t_coords) == 0: continue + tcls_global = torch.stack( + [ + sample.streams_data[stream_name].target_coords_lens[step] + for sample in batch.samples + ] + ) + t_coords, tcls = select_packed_cell_shard( + t_coords, + tcls_global.flatten(), + self.num_healpix_cells, + spatial_parallel.cell_start, + spatial_parallel.cell_end, + ) + + local_coords_empty = len(t_coords) == 0 + tcls_compute = tcls + if local_coords_empty: + # FSDP-wrapped decoder modules must be entered by every spatial + # rank. Run one dummy coordinate through them, then discard it. + t_coords = tokens.new_zeros((1, self.target_coords_size_by_stream[stream_name])) + tcls_compute = torch.zeros_like(tcls) + tcls_compute[0] = 1 # embed token coords tc_embed = self.embed_target_coords[stream_name] tc_tokens = checkpoint(tc_embed, t_coords, use_reentrant=False) - # skip when coordinate embeddings yields nan (i.e. the coord embedding network diverged) - if torch.isnan(tc_tokens).any(): - logger.warning( - ( - f"Skipping prediction for {stream_name} because", - f" of {torch.isnan(tc_tokens).sum()} NaN in tc_tokens.", - ) + # A rank-local NaN must stop every peer before any rank enters the + # FSDP-wrapped prediction engine and head. + has_nan = torch.isnan(tc_tokens).any().to(dtype=torch.int32) + if spatial_parallel.size > 1: + dist.all_reduce( + has_nan, + op=dist.ReduceOp.MAX, + group=spatial_parallel.group, + ) + if has_nan.item(): + raise FloatingPointError( + f"NaN in target-coordinate embedding for {stream_name} at step {step}" ) - pred = torch.tensor([], device=tc_tokens.device) - # skip empty lengths - elif tc_tokens.shape[0] == 0: - pred = torch.tensor([], device=tc_tokens.device) + # lens for varlen attention + tcs_lens = torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=tcls.device), + tcls_compute, + ] + ) + if self.cf.decoder_type == "Linear": + if self.cf.ae_local_num_queries != 1: + raise NotImplementedError("Linear decoder requires ae_local_num_queries == 1") + local_tokens = tokens.reshape(s)[ + :, + spatial_parallel.cell_start : spatial_parallel.cell_end, + ].reshape(-1, s[-1]) + pred = self.target_token_engines[stream_name]( + tc_tokens, + local_tokens, + tcs_lens, + ).unsqueeze(0) else: - # lens for varlen attention - tcls = torch.cat( - [ - sample.streams_data[stream_name].target_coords_lens[step] - for sample in batch.samples - ] + tc_tokens = self.target_token_engines[stream_name]( + latent=tokens_nbors, + output=tc_tokens, + latent_lens=tokens_nbors_lens, + output_lens=tcs_lens, + coordinates=t_coords, ) - tcs_lens = torch.cat([torch.zeros(1, dtype=torch.int32, device=tcls.device), tcls]) - - if self.cf.decoder_type == "Linear": - pred = self.target_token_engines[stream_name]( - tc_tokens, - tokens.reshape(-1, s[-1]), # collapse the batch and token dimensions - tcs_lens, - ).unsqueeze(0) # add ensemble dim: shape is then [1, preds_per_coord, channels] - else: - tc_tokens = self.target_token_engines[stream_name]( - latent=tokens_nbors, - output=tc_tokens, - latent_lens=tokens_nbors_lens, - output_lens=tcs_lens, - coordinates=t_coords, - ) - # final prediction head to map back to physical space - pred = self.pred_heads[stream_name](tc_tokens) + # final prediction head to map back to physical space + pred = self.pred_heads[stream_name](tc_tokens) + + if local_coords_empty: + pred = pred[:, :0] + pred = self._gather_decoder_predictions(pred, tcls_global) # recover batch dimension (ragged, so as list) pred = torch.split(pred, t_coords_lens, dim=1) output.add_physical_prediction(step, stream_name, pred) return output + + def _gather_decoder_predictions( + self, pred: torch.Tensor, global_cell_lens: torch.Tensor + ) -> torch.Tensor: + """Gather packed cell-level decoder output in sample/cell order.""" + + assert self.spatial_parallel is not None + spatial_parallel = self.spatial_parallel + if spatial_parallel.size == 1: + return pred + + gathered_lens = split_cell_lens_by_shard( + global_cell_lens, + spatial_parallel.size, + ) + packed_lens = [int(lens.sum().item()) for lens in gathered_lens] + max_len = max(packed_lens) + if pred.shape[1] < max_len: + pred = torch.cat( + [ + pred, + pred.new_zeros( + pred.shape[0], + max_len - pred.shape[1], + pred.shape[2], + ), + ], + dim=1, + ) + gathered_pred = all_gather(pred, group=spatial_parallel.group) + + return reassemble_packed_cell_shards( + list(gathered_pred), + gathered_lens, + spatial_parallel.local_num_cells, + ) diff --git a/src/weathergen/model/model_interface.py b/src/weathergen/model/model_interface.py index c0b475b156..dccd9c3f78 100644 --- a/src/weathergen/model/model_interface.py +++ b/src/weathergen/model/model_interface.py @@ -48,10 +48,11 @@ def init_model_and_shard( mini_epoch_contd, training_mode, device, - with_ddp, - with_fsdp, + data_parallel_cfg, overrides={}, ): + with_ddp = data_parallel_cfg["with_ddp"] + with_fsdp = data_parallel_cfg["with_fsdp"] model_creation_device = "meta" if with_ddp and with_fsdp else "cuda" with torch.device(model_creation_device): model = get_model(cf, training_mode, dataset, overrides) @@ -72,7 +73,7 @@ def init_model_and_shard( model = torch.nn.parallel.DistributedDataParallel( model, broadcast_buffers=True, - find_unused_parameters=cf.get("ddp_find_unused_parameters", True), + find_unused_parameters=data_parallel_cfg.get("find_unused_parameters", True), gradient_as_bucket_view=True, bucket_cap_mb=512, ) @@ -191,7 +192,9 @@ def load_model(cf, model, device, run_id: str, mini_epoch=-1): path_run / filename, map_location=torch.device("cpu"), mmap=True, weights_only=True ) - is_model_sharded = cf.with_ddp and cf.with_fsdp + is_model_sharded = ( + cf.distributed.data_parallel.with_ddp and cf.distributed.data_parallel.with_fsdp + ) if is_model_sharded: meta_sharded_sd = model.state_dict() maybe_sharded_sd = {} diff --git a/src/weathergen/model/spatial_parallel.py b/src/weathergen/model/spatial_parallel.py new file mode 100644 index 0000000000..04c9f4f021 --- /dev/null +++ b/src/weathergen/model/spatial_parallel.py @@ -0,0 +1,123 @@ +# (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. + +import torch + + +def select_packed_cell_shard( + tokens: torch.Tensor, + cell_lens: torch.Tensor, + num_cells: int, + cell_start: int, + cell_end: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Select a HEALPix-cell range from a cell-packed token tensor. + + ``cell_lens`` is flattened in ``(input_step, sample, cell)`` order. Token + counts vary by cell, so slicing the first tensor dimension directly would + split cells. The returned tokens remain packed in the same order, restricted + to ``[cell_start, cell_end)`` for every input-step/sample row. + """ + + if cell_lens.numel() % num_cells: + raise ValueError("cell_lens does not contain a whole number of HEALPix grids") + if not 0 <= cell_start < cell_end <= num_cells: + raise ValueError( + f"invalid HEALPix cell range [{cell_start}, {cell_end}) for {num_cells} cells" + ) + + cell_lens_2d = cell_lens.reshape(-1, num_cells) + selected_cells = torch.zeros_like(cell_lens_2d, dtype=torch.bool) + selected_cells[:, cell_start:cell_end] = True + selected_tokens = torch.repeat_interleave( + selected_cells.flatten(), cell_lens.to(dtype=torch.long) + ) + if selected_tokens.numel() != tokens.shape[0]: + raise ValueError( + f"packed token length mismatch: cell_lens describes {selected_tokens.numel()} " + f"tokens, tensor has {tokens.shape[0]}" + ) + + return tokens[selected_tokens], cell_lens_2d[:, cell_start:cell_end].flatten() + + +def select_healpix_neighborhood_shard( + tokens: torch.Tensor, + hp_neighbours: torch.Tensor, + cell_start: int, + cell_end: int, +) -> torch.Tensor: + """Pack nine-token HEALPix neighbourhoods for one contiguous cell shard. + + Args: + tokens: Dense tensor in ``(sample, cell, query, channel)`` order. + hp_neighbours: Global ``(cell, 9)`` lookup containing self and its + eight neighbours. + cell_start: First cell owned by the spatial rank. + cell_end: Exclusive end of the rank's cell range. + """ + + if tokens.ndim != 4: + raise ValueError("tokens must have shape (sample, cell, query, channel)") + num_cells = tokens.shape[1] + if hp_neighbours.shape != (num_cells, 9): + raise ValueError( + f"expected HEALPix neighbour lookup with shape ({num_cells}, 9), " + f"got {tuple(hp_neighbours.shape)}" + ) + if not 0 <= cell_start < cell_end <= num_cells: + raise ValueError( + f"invalid HEALPix cell range [{cell_start}, {cell_end}) for {num_cells} cells" + ) + + local_neighbours = hp_neighbours[cell_start:cell_end].to(dtype=torch.long) + return tokens[:, local_neighbours].flatten(0, 3) + + +def reassemble_packed_cell_shards( + shards: list[torch.Tensor], + shard_cell_lens: list[torch.Tensor], + cells_per_shard: int, +) -> torch.Tensor: + """Reassemble rank-packed predictions into sample-major HEALPix order.""" + + if len(shards) != len(shard_cell_lens) or not shards: + raise ValueError("shards and shard_cell_lens must be non-empty and have equal length") + if any(lens.numel() % cells_per_shard for lens in shard_cell_lens): + raise ValueError("shard_cell_lens does not contain whole local HEALPix grids") + + batch_size = shard_cell_lens[0].numel() // cells_per_shard + if any(lens.numel() != batch_size * cells_per_shard for lens in shard_cell_lens): + raise ValueError("all spatial shards must describe the same batch size") + + rank_offsets = [0] * len(shards) + sample_predictions = [] + for sample_idx in range(batch_size): + sample_shards = [] + for rank, (rank_pred, rank_lens) in enumerate(zip(shards, shard_cell_lens, strict=True)): + lens = rank_lens.reshape(batch_size, cells_per_shard) + sample_len = int(lens[sample_idx].sum().item()) + offset = rank_offsets[rank] + sample_shards.append(rank_pred[:, offset : offset + sample_len]) + rank_offsets[rank] += sample_len + sample_predictions.append(torch.cat(sample_shards, dim=1)) + return torch.cat(sample_predictions, dim=1) + + +def split_cell_lens_by_shard(cell_lens: torch.Tensor, num_shards: int) -> list[torch.Tensor]: + """Split global ``(sample, cell)`` lengths into contiguous rank-local grids.""" + + if cell_lens.ndim != 2: + raise ValueError("cell_lens must have shape (sample, cell)") + if num_shards < 1 or cell_lens.shape[1] % num_shards: + raise ValueError( + f"number of cells ({cell_lens.shape[1]}) must be divisible by num_shards ({num_shards})" + ) + + return [shard.flatten() for shard in torch.chunk(cell_lens, num_shards, dim=1)] diff --git a/src/weathergen/train/target_and_aux_ssl_teacher.py b/src/weathergen/train/target_and_aux_ssl_teacher.py index edd8e53b6c..b50bb030b9 100644 --- a/src/weathergen/train/target_and_aux_ssl_teacher.py +++ b/src/weathergen/train/target_and_aux_ssl_teacher.py @@ -27,6 +27,7 @@ load_encoder_from_checkpoint, prepare_encoder_teacher, ) +from weathergen.utils.distributed import normalize_distributed_config logger = logging.getLogger(__name__) @@ -142,7 +143,15 @@ def from_pretrained(cls, cf: Config, dataset, device, params: dict) -> FrozenTea # Load teacher's config, create model with teacher's architecture teacher_config = load_run_config(teacher_run_id, teacher_mini_epoch, model_path=None) - teacher_config = merge_configs(teacher_config, {"with_ddp": False, "with_fsdp": False}) + teacher_config = normalize_distributed_config(teacher_config) + teacher_config = merge_configs( + teacher_config, + { + "distributed": { + "data_parallel": {"with_ddp": False, "with_fsdp": False} + } + }, + ) teacher_model = get_model(teacher_config, "student", dataset, {}) diff --git a/src/weathergen/train/target_and_aux_utils.py b/src/weathergen/train/target_and_aux_utils.py index efaff18bce..0d8d0cac2b 100644 --- a/src/weathergen/train/target_and_aux_utils.py +++ b/src/weathergen/train/target_and_aux_utils.py @@ -35,7 +35,9 @@ def get_target_aux_calculator( elif target_and_aux_calc == "EMATeacher": # work around for problems with FSDP2 - assert not cf.with_fsdp, "EMATeacher not supported with FSDP(2) at the moment" + assert not cf.distributed.data_parallel.with_fsdp, ( + "EMATeacher not supported with FSDP(2) at the moment" + ) meta_ema_model, _ = init_model_and_shard( cf, @@ -44,8 +46,7 @@ def get_target_aux_calculator( None, "student", device, - with_ddp=False, - with_fsdp=False, + {"with_ddp": False, "with_fsdp": False}, overrides=target_and_aux_calc_params.get("model_param_overrides", {}), ) @@ -60,7 +61,10 @@ def get_target_aux_calculator( meta_ema_model, halflife_steps=target_and_aux_calc_params.get("ema_halflife_in_thousands", 1e-3), rampup_ratio=target_and_aux_calc_params.get("ema_ramp_up_ratio", 0.09), - is_model_sharded=(cf.with_ddp and cf.with_fsdp), + is_model_sharded=( + cf.distributed.data_parallel.with_ddp + and cf.distributed.data_parallel.with_fsdp + ), ) batch_size = cf.get("world_size_original", cf.get("world_size")) * batch_size_per_gpu diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py index 276da0bd67..7001afcc45 100644 --- a/src/weathergen/train/trainer.py +++ b/src/weathergen/train/trainer.py @@ -47,7 +47,7 @@ get_batch_size_from_config, get_target_idxs_from_cfg, ) -from weathergen.utils.distributed import is_root +from weathergen.utils.distributed import get_spatial_parallel_size, is_root from weathergen.utils.performance import NullThroughputTracker, ThroughputTracker, nvtx_range from weathergen.utils.train_logger import TrainLogger, prepare_losses_for_logging from weathergen.utils.utils import get_dtype @@ -95,7 +95,7 @@ def get_batch_size_total(self, batch_size_per_gpu) -> int: """ Get total, effective batch size across all DDP ranks """ - return self.world_size_original * batch_size_per_gpu + return self.data_parallel_world_size_original * batch_size_per_gpu def init(self, cf: Config, devices): # pylint: disable=attribute-defined-outside-init @@ -151,6 +151,20 @@ def init(self, cf: Config, devices): # world_size gets overwritten by current setting during init_ddp() self.world_size_original = cf.get("world_size_original", cf.get("world_size", None)) cf.world_size_original = self.world_size_original + spatial_cfg = cf.distributed.spatial_parallel + data_parallel_cfg = cf.distributed.data_parallel + spatial_parallel_size = get_spatial_parallel_size(spatial_cfg) + data_parallel_cfg.world_size = cf.world_size // spatial_parallel_size + spatial_parallel_size_original = spatial_cfg.get("size_original", spatial_parallel_size) + if self.world_size_original % spatial_parallel_size_original: + raise ValueError( + "world_size_original must be divisible by " + "distributed.spatial_parallel.size_original" + ) + self.data_parallel_world_size_original = ( + self.world_size_original // spatial_parallel_size_original + ) + spatial_cfg.size_original = spatial_parallel_size_original self.log_grad_norms = cf.train_logging.get("log_grad_norms", False) @@ -229,8 +243,7 @@ def inference(self, cf, devices, run_id_contd, mini_epoch_contd): mini_epoch_contd, self.test_cfg.training_mode, devices[0], - cf.with_ddp, - cf.with_fsdp, + cf.distributed.data_parallel, ) # get target_aux calculators for different loss terms @@ -280,8 +293,7 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None): mini_epoch_contd, self.training_cfg.training_mode, devices[0], - cf.with_ddp, - cf.with_fsdp, + cf.distributed.data_parallel, ) validate_with_ema_cfg = self.validation_cfg.get("validate_with_ema") @@ -299,15 +311,17 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None): mini_epoch_contd, cf.training_config.training_mode, devices[0], - cf.with_ddp, - cf.with_fsdp, + cf.distributed.data_parallel, ) self.ema_model = EMAModel( self.model, meta_ema_model, halflife_steps=validate_with_ema_cfg.get("ema_halflife_in_thousands", 1e-3), rampup_ratio=validate_with_ema_cfg.get("ema_ramp_up_ratio", 0.09), - is_model_sharded=(cf.with_ddp and cf.with_fsdp), + is_model_sharded=( + cf.distributed.data_parallel.with_ddp + and cf.distributed.data_parallel.with_fsdp + ), ) # get target_aux calculators for different loss terms @@ -317,7 +331,7 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None): # if with_fsdp then parameter count is unreliable if is_root(): # ddp-wrapped model does not expose this function - if not cf.with_ddp: + if not cf.distributed.data_parallel.with_ddp: self.model.print_num_parameters() # https://www.cs.princeton.edu/~smalladi/blog/2024/01/22/SDEs-ScalingRules/ @@ -349,7 +363,7 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None): self.lr_scheduler = LearningRateScheduler( self.optimizer, self.batch_size_per_gpu, - cf.world_size, + cf.distributed.data_parallel.world_size, cf.general.istep, lr_steps, self.training_cfg.learning_rate_scheduling, @@ -368,13 +382,17 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None): mini_epoch_base = int(self.cf.general.istep / len(self.data_loader)) else: len_per_rank = ( - max(1, len(self.dataset) // (self.world_size_original * self.batch_size_per_gpu)) + max( + 1, + len(self.dataset) + // (self.data_parallel_world_size_original * self.batch_size_per_gpu), + ) ) * self.batch_size_per_gpu mini_epoch_base = int( self.cf.general.istep / ( min(len_per_rank, self.training_cfg.samples_per_mini_epoch) - * self.world_size_original + * self.data_parallel_world_size_original ) ) @@ -585,7 +603,8 @@ def validate(self, mini_epoch, mode_cfg, batch_size): with torch.no_grad(): # print progress bar but only in interactive mode, i.e. when without ddp with tqdm.tqdm( - total=len(self.data_loader_validation), disable=self.cf.with_ddp + total=len(self.data_loader_validation), + disable=self.cf.distributed.data_parallel.with_ddp, ) as pbar: for bidx, batch in enumerate(dataset_val_iter): batch.to_device(self.device) @@ -659,7 +678,10 @@ def _get_full_model_state_dict(self): maybe_sharded_sd = ( self.model.state_dict() if self.ema_model is None else self.ema_model.state_dict() ) - if self.cf.with_ddp and self.cf.with_fsdp: + if ( + self.cf.distributed.data_parallel.with_ddp + and self.cf.distributed.data_parallel.with_fsdp + ): cpu_state_dict = {} for param_name, sharded_param in maybe_sharded_sd.items(): full_param = sharded_param.full_tensor() diff --git a/src/weathergen/train/trainer_base.py b/src/weathergen/train/trainer_base.py index d0359892a3..eb705f96e9 100644 --- a/src/weathergen/train/trainer_base.py +++ b/src/weathergen/train/trainer_base.py @@ -17,7 +17,7 @@ from weathergen.common.config import Config from weathergen.train.utils import str_to_tensor, tensor_to_str -from weathergen.utils.distributed import is_root +from weathergen.utils.distributed import is_root, normalize_distributed_config PORT = 1345 @@ -65,6 +65,7 @@ def init_torch(use_cuda=True, num_accs_per_task=1, multiprocessing_method="fork" @staticmethod def init_ddp(cf): """Initializes the distributed environment.""" + cf = normalize_distributed_config(cf) rank = 0 local_rank = 0 @@ -139,6 +140,6 @@ def init_ddp(cf): cf.world_size = world_size cf.rank = rank cf.local_rank = local_rank - cf.with_ddp = world_size > 1 + cf.distributed.data_parallel.with_ddp = world_size > 1 return cf diff --git a/src/weathergen/utils/distributed.py b/src/weathergen/utils/distributed.py index af467a3a94..5e90786f9a 100644 --- a/src/weathergen/utils/distributed.py +++ b/src/weathergen/utils/distributed.py @@ -8,10 +8,55 @@ # nor does it submit to any jurisdiction. +import dataclasses + import torch import torch.distributed as dist SYNC_TIMEOUT_SEC = 60 * 60 # 1 hour +_SPATIAL_GROUPS: dict[int, tuple[dist.ProcessGroup, int]] = {} + + +def normalize_distributed_config(cf): + """Move legacy top-level distributed settings into their shared sections.""" + + if cf.get("distributed") is None: + cf["distributed"] = {} + distributed_cfg = cf["distributed"] + + if distributed_cfg.get("data_parallel") is None: + distributed_cfg["data_parallel"] = {} + data_parallel_cfg = distributed_cfg["data_parallel"] + if "with_ddp" not in data_parallel_cfg: + data_parallel_cfg["with_ddp"] = cf.get("with_ddp", False) + if "with_fsdp" not in data_parallel_cfg: + data_parallel_cfg["with_fsdp"] = cf.get("with_fsdp", False) + if "find_unused_parameters" not in data_parallel_cfg: + data_parallel_cfg["find_unused_parameters"] = cf.get( + "ddp_find_unused_parameters", True + ) + if "world_size" not in data_parallel_cfg and cf.get("data_parallel_world_size") is not None: + data_parallel_cfg["world_size"] = cf["data_parallel_world_size"] + + if distributed_cfg.get("spatial_parallel") is None: + distributed_cfg["spatial_parallel"] = {} + spatial_cfg = distributed_cfg["spatial_parallel"] + if "size" not in spatial_cfg: + spatial_cfg["size"] = cf.get("spatial_parallel_size", 1) + if "size_original" not in spatial_cfg and cf.get("spatial_parallel_size_original") is not None: + spatial_cfg["size_original"] = cf["spatial_parallel_size_original"] + + for legacy_key in ( + "with_ddp", + "with_fsdp", + "data_parallel_world_size", + "ddp_find_unused_parameters", + "spatial_parallel_size", + "spatial_parallel_size_original", + ): + cf.pop(legacy_key, None) + + return cf def is_root(pg: dist.ProcessGroup | None = None) -> bool: @@ -59,6 +104,115 @@ def get_rank() -> int: return dist.get_rank() +def get_spatial_parallel_size(spatial_cfg) -> int: + """Return and validate the configured spatial-parallel size.""" + + size = int(spatial_cfg.get("size", 1)) + if size < 1: + raise ValueError("distributed.spatial_parallel.size must be at least 1") + + world_size = get_world_size() + if size > world_size: + raise ValueError( + f"distributed.spatial_parallel.size ({size}) exceeds world_size ({world_size})" + ) + if world_size % size: + raise ValueError( + "world_size " + f"({world_size}) must be divisible by distributed.spatial_parallel.size ({size})" + ) + return size + + +def get_spatial_parallel_group(spatial_cfg) -> tuple[dist.ProcessGroup | None, int]: + """Create the consecutive-rank process groups used to shard HEALPix cells. + + All ranks call ``new_group`` in the same order. The returned rank is local to + the spatial group. A size of one deliberately avoids creating a process group. + """ + + size = get_spatial_parallel_size(spatial_cfg) + if size == 1: + return None, 0 + if not _is_distributed_initialized(): + raise RuntimeError("spatial parallelism requires torch.distributed") + + cached = _SPATIAL_GROUPS.get(size) + if cached is not None: + return cached + + world_size = dist.get_world_size() + global_rank = dist.get_rank() + own_group = None + for first_rank in range(0, world_size, size): + ranks = list(range(first_rank, first_rank + size)) + group = dist.new_group(ranks=ranks) + if global_rank in ranks: + own_group = group + + assert own_group is not None + result = (own_group, global_rank % size) + _SPATIAL_GROUPS[size] = result + return result + + +@dataclasses.dataclass(frozen=True) +class SpatialParallelContext: + """Distributed topology and HEALPix ownership for one spatial rank.""" + + size: int + rank: int + group: dist.ProcessGroup | None + ddp_rank: int + ddp_world_size: int + num_cells: int + local_num_cells: int + cell_start: int + cell_end: int + + @classmethod + def from_config( + cls, + cf, + num_cells: int, + *, + create_process_group: bool = False, + ) -> "SpatialParallelContext": + """Build and validate the shared spatial-parallel topology.""" + + spatial_cfg = cf.distributed.spatial_parallel + size = get_spatial_parallel_size(spatial_cfg) + if num_cells % size: + raise ValueError( + f"number of HEALPix cells ({num_cells}) must be divisible by " + f"distributed.spatial_parallel.size ({size})" + ) + + rank = cf.rank % size + group = None + if create_process_group: + group, group_rank = get_spatial_parallel_group(spatial_cfg) + if group_rank != rank: + raise RuntimeError( + f"spatial process-group rank ({group_rank}) does not match " + f"topology rank ({rank})" + ) + + local_num_cells = num_cells // size + cell_start = rank * local_num_cells + return cls( + size=size, + rank=rank, + group=group, + ddp_rank=cf.rank // size, + ddp_world_size=cf.world_size // size, + num_cells=num_cells, + local_num_cells=local_num_cells, + cell_start=cell_start, + cell_end=cell_start + local_num_cells, + ) + + def ddp_average(data: torch.Tensor) -> torch.Tensor: """ Average a tensor across DDP ranks diff --git a/tests/test_spatial_parallel.py b/tests/test_spatial_parallel.py new file mode 100644 index 0000000000..eba3c46a75 --- /dev/null +++ b/tests/test_spatial_parallel.py @@ -0,0 +1,245 @@ +# (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. + +import types + +import numpy as np +import pytest +import torch + +from weathergen.datasets.healpix_domain import build_local_healpix_cell_splits +from weathergen.model.spatial_parallel import ( + reassemble_packed_cell_shards, + select_healpix_neighborhood_shard, + select_packed_cell_shard, + split_cell_lens_by_shard, +) +from weathergen.utils import distributed + + +def test_local_healpix_construction_matches_global_cell_slices(): + num_cells = 48 + cell_ids = np.repeat(np.arange(num_cells), np.arange(num_cells) % 3 + 1) + rng = np.random.default_rng(7) + cell_ids = cell_ids[rng.permutation(len(cell_ids))] + global_cells = build_local_healpix_cell_splits( + cell_ids, + num_cells, + cell_start=0, + cell_end=num_cells, + ) + cells_per_rank = len(global_cells) // 4 + + local_cells_all = [] + for spatial_rank in range(4): + cell_start = spatial_rank * cells_per_rank + cell_end = cell_start + cells_per_rank + local_cells = build_local_healpix_cell_splits( + cell_ids, + num_cells, + cell_start=cell_start, + cell_end=cell_end, + ) + + assert len(local_cells) == cells_per_rank + for local_cell, global_cell in zip( + local_cells, + global_cells[cell_start:cell_end], + strict=True, + ): + np.testing.assert_array_equal(local_cell, global_cell) + local_cells_all.extend(local_cells) + + assert len(local_cells_all) == len(global_cells) + for local_cell, global_cell in zip(local_cells_all, global_cells, strict=True): + np.testing.assert_array_equal(local_cell, global_cell) + + +def test_local_healpix_construction_rejects_invalid_range(): + with pytest.raises(ValueError, match="invalid HEALPix cell range"): + build_local_healpix_cell_splits( + np.arange(12), + num_cells=12, + cell_start=6, + cell_end=13, + ) + + +def test_local_healpix_construction_handles_empty_domain(): + cell_splits = build_local_healpix_cell_splits( + np.array([0, 1, 2], dtype=np.int64), + num_cells=12, + cell_start=6, + cell_end=9, + ) + + assert len(cell_splits) == 3 + assert all(cell.dtype == np.int64 and cell.size == 0 for cell in cell_splits) + + +def test_select_packed_cell_shard_preserves_cell_boundaries_across_rows(): + cell_lens = torch.tensor( + [ + [1, 0, 2, 1, 3, 0, 1, 2], + [0, 2, 1, 0, 1, 2, 0, 1], + ], + dtype=torch.int32, + ) + tokens = torch.arange(cell_lens.sum(), dtype=torch.float32).unsqueeze(1) + + shard, shard_lens = select_packed_cell_shard( + tokens, cell_lens.flatten(), num_cells=8, cell_start=2, cell_end=4 + ) + + assert shard_lens.tolist() == [2, 1, 1, 0] + assert shard.squeeze(1).tolist() == [1, 2, 3, 12] + + +def test_eight_shards_cover_every_packed_token_once_and_keep_gradients(): + num_cells = 16 + cell_lens = torch.tensor( + [ + [0, 1, 2, 0, 1, 3, 0, 2, 1, 0, 2, 1, 0, 1, 2, 1], + [1, 0, 1, 2, 0, 1, 2, 0, 3, 1, 0, 1, 2, 0, 1, 1], + ], + dtype=torch.int32, + ) + tokens = torch.arange(cell_lens.sum(), dtype=torch.float32, requires_grad=True) + shard_width = num_cells // 8 + + selected = [] + for rank in range(8): + shard, _ = select_packed_cell_shard( + tokens, + cell_lens.flatten(), + num_cells, + rank * shard_width, + (rank + 1) * shard_width, + ) + selected.append(shard) + shard.sum().backward(retain_graph=rank < 7) + + assert sum(shard.numel() for shard in selected) == tokens.numel() + assert torch.equal(tokens.grad, torch.ones_like(tokens)) + + +@pytest.mark.parametrize( + ("num_cells", "cell_start", "cell_end"), + [(8, -1, 1), (8, 3, 3), (8, 0, 9)], +) +def test_select_packed_cell_shard_rejects_invalid_ranges(num_cells, cell_start, cell_end): + with pytest.raises(ValueError, match="invalid HEALPix cell range"): + select_packed_cell_shard( + torch.arange(num_cells), + torch.ones(num_cells, dtype=torch.int32), + num_cells, + cell_start, + cell_end, + ) + + +def test_spatial_parallel_size_requires_whole_rank_groups(monkeypatch): + monkeypatch.setattr(distributed, "get_world_size", lambda: 16) + assert distributed.get_spatial_parallel_size({"size": 4}) == 4 + assert distributed.get_spatial_parallel_size({"size": 8}) == 8 + + with pytest.raises(ValueError, match="must be divisible"): + distributed.get_spatial_parallel_size({"size": 6}) + + +def test_normalize_distributed_config_moves_legacy_settings(): + config = { + "with_ddp": True, + "with_fsdp": True, + "data_parallel_world_size": 2, + "ddp_find_unused_parameters": False, + "spatial_parallel_size": 4, + "spatial_parallel_size_original": 8, + } + + result = distributed.normalize_distributed_config(config) + + assert result == { + "distributed": { + "data_parallel": { + "with_ddp": True, + "with_fsdp": True, + "find_unused_parameters": False, + "world_size": 2, + }, + "spatial_parallel": {"size": 4, "size_original": 8}, + } + } + + +def test_spatial_parallel_context_collects_topology_and_cell_ownership(monkeypatch): + monkeypatch.setattr(distributed, "get_world_size", lambda: 16) + config = types.SimpleNamespace( + distributed=types.SimpleNamespace(spatial_parallel={"size": 4}), + rank=6, + world_size=16, + ) + + context = distributed.SpatialParallelContext.from_config(config, num_cells=48) + + assert context.size == 4 + assert context.rank == 2 + assert context.ddp_rank == 1 + assert context.ddp_world_size == 4 + assert context.local_num_cells == 12 + assert (context.cell_start, context.cell_end) == (24, 36) + assert context.group is None + + +def test_decoder_selects_nine_neighbours_for_each_rank_local_cell(): + tokens = torch.arange(2 * 4 * 1 * 2).reshape(2, 4, 1, 2) + neighbours = torch.tensor( + [ + [0, 1, 2, 3, 0, 1, 2, 3, 0], + [1, 0, 2, 3, 1, 0, 2, 3, 1], + [2, 0, 1, 3, 2, 0, 1, 3, 2], + [3, 0, 1, 2, 3, 0, 1, 2, 3], + ] + ) + + selected = select_healpix_neighborhood_shard(tokens, neighbours, 1, 3) + expected = tokens[:, neighbours[1:3]].flatten(0, 3) + + assert selected.shape == (2 * 2 * 9, 2) + assert torch.equal(selected, expected) + + +def test_decoder_reassembles_rank_shards_in_sample_cell_order(): + # Rank packing is [sample 0 local cells, sample 1 local cells]. + rank_0 = torch.tensor([[[0.0], [1.0], [4.0]]]) + rank_1 = torch.tensor([[[2.0], [3.0], [5.0], [6.0]]]) + rank_0_lens = torch.tensor([1, 1, 1, 0], dtype=torch.int32) + rank_1_lens = torch.tensor([1, 1, 1, 2], dtype=torch.int32) + + assembled = reassemble_packed_cell_shards( + [rank_0, rank_1], + [rank_0_lens, rank_1_lens], + cells_per_shard=2, + ) + + assert assembled.squeeze().tolist() == [0, 1, 2, 3, 4, 5, 6] + + +def test_decoder_derives_rank_lens_from_global_target_lens(): + global_lens = torch.tensor( + [ + [1, 0, 2, 3], + [0, 4, 1, 2], + ], + dtype=torch.int32, + ) + + shards = split_cell_lens_by_shard(global_lens, num_shards=2) + + assert [shard.tolist() for shard in shards] == [[1, 0, 0, 4], [2, 3, 1, 2]]