Feature/encoder and decoder healpix cell parallelism by pkuyyj · Pull Request #2705 · ecmwf/WeatherGenerator

pkuyyj · 2026-07-31T09:58:01Z

Description

Based on #2700, we further implemented the spatial parallelism of the decoder. This is also done during the Hackathon in Zurich.

Memory allocation profiling results:

Before:

After encoder parallelization (16 ranks):

After encoder and decoder parallelization (16 ranks):

Issue Number

Is this PR a draft? Mark it as draft.

Checklist before asking for review

I have performed a self-review of my code
My changes comply with basic sanity checks:
- I have fixed formatting issues with ./scripts/actions.sh lint
- I have run unit tests with ./scripts/actions.sh unit-test
- I have documented my code and I have updated the docstrings.
- I have added unit tests, if relevant
I have tried my changes with data and code:
- I have run the integration tests with ./scripts/actions.sh integration-test
- (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
- (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
I have informed and aligned with people impacted by my change:
- for config changes: the MatterMost channels and/or a design doc
- for changes of dependencies: the MatterMost software development channel

…atial-parallel

clessig

Didn't go through the code in detail but here are some high level comments, in addition to what we already discussed yesterday. For some pieces I probably also need an explanation at some point.

clessig · 2026-08-05T09:38:26Z


+# Split the encoder's HEALPix/location dimension across this many consecutive
+# distributed ranks through local assimilation.
+encoder_spatial_parallel_size: 4


We should have a section distributed where we collect everything related to parallelization, including with_ddp, with_fsdp etc. Please try to structure it well so that we can pass sub-dicts of it to the relevant parts

Remove encoder_ since we will have homogeneous domain decomposition througout the model

Agreed. The corresponding configs are also updated.

clessig · 2026-08-05T09:43:15Z

 import torch.distributed as dist

 SYNC_TIMEOUT_SEC = 60 * 60  # 1 hour
+_ENCODER_SPATIAL_GROUPS: dict[int, tuple[dist.ProcessGroup, int]] = {}


Remove ENCODER

clessig · 2026-08-05T09:43:25Z

    return dist.get_rank()


+def get_encoder_spatial_parallel_size(cf) -> int:


Remove _encoder

clessig · 2026-08-06T12:16:37Z

+    num_cells: int,
+    cell_start: int,
+    cell_end: int,
+) -> list[np.typing.NDArray[np.int64]]:


What are the assumptions on cell_start and cell_end? The goal of the function is to split the data in all healpix cells according to regions?

The assumptions are based on the nested Healpix cell numbering. Detailed explanation can be found below.

clessig · 2026-08-06T12:18:19Z

-        self.world_size = cf.world_size
+        # Ranks in one encoder-spatial group must consume the same batch. Data
+        # parallelism therefore operates across groups, not across individual ranks.
+        spatial_parallel_size = get_encoder_spatial_parallel_size(cf)


This information should be organized in a small class that is defined in utils.distributed. This will also be need in model.

clessig · 2026-08-06T12:30:43Z

        """

-        cell_lens = torch.sum(batch.tokens_lens, 2).flatten()
+        tokens_lens_global = batch.tokens_lens


Can we encapsulate this into a separate function

clessig · 2026-08-06T12:31:50Z


-        if batch.tokens_lens.shape[2] == 1:
+        if not x_embeds:
+            # A spatial rank can legitimately own a domain with no observations


Tighten comment

clessig · 2026-08-06T12:32:23Z

        self.q_cells: torch.Tensor | None = None
        self.streams: dict[str, typing.Any] = cf.streams
        self.target_token_engines = None
+        self.decoder_spatial_parallel_group = None


Same as above in MultiStreamDataSampler: we want a class for this

clessig · 2026-08-06T12:32:32Z

+        # Decoder cells follow the encoder's contiguous spatial-rank ownership.
+        # Keeping the same group also means a cell's nine-token HEALPix
+        # neighbourhood is sent to the rank that already owns that cell.
+        self.decoder_spatial_parallel_group = self.encoder.spatial_parallel_group


Same as above

clessig · 2026-08-06T12:33:50Z

+                ],
+                dim=1,
+            )
+        gathered_pred = all_gather(pred, group=self.decoder_spatial_parallel_group)


Anemoi has this function wrapped. Worth to have a discussion of what the pros and cons are.

Agreed. A wrapper could centralize the process-group handling and hide the tensor-collective details. For now, I would keep the explicit call because this is the only use site and the required operation is clear, but we can revisit this if the pattern appears elsewhere.

pkuyyj and others added 8 commits July 28, 2026 15:14

Add HEALPix encoder spatial parallelism

15dcff5

Build encoder inputs on local HEALPix domains

db4b477

Merge remote-tracking branch 'origin/develop' into feature/encoder-sp…

00cb384

…atial-parallel

Document encoder spatial parallelism

b79bb02

Document contiguous HEALPix rank ownership

dd3b9b9

Handle empty spatial HEALPix domains

760fb25

Parallelize decoder by HEALPix cell

d753973

Fix spatial decoder edge cases

6f08777

github-project-automation Bot added this to WeatherGen-dev Jul 31, 2026

Document decoder spatial parallelism

2b721dd

github-actions Bot added the model Related to model training or definition (not generic infra) label Jul 31, 2026

This was referenced Aug 5, 2026

Clean up the encoder/decoder PR #2730

Open

Investigate memory-efficient data reading #2732

Open

clessig reviewed Aug 6, 2026

View reviewed changes

pkuyyj added 6 commits August 11, 2026 14:23

Rename shared spatial parallel settings

52d2e1d

Group distributed configuration settings

ad1e12e

Centralize spatial parallel topology

1cca130

Encapsulate local encoder token selection

4af42ce

Tighten empty-observation comment

b3499b3

Encapsulate global encoder token lengths

2e3f59a

pkuyyj force-pushed the feature/decoder-healpix-cell-parallelism branch from 211eb4a to 2e3f59a Compare August 11, 2026 14:33

		return dist.get_rank()


		def get_encoder_spatial_parallel_size(cf) -> int:

Uh oh!

Conversation

pkuyyj commented Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Memory allocation profiling results:

Issue Number

Checklist before asking for review

Uh oh!

clessig left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pkuyyj commented Jul 31, 2026 •

edited

Loading